diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 7a1add3..6b8baf8 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -99,6 +99,7 @@
         "framework_graphics_flags_java_lib",
         "hwui_flags_java_lib",
         "libcore_exported_aconfig_flags_lib",
+        "libgui_flags_java_lib",
         "power_flags_lib",
         "sdk_sandbox_flags_lib",
         "surfaceflinger_flags_java_lib",
@@ -1208,6 +1209,12 @@
     defaults: ["framework-minus-apex-aconfig-java-defaults"],
 }
 
+java_aconfig_library {
+    name: "libgui_flags_java_lib",
+    aconfig_declarations: "libgui_flags",
+    defaults: ["framework-minus-apex-aconfig-java-defaults"],
+}
+
 // Content Capture
 aconfig_declarations {
     name: "android.view.contentcapture.flags-aconfig",
diff --git a/Android.bp b/Android.bp
index cf73451..7f4871f 100644
--- a/Android.bp
+++ b/Android.bp
@@ -401,6 +401,7 @@
     ],
     sdk_version: "core_platform",
     static_libs: [
+        "aconfig_storage_reader_java",
         "android.hardware.common.fmq-V1-java",
         "bouncycastle-repackaged-unbundled",
         "com.android.sysprop.foldlockbehavior",
@@ -636,7 +637,6 @@
         "core/java/com/android/internal/util/AsyncService.java",
         "core/java/com/android/internal/util/Protocol.java",
         "telephony/java/android/telephony/Annotation.java",
-        ":net-utils-framework-wifi-common-srcs",
     ],
     libs: [
         "framework-annotations-lib",
diff --git a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
index 857154f..9a178e5 100644
--- a/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/alarm/AlarmManagerService.java
@@ -4495,8 +4495,9 @@
                             final int[] userIds =
                                     mUserWakeupStore.getUserIdsToWakeup(nowELAPSED);
                             for (int i = 0; i < userIds.length; i++) {
-                                if (!mActivityManagerInternal.startUserInBackground(
-                                        userIds[i])) {
+                                if (mActivityManagerInternal.isUserRunning(userIds[i], 0)
+                                        || !mActivityManagerInternal.startUserInBackground(
+                                                userIds[i])) {
                                     mUserWakeupStore.removeUserWakeup(userIds[i]);
                                 }
                             }
diff --git a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
index c3fe031..d92351d 100644
--- a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
+++ b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
@@ -1990,7 +1990,7 @@
             }
         }
         if (android.app.admin.flags.Flags.disallowUserControlBgUsageFix()) {
-            if (!Flags.avoidIdleCheck()) {
+            if (!Flags.avoidIdleCheck() || mInjector.getBootPhase() >= PHASE_BOOT_COMPLETED) {
                 postCheckIdleStates(userId);
             }
         }
diff --git a/core/api/current.txt b/core/api/current.txt
index 69c409b..d610f4c 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -54862,8 +54862,6 @@
     method @Deprecated public void addAction(int);
     method public void addChild(android.view.View);
     method public void addChild(android.view.View, int);
-    method @FlaggedApi("android.view.accessibility.support_multiple_labeledby") public void addLabeledBy(@NonNull android.view.View);
-    method @FlaggedApi("android.view.accessibility.support_multiple_labeledby") public void addLabeledBy(@NonNull android.view.View, int);
     method public boolean canOpenPopup();
     method public int describeContents();
     method public java.util.List<android.view.accessibility.AccessibilityNodeInfo> findAccessibilityNodeInfosByText(String);
@@ -54892,7 +54890,6 @@
     method public int getInputType();
     method public android.view.accessibility.AccessibilityNodeInfo getLabelFor();
     method public android.view.accessibility.AccessibilityNodeInfo getLabeledBy();
-    method @FlaggedApi("android.view.accessibility.support_multiple_labeledby") @NonNull public java.util.List<android.view.accessibility.AccessibilityNodeInfo> getLabeledByList();
     method public int getLiveRegion();
     method public int getMaxTextLength();
     method @NonNull public java.time.Duration getMinDurationBetweenContentChanges();
@@ -54953,8 +54950,6 @@
     method public boolean removeAction(android.view.accessibility.AccessibilityNodeInfo.AccessibilityAction);
     method public boolean removeChild(android.view.View);
     method public boolean removeChild(android.view.View, int);
-    method @FlaggedApi("android.view.accessibility.support_multiple_labeledby") public boolean removeLabeledBy(@NonNull android.view.View);
-    method @FlaggedApi("android.view.accessibility.support_multiple_labeledby") public boolean removeLabeledBy(@NonNull android.view.View, int);
     method public void setAccessibilityDataSensitive(boolean);
     method public void setAccessibilityFocused(boolean);
     method public void setAvailableExtraData(java.util.List<java.lang.String>);
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index fd0262e..7674246 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -6718,6 +6718,14 @@
     field public static final int STATUS_OK = 0; // 0x0
   }
 
+  @FlaggedApi("android.media.soundtrigger.sound_trigger_generic_model_api") public static final class SoundTrigger.GenericSoundModel extends android.hardware.soundtrigger.SoundTrigger.SoundModel implements android.os.Parcelable {
+    ctor public SoundTrigger.GenericSoundModel(@NonNull java.util.UUID, @NonNull java.util.UUID, @Nullable byte[], int);
+    ctor public SoundTrigger.GenericSoundModel(@NonNull java.util.UUID, @NonNull java.util.UUID, @Nullable byte[]);
+    method public int describeContents();
+    method public void writeToParcel(@NonNull android.os.Parcel, int);
+    field @NonNull public static final android.os.Parcelable.Creator<android.hardware.soundtrigger.SoundTrigger.GenericSoundModel> CREATOR;
+  }
+
   public static final class SoundTrigger.Keyphrase implements android.os.Parcelable {
     ctor public SoundTrigger.Keyphrase(int, int, @NonNull java.util.Locale, @NonNull String, @Nullable int[]);
     method public int describeContents();
diff --git a/core/java/android/app/TaskInfo.java b/core/java/android/app/TaskInfo.java
index 531537c..c83dd65 100644
--- a/core/java/android/app/TaskInfo.java
+++ b/core/java/android/app/TaskInfo.java
@@ -304,6 +304,12 @@
     public boolean isTopActivityStyleFloating;
 
     /**
+     * The last non-fullscreen bounds the task was launched in or resized to.
+     * @hide
+     */
+    public Rect lastNonFullscreenBounds;
+
+    /**
      * The URI of the intent that generated the top-most activity opened using a URL.
      * @hide
      */
@@ -450,6 +456,7 @@
                 && Objects.equals(topActivity, that.topActivity)
                 && isTopActivityTransparent == that.isTopActivityTransparent
                 && isTopActivityStyleFloating == that.isTopActivityStyleFloating
+                && lastNonFullscreenBounds == this.lastNonFullscreenBounds
                 && Objects.equals(capturedLink, that.capturedLink)
                 && capturedLinkTimestamp == that.capturedLinkTimestamp
                 && appCompatTaskInfo.equalsForTaskOrganizer(that.appCompatTaskInfo);
@@ -522,6 +529,7 @@
         displayAreaFeatureId = source.readInt();
         isTopActivityTransparent = source.readBoolean();
         isTopActivityStyleFloating = source.readBoolean();
+        lastNonFullscreenBounds = source.readTypedObject(Rect.CREATOR);
         capturedLink = source.readTypedObject(Uri.CREATOR);
         capturedLinkTimestamp = source.readLong();
         appCompatTaskInfo = source.readTypedObject(AppCompatTaskInfo.CREATOR);
@@ -572,6 +580,7 @@
         dest.writeInt(displayAreaFeatureId);
         dest.writeBoolean(isTopActivityTransparent);
         dest.writeBoolean(isTopActivityStyleFloating);
+        dest.writeTypedObject(lastNonFullscreenBounds, flags);
         dest.writeTypedObject(capturedLink, flags);
         dest.writeLong(capturedLinkTimestamp);
         dest.writeTypedObject(appCompatTaskInfo, flags);
@@ -612,6 +621,7 @@
                 + " displayAreaFeatureId=" + displayAreaFeatureId
                 + " isTopActivityTransparent=" + isTopActivityTransparent
                 + " isTopActivityStyleFloating=" + isTopActivityStyleFloating
+                + " lastNonFullscreenBounds=" + lastNonFullscreenBounds
                 + " capturedLink=" + capturedLink
                 + " capturedLinkTimestamp=" + capturedLinkTimestamp
                 + " appCompatTaskInfo=" + appCompatTaskInfo
diff --git a/core/java/android/companion/virtual/flags/flags.aconfig b/core/java/android/companion/virtual/flags/flags.aconfig
index 55278f6..19de793 100644
--- a/core/java/android/companion/virtual/flags/flags.aconfig
+++ b/core/java/android/companion/virtual/flags/flags.aconfig
@@ -117,3 +117,11 @@
     description: "Enable high resolution scroll"
     bug: "335160780"
 }
+
+flag {
+  name: "camera_multiple_input_streams"
+  is_exported: true
+  namespace: "virtual_devices"
+  description: "Expose multiple surface for the virtual camera owner for different stream resolution"
+  bug: "341083465"
+}
diff --git a/core/java/android/content/pm/multiuser.aconfig b/core/java/android/content/pm/multiuser.aconfig
index ce241c1..88fbbdd 100644
--- a/core/java/android/content/pm/multiuser.aconfig
+++ b/core/java/android/content/pm/multiuser.aconfig
@@ -190,6 +190,18 @@
   }
 }
 
+flag {
+    name: "cache_user_serial_number_read_only"
+    namespace: "multiuser"
+    description: "Optimise user serial number retrieval. Read only flag, so that it can be used before the flags are initialized."
+    bug: "353134536"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+  }
+  is_fixed_read_only: true
+}
+
+
 # This flag guards the private space feature and all its implementations excluding the APIs. APIs are guarded by android.os.Flags.allow_private_profile.
 flag {
     name: "enable_private_space_features"
diff --git a/core/java/android/hardware/input/VirtualKeyEvent.java b/core/java/android/hardware/input/VirtualKeyEvent.java
index c0102bf..da959af 100644
--- a/core/java/android/hardware/input/VirtualKeyEvent.java
+++ b/core/java/android/hardware/input/VirtualKeyEvent.java
@@ -286,6 +286,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualKeyEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/input/VirtualMouseButtonEvent.java b/core/java/android/hardware/input/VirtualMouseButtonEvent.java
index fc42b15..333c3c7 100644
--- a/core/java/android/hardware/input/VirtualMouseButtonEvent.java
+++ b/core/java/android/hardware/input/VirtualMouseButtonEvent.java
@@ -197,6 +197,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualMouseButtonEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/input/VirtualMouseRelativeEvent.java b/core/java/android/hardware/input/VirtualMouseRelativeEvent.java
index 2a42cfc..86d759d 100644
--- a/core/java/android/hardware/input/VirtualMouseRelativeEvent.java
+++ b/core/java/android/hardware/input/VirtualMouseRelativeEvent.java
@@ -135,6 +135,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualMouseRelativeEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/input/VirtualMouseScrollEvent.java b/core/java/android/hardware/input/VirtualMouseScrollEvent.java
index c89c188..a4958c7 100644
--- a/core/java/android/hardware/input/VirtualMouseScrollEvent.java
+++ b/core/java/android/hardware/input/VirtualMouseScrollEvent.java
@@ -146,6 +146,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualMouseScrollEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/input/VirtualRotaryEncoderScrollEvent.java b/core/java/android/hardware/input/VirtualRotaryEncoderScrollEvent.java
index 8c98abd..033b1c1 100644
--- a/core/java/android/hardware/input/VirtualRotaryEncoderScrollEvent.java
+++ b/core/java/android/hardware/input/VirtualRotaryEncoderScrollEvent.java
@@ -129,6 +129,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualRotaryEncoderScrollEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/input/VirtualStylusButtonEvent.java b/core/java/android/hardware/input/VirtualStylusButtonEvent.java
index 97a4cd0..8fcf561b 100644
--- a/core/java/android/hardware/input/VirtualStylusButtonEvent.java
+++ b/core/java/android/hardware/input/VirtualStylusButtonEvent.java
@@ -187,6 +187,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualStylusButtonEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/input/VirtualStylusMotionEvent.java b/core/java/android/hardware/input/VirtualStylusMotionEvent.java
index 2ab76ae..0ac6f3a 100644
--- a/core/java/android/hardware/input/VirtualStylusMotionEvent.java
+++ b/core/java/android/hardware/input/VirtualStylusMotionEvent.java
@@ -377,6 +377,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualStylusMotionEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/input/VirtualTouchEvent.java b/core/java/android/hardware/input/VirtualTouchEvent.java
index 7936dfe..0cccd25 100644
--- a/core/java/android/hardware/input/VirtualTouchEvent.java
+++ b/core/java/android/hardware/input/VirtualTouchEvent.java
@@ -354,6 +354,10 @@
          * obtained from {@link SystemClock#uptimeMillis()} (with nanosecond precision instead of
          * millisecond), but can be different depending on the use case.
          * This field is optional and can be omitted.
+         * <p>
+         * If this field is unset, then the time at which this event is sent to the framework would
+         * be considered as the event time (even though
+         * {@link VirtualTouchEvent#getEventTimeNanos()}) would return {@code 0L}).
          *
          * @return this builder, to allow for chaining of calls
          * @see InputEvent#getEventTime()
diff --git a/core/java/android/hardware/soundtrigger/SoundTrigger.java b/core/java/android/hardware/soundtrigger/SoundTrigger.java
index bfff4db..e33a5c9 100644
--- a/core/java/android/hardware/soundtrigger/SoundTrigger.java
+++ b/core/java/android/hardware/soundtrigger/SoundTrigger.java
@@ -29,6 +29,7 @@
 import static java.util.Objects.requireNonNull;
 
 import android.annotation.ElapsedRealtimeLong;
+import android.annotation.FlaggedApi;
 import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
@@ -874,10 +875,9 @@
     /*****************************************************************************
      * A GenericSoundModel is a specialized {@link SoundModel} for non-voice sound
      * patterns.
-     *
-     * @hide
      ****************************************************************************/
-    public static class GenericSoundModel extends SoundModel implements Parcelable {
+    @FlaggedApi(android.media.soundtrigger.Flags.FLAG_SOUND_TRIGGER_GENERIC_MODEL_API)
+    public static final class GenericSoundModel extends SoundModel implements Parcelable {
 
         public static final @android.annotation.NonNull Parcelable.Creator<GenericSoundModel> CREATOR
                 = new Parcelable.Creator<GenericSoundModel>() {
@@ -890,12 +890,27 @@
             }
         };
 
+        /**
+         * Constructor for {@link GenericSoundModel} with version.
+         *
+         * @param uuid Unique identifier for this sound model.
+         * @param vendorUuid Unique vendor identifier for this sound model.
+         * @param data Opaque data for this sound model.
+         * @param version Vendor-specific version number of this sound model.
+         */
         public GenericSoundModel(@NonNull UUID uuid, @NonNull UUID vendorUuid,
                 @Nullable byte[] data, int version) {
-            super(uuid, vendorUuid, TYPE_GENERIC_SOUND, data, version);
+            super(uuid, Objects.requireNonNull(vendorUuid, "vendorUuid cannot be null"),
+                    TYPE_GENERIC_SOUND, data, version);
         }
 
-        @UnsupportedAppUsage
+        /**
+         * Constructor for {@link GenericSoundModel} without version. The version is set to -1.
+         *
+         * @param uuid Unique identifier for this sound model.
+         * @param vendorUuid Unique vendor identifier for this sound model.
+         * @param data Opaque data for this sound model.
+         */
         public GenericSoundModel(@NonNull UUID uuid, @NonNull UUID vendorUuid,
                 @Nullable byte[] data) {
             this(uuid, vendorUuid, data, -1);
@@ -919,7 +934,7 @@
         }
 
         @Override
-        public void writeToParcel(Parcel dest, int flags) {
+        public void writeToParcel(@NonNull Parcel dest, int flags) {
             dest.writeString(getUuid().toString());
             if (getVendorUuid() == null) {
                 dest.writeInt(-1);
diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index 2447ff9..623196b 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -66,6 +66,7 @@
             POWER_COMPONENT_WAKELOCK,
             POWER_COMPONENT_MEMORY,
             POWER_COMPONENT_PHONE,
+            POWER_COMPONENT_AMBIENT_DISPLAY,
             POWER_COMPONENT_IDLE,
             POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS,
     })
@@ -208,7 +209,8 @@
                 POWER_COMPONENT_VIDEO,
                 POWER_COMPONENT_FLASHLIGHT,
                 POWER_COMPONENT_CAMERA,
-                POWER_COMPONENT_GNSS};
+                POWER_COMPONENT_GNSS,
+                POWER_COMPONENT_SENSORS};
         Arrays.sort(supportedPowerComponents);
         SUPPORTED_POWER_COMPONENTS_PER_PROCESS_STATE = IntArray.wrap(supportedPowerComponents);
     };
diff --git a/core/java/android/os/BatteryStats.java b/core/java/android/os/BatteryStats.java
index b2f333a..c7751e3 100644
--- a/core/java/android/os/BatteryStats.java
+++ b/core/java/android/os/BatteryStats.java
@@ -2491,7 +2491,7 @@
     public static final int SCREEN_BRIGHTNESS_LIGHT = 3;
     public static final int SCREEN_BRIGHTNESS_BRIGHT = 4;
 
-    static final String[] SCREEN_BRIGHTNESS_NAMES = {
+    public static final String[] SCREEN_BRIGHTNESS_NAMES = {
         "dark", "dim", "medium", "light", "bright"
     };
 
@@ -3077,7 +3077,7 @@
     public static final String[] HISTORY_EVENT_NAMES = new String[] {
             "null", "proc", "fg", "top", "sync", "wake_lock_in", "job", "user", "userfg", "conn",
             "active", "pkginst", "pkgunin", "alarm", "stats", "pkginactive", "pkgactive",
-            "tmpwhitelist", "screenwake", "wakeupap", "longwake", "est_capacity", "state"
+            "tmpwhitelist", "screenwake", "wakeupap", "longwake", "state"
     };
 
     public static final String[] HISTORY_EVENT_CHECKIN_NAMES = new String[] {
diff --git a/core/java/android/os/ServiceManager.java b/core/java/android/os/ServiceManager.java
index e95c6a4..8aec7eb 100644
--- a/core/java/android/os/ServiceManager.java
+++ b/core/java/android/os/ServiceManager.java
@@ -425,7 +425,7 @@
     private static IBinder rawGetService(String name) throws RemoteException {
         final long start = sStatLogger.getTime();
 
-        final IBinder binder = getIServiceManager().getService(name).getBinder();
+        final IBinder binder = getIServiceManager().getService2(name).getBinder();
 
         final int time = (int) sStatLogger.logDurationStat(Stats.GET_SERVICE, start);
 
diff --git a/core/java/android/os/ServiceManagerNative.java b/core/java/android/os/ServiceManagerNative.java
index 6c9a5c7..5a9c878 100644
--- a/core/java/android/os/ServiceManagerNative.java
+++ b/core/java/android/os/ServiceManagerNative.java
@@ -57,8 +57,14 @@
         return mRemote;
     }
 
+    // TODO(b/355394904): This function has been deprecated, please use getService2 instead.
     @UnsupportedAppUsage
-    public Service getService(String name) throws RemoteException {
+    public IBinder getService(String name) throws RemoteException {
+        // Same as checkService (old versions of servicemanager had both methods).
+        return checkService(name).getBinder();
+    }
+
+    public Service getService2(String name) throws RemoteException {
         // Same as checkService (old versions of servicemanager had both methods).
         return checkService(name);
     }
diff --git a/core/java/android/os/UserManager.java b/core/java/android/os/UserManager.java
index 3aa42c6..392b6eb 100644
--- a/core/java/android/os/UserManager.java
+++ b/core/java/android/os/UserManager.java
@@ -6442,7 +6442,11 @@
      */
     @UnsupportedAppUsage
     public int getUserSerialNumber(@UserIdInt int userId) {
-        if (android.multiuser.Flags.cacheUserSerialNumber()) {
+        // Read only flag should is to fix early access to this API
+        // cacheUserSerialNumber to be removed after the
+        // cacheUserSerialNumberReadOnly is fully rolled out
+        if (android.multiuser.Flags.cacheUserSerialNumberReadOnly()
+                || android.multiuser.Flags.cacheUserSerialNumber()) {
             // System user serial number is always 0, and it always exists.
             // There is no need to call binder for that.
             if (userId == UserHandle.USER_SYSTEM) {
diff --git a/core/java/android/os/VibrationEffect.java b/core/java/android/os/VibrationEffect.java
index 44edf29..475984e 100644
--- a/core/java/android/os/VibrationEffect.java
+++ b/core/java/android/os/VibrationEffect.java
@@ -345,7 +345,10 @@
     @FlaggedApi(FLAG_VENDOR_VIBRATION_EFFECTS)
     @RequiresPermission(android.Manifest.permission.VIBRATE_VENDOR_EFFECTS)
     public static VibrationEffect createVendorEffect(@NonNull PersistableBundle effect) {
-        return new VendorEffect(effect, VendorEffect.DEFAULT_STRENGTH, VendorEffect.DEFAULT_SCALE);
+        VibrationEffect vendorEffect = new VendorEffect(effect, VendorEffect.DEFAULT_STRENGTH,
+                VendorEffect.DEFAULT_SCALE);
+        vendorEffect.validate();
+        return vendorEffect;
     }
 
     /**
@@ -1204,9 +1207,7 @@
             }
             return mEffectStrength == other.mEffectStrength
                     && (Float.compare(mLinearScale, other.mLinearScale) == 0)
-                    // Make sure it calls unparcel for both before calling BaseBundle.kindofEquals.
-                    && mVendorData.size() == other.mVendorData.size()
-                    && BaseBundle.kindofEquals(mVendorData, other.mVendorData);
+                    && isPersistableBundleEquals(mVendorData, other.mVendorData);
         }
 
         @Override
@@ -1243,6 +1244,55 @@
             out.writeFloat(mLinearScale);
         }
 
+        /**
+         * Compares two {@link PersistableBundle} objects are equals.
+         */
+        private static boolean isPersistableBundleEquals(
+                PersistableBundle first, PersistableBundle second) {
+            if (first == second) {
+                return true;
+            }
+            if (first == null || second == null || first.size() != second.size()) {
+                return false;
+            }
+            for (String key : first.keySet()) {
+                if (!isPersistableBundleSupportedValueEquals(first.get(key), second.get(key))) {
+                    return false;
+                }
+            }
+            return true;
+        }
+
+        /**
+         * Compares two values which type is supported by {@link PersistableBundle}.
+         *
+         * <p>If the type isn't supported. The equality is done by {@link Object#equals(Object)}.
+         */
+        private static boolean isPersistableBundleSupportedValueEquals(
+                Object first, Object second) {
+            if (first == second) {
+                return true;
+            } else if (first == null || second == null
+                    || !first.getClass().equals(second.getClass())) {
+                return false;
+            } else if (first instanceof PersistableBundle) {
+                return isPersistableBundleEquals(
+                        (PersistableBundle) first, (PersistableBundle) second);
+            } else if (first instanceof int[]) {
+                return Arrays.equals((int[]) first, (int[]) second);
+            } else if (first instanceof long[]) {
+                return Arrays.equals((long[]) first, (long[]) second);
+            } else if (first instanceof double[]) {
+                return Arrays.equals((double[]) first, (double[]) second);
+            } else if (first instanceof boolean[]) {
+                return Arrays.equals((boolean[]) first, (boolean[]) second);
+            } else if (first instanceof String[]) {
+                return Arrays.equals((String[]) first, (String[]) second);
+            } else {
+                return Objects.equals(first, second);
+            }
+        }
+
         @NonNull
         public static final Creator<VendorEffect> CREATOR =
                 new Creator<VendorEffect>() {
diff --git a/core/java/android/os/Vibrator.java b/core/java/android/os/Vibrator.java
index 99bd67b..161cce0 100644
--- a/core/java/android/os/Vibrator.java
+++ b/core/java/android/os/Vibrator.java
@@ -516,7 +516,7 @@
      */
     @RequiresPermission(android.Manifest.permission.VIBRATE)
     public void vibrate(@NonNull VibrationEffect vibe,
-            @NonNull VibrationAttributes attributes, @NonNull String reason) {
+            @NonNull VibrationAttributes attributes, String reason) {
         vibrate(Process.myUid(), mPackageName, vibe, reason, attributes);
     }
 
diff --git a/core/java/android/os/vibrator/flags.aconfig b/core/java/android/os/vibrator/flags.aconfig
index f4e2a7e..62b3682 100644
--- a/core/java/android/os/vibrator/flags.aconfig
+++ b/core/java/android/os/vibrator/flags.aconfig
@@ -64,3 +64,13 @@
         purpose: PURPOSE_FEATURE
     }
 }
+
+flag {
+    namespace: "haptics"
+    name: "throttle_vibration_params_requests"
+    description: "Control the frequency of vibration params requests to prevent overloading the vendor service"
+    bug: "355320860"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/core/java/android/provider/Settings.java b/core/java/android/provider/Settings.java
index 850b979..c186538 100644
--- a/core/java/android/provider/Settings.java
+++ b/core/java/android/provider/Settings.java
@@ -12560,6 +12560,14 @@
                 "contextual_screen_timeout_enabled";
 
         /**
+         * Whether hinge angle lidevent is enabled.
+         *
+         * @hide
+         */
+        public static final String HINGE_ANGLE_LIDEVENT_ENABLED =
+                "hinge_angle_lidevent_enabled";
+
+        /**
          * Whether lockscreen weather is enabled.
          *
          * @hide
diff --git a/core/java/android/service/dreams/DreamService.java b/core/java/android/service/dreams/DreamService.java
index fbeab84..74545a8 100644
--- a/core/java/android/service/dreams/DreamService.java
+++ b/core/java/android/service/dreams/DreamService.java
@@ -266,9 +266,9 @@
     private boolean mDozing;
     private boolean mWindowless;
     private boolean mPreviewMode;
-    private int mDozeScreenState = Display.STATE_UNKNOWN;
-    private @Display.StateReason int mDozeScreenStateReason = Display.STATE_REASON_UNKNOWN;
-    private int mDozeScreenBrightness = PowerManager.BRIGHTNESS_DEFAULT;
+    private volatile int mDozeScreenState = Display.STATE_UNKNOWN;
+    private volatile @Display.StateReason int mDozeScreenStateReason = Display.STATE_REASON_UNKNOWN;
+    private volatile int mDozeScreenBrightness = PowerManager.BRIGHTNESS_DEFAULT;
 
     private boolean mDebug = false;
 
@@ -1346,7 +1346,11 @@
                     Slog.w(mTag, "WakeUp was called before the dream was attached.");
                 } else {
                     try {
-                        mDreamManager.finishSelf(mDreamToken, false /*immediate*/);
+                        if (startAndStopDozingInBackground()) {
+                            mDreamManager.finishSelfOneway(mDreamToken, false /*immediate*/);
+                        } else {
+                            mDreamManager.finishSelf(mDreamToken, false /*immediate*/);
+                        }
                     } catch (RemoteException ex) {
                         // system server died
                     }
@@ -1497,7 +1501,11 @@
         if (mFinished || mWaking) {
             Slog.w(mTag, "attach() called after dream already finished");
             try {
-                mDreamManager.finishSelf(dreamToken, true /*immediate*/);
+                if (startAndStopDozingInBackground()) {
+                    mDreamManager.finishSelfOneway(dreamToken, true /*immediate*/);
+                } else {
+                    mDreamManager.finishSelf(dreamToken, true /*immediate*/);
+                }
             } catch (RemoteException ex) {
                 // system server died
             }
diff --git a/core/java/android/service/dreams/IDreamManager.aidl b/core/java/android/service/dreams/IDreamManager.aidl
index 620eef6..76f6363 100644
--- a/core/java/android/service/dreams/IDreamManager.aidl
+++ b/core/java/android/service/dreams/IDreamManager.aidl
@@ -39,7 +39,9 @@
     @UnsupportedAppUsage
     boolean isDreamingOrInPreview();
     boolean canStartDreaming(boolean isScreenOn);
+    /** @deprecated Please use finishSelfOneway instead. */
     void finishSelf(in IBinder token, boolean immediate);
+    /** @deprecated Please use startDozingOneway instead. */
     void startDozing(in IBinder token, int screenState, int reason, int screenBrightness);
     void stopDozing(in IBinder token);
     void forceAmbientDisplayEnabled(boolean enabled);
diff --git a/core/java/android/telephony/PhoneStateListener.java b/core/java/android/telephony/PhoneStateListener.java
index 4281da1..5ac0c50 100644
--- a/core/java/android/telephony/PhoneStateListener.java
+++ b/core/java/android/telephony/PhoneStateListener.java
@@ -1689,6 +1689,10 @@
         public final void onCarrierRoamingNtnModeChanged(boolean active) {
             // not supported on the deprecated interface - Use TelephonyCallback instead
         }
+
+        public final void onCarrierRoamingNtnEligibleStateChanged(boolean eligible) {
+            // not supported on the deprecated interface - Use TelephonyCallback instead
+        }
     }
 
     private void log(String s) {
diff --git a/core/java/android/telephony/TelephonyCallback.java b/core/java/android/telephony/TelephonyCallback.java
index b8b84d9..c360e64 100644
--- a/core/java/android/telephony/TelephonyCallback.java
+++ b/core/java/android/telephony/TelephonyCallback.java
@@ -653,6 +653,27 @@
     public static final int EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED = 42;
 
     /**
+     * Event for listening to changes in carrier roaming non-terrestrial network eligibility.
+     *
+     * @see CarrierRoamingNtnModeListener
+     *
+     * Device is eligible for satellite communication if all the following conditions are met:
+     * <ul>
+     * <li>Any subscription on the device supports P2P satellite messaging which is defined by
+     * {@link CarrierConfigManager#KEY_SATELLITE_ATTACH_SUPPORTED_BOOL} </li>
+     * <li>{@link CarrierConfigManager#KEY_CARRIER_ROAMING_NTN_CONNECT_TYPE_INT} set to
+     * {@link CarrierConfigManager#CARRIER_ROAMING_NTN_CONNECT_MANUAL} </li>
+     * <li>The device is in {@link ServiceState#STATE_OUT_OF_SERVICE}, not connected to Wi-Fi,
+     * and the hysteresis timer defined by {@link CarrierConfigManager
+     * #KEY_CARRIER_SUPPORTED_SATELLITE_NOTIFICATION_HYSTERESIS_SEC_INT} is expired.
+     * </li>
+     * </ul>
+     *
+     * @hide
+     */
+    public static final int EVENT_CARRIER_ROAMING_NTN_ELIGIBLE_STATE_CHANGED = 43;
+
+    /**
      * @hide
      */
     @IntDef(prefix = {"EVENT_"}, value = {
@@ -697,7 +718,8 @@
             EVENT_MEDIA_QUALITY_STATUS_CHANGED,
             EVENT_EMERGENCY_CALLBACK_MODE_CHANGED,
             EVENT_SIMULTANEOUS_CELLULAR_CALLING_SUBSCRIPTIONS_CHANGED,
-            EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED
+            EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED,
+            EVENT_CARRIER_ROAMING_NTN_ELIGIBLE_STATE_CHANGED
     })
     @Retention(RetentionPolicy.SOURCE)
     public @interface TelephonyEvent {
@@ -1711,6 +1733,23 @@
          *                           {code false} otherwise.
          */
         void onCarrierRoamingNtnModeChanged(boolean active);
+
+        /**
+         * Callback invoked when carrier roaming non-terrestrial network eligibility changes.
+         *
+         * @param eligible {@code true} when the device is eligible for satellite
+         * communication if all the following conditions are met:
+         * <ul>
+         * <li>Any subscription on the device supports P2P satellite messaging which is defined by
+         * {@link CarrierConfigManager#KEY_SATELLITE_ATTACH_SUPPORTED_BOOL} </li>
+         * <li>{@link CarrierConfigManager#KEY_CARRIER_ROAMING_NTN_CONNECT_TYPE_INT} set to
+         * {@link CarrierConfigManager#CARRIER_ROAMING_NTN_CONNECT_MANUAL} </li>
+         * <li>The device is in {@link ServiceState#STATE_OUT_OF_SERVICE}, not connected to Wi-Fi,
+         * and the hysteresis timer defined by {@link CarrierConfigManager
+         * #KEY_CARRIER_SUPPORTED_SATELLITE_NOTIFICATION_HYSTERESIS_SEC_INT} is expired. </li>
+         * </ul>
+         */
+        default void onCarrierRoamingNtnEligibleStateChanged(boolean eligible) {}
     }
 
     /**
@@ -2125,5 +2164,16 @@
             Binder.withCleanCallingIdentity(
                     () -> mExecutor.execute(() -> listener.onCarrierRoamingNtnModeChanged(active)));
         }
+
+        public void onCarrierRoamingNtnEligibleStateChanged(boolean eligible) {
+            if (!Flags.carrierRoamingNbIotNtn()) return;
+
+            CarrierRoamingNtnModeListener listener =
+                    (CarrierRoamingNtnModeListener) mTelephonyCallbackWeakRef.get();
+            if (listener == null) return;
+
+            Binder.withCleanCallingIdentity(() -> mExecutor.execute(
+                    () -> listener.onCarrierRoamingNtnEligibleStateChanged(eligible)));
+        }
     }
 }
diff --git a/core/java/android/telephony/TelephonyRegistryManager.java b/core/java/android/telephony/TelephonyRegistryManager.java
index 6160fdb..10f03c1 100644
--- a/core/java/android/telephony/TelephonyRegistryManager.java
+++ b/core/java/android/telephony/TelephonyRegistryManager.java
@@ -1091,6 +1091,34 @@
     }
 
     /**
+     * Notify external listeners that device eligibility to connect to carrier roaming
+     * non-terrestrial network changed.
+     *
+     * @param subId subscription ID.
+     * @param eligible {@code true} when the device is eligible for satellite
+     * communication if all the following conditions are met:
+     * <ul>
+     * <li>Any subscription supports P2P satellite messaging which is defined by
+     * {@link CarrierConfigManager#KEY_SATELLITE_ATTACH_SUPPORTED_BOOL} </li>
+     * <li>{@link CarrierConfigManager#KEY_CARRIER_ROAMING_NTN_CONNECT_TYPE_INT} set to
+     * {@link CarrierConfigManager#CARRIER_ROAMING_NTN_CONNECT_MANUAL} </li>
+     * <li>The device is in {@link ServiceState#STATE_OUT_OF_SERVICE}, not connected to Wi-Fi,
+     * and the hysteresis timer defined by {@link CarrierConfigManager
+     * #KEY_CARRIER_SUPPORTED_SATELLITE_NOTIFICATION_HYSTERESIS_SEC_INT} is expired. </li>
+     * </ul>
+     *
+     * @hide
+     */
+    public void notifyCarrierRoamingNtnEligibleStateChanged(int subId, boolean eligible) {
+        try {
+            sRegistry.notifyCarrierRoamingNtnEligibleStateChanged(subId, eligible);
+        } catch (RemoteException ex) {
+            // system server crash
+            throw ex.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Processes potential event changes from the provided {@link TelephonyCallback}.
      *
      * @param telephonyCallback callback for monitoring callback changes to the telephony state.
@@ -1246,6 +1274,11 @@
         if (telephonyCallback instanceof TelephonyCallback.CarrierRoamingNtnModeListener) {
             eventList.add(TelephonyCallback.EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED);
         }
+
+        if (telephonyCallback instanceof TelephonyCallback.CarrierRoamingNtnModeListener) {
+            eventList.add(TelephonyCallback.EVENT_CARRIER_ROAMING_NTN_ELIGIBLE_STATE_CHANGED);
+        }
+
         return eventList;
     }
 
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index 634469d..cf329d3 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -164,6 +164,9 @@
             float width, float height, float vecX, float vecY,
             float maxStretchAmountX, float maxStretchAmountY, float childRelativeLeft,
             float childRelativeTop, float childRelativeRight, float childRelativeBottom);
+    private static native void nativeSetEdgeExtensionEffect(long transactionObj, long nativeObj,
+                                                            boolean leftEdge, boolean rightEdge,
+                                                            boolean topEdge, boolean bottomEdge);
     private static native void nativeSetTrustedOverlay(long transactionObj, long nativeObject,
             int isTrustedOverlay);
     private static native void nativeSetDropInputMode(
@@ -3513,6 +3516,19 @@
         /**
          * @hide
          */
+        public Transaction setEdgeExtensionEffect(SurfaceControl sc, int edge) {
+            checkPreconditions(sc);
+
+            nativeSetEdgeExtensionEffect(
+                    mNativeObject, sc.mNativeObject,
+                    (edge & WindowInsets.Side.LEFT) != 0, (edge & WindowInsets.Side.RIGHT) != 0,
+                    (edge & WindowInsets.Side.TOP) != 0, (edge & WindowInsets.Side.BOTTOM) != 0);
+            return this;
+        }
+
+        /**
+         * @hide
+         */
         @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.O)
         public Transaction setLayerStack(SurfaceControl sc, int layerStack) {
             checkPreconditions(sc);
@@ -4882,4 +4898,5 @@
     public static void notifyShutdown() {
         nativeNotifyShutdown();
     }
+
 }
diff --git a/core/java/android/view/View.java b/core/java/android/view/View.java
index 64c7766..5f8bea1 100644
--- a/core/java/android/view/View.java
+++ b/core/java/android/view/View.java
@@ -34071,14 +34071,21 @@
     }
 
     private float convertVelocityToFrameRate(float velocityPps) {
-        // From UXR study, premium experience is:
-        // 1500+    dp/s: 120fps
-        // 0 - 1500 dp/s:  80fps
-        // OEMs are likely to modify this to balance battery and user experience for their
-        // specific device.
+        // Internal testing has shown that this gives a premium experience:
+        // above 300dp/s => 120fps
+        // between 300dp/s and 125fps => 80fps
+        // below 125dp/s => 60fps
         float density = mAttachInfo.mDensity;
         float velocityDps = velocityPps / density;
-        return (velocityDps >= 1500f) ? MAX_FRAME_RATE : 80f;
+        float frameRate;
+        if (velocityDps > 300f) {
+            frameRate = MAX_FRAME_RATE; // Use maximum at fast motion
+        } else if (velocityDps > 125f) {
+            frameRate = 80f; // Use medium frame rate when motion is slower
+        } else {
+            frameRate = 60f; // Use minimum frame rate when motion is very slow
+        }
+        return frameRate;
     }
 
     /**
diff --git a/core/java/android/view/accessibility/AccessibilityCache.java b/core/java/android/view/accessibility/AccessibilityCache.java
index f3cde43..376e66f 100644
--- a/core/java/android/view/accessibility/AccessibilityCache.java
+++ b/core/java/android/view/accessibility/AccessibilityCache.java
@@ -19,6 +19,7 @@
 
 import static android.view.accessibility.AccessibilityNodeInfo.FOCUS_ACCESSIBILITY;
 
+import android.annotation.Nullable;
 import android.os.Build;
 import android.os.SystemClock;
 import android.util.ArraySet;
@@ -48,6 +49,8 @@
 
     private boolean mEnabled = true;
 
+    private final SparseArray<String> mWindowIdToEventSourceClassName = new SparseArray<>();
+
     /**
      * {@link AccessibilityEvent} types that are critical for the cache to stay up to date
      *
@@ -273,8 +276,11 @@
                     clearSubTreeLocked(event.getWindowId(), event.getSourceNodeId());
                 } break;
 
-                case AccessibilityEvent.TYPE_WINDOWS_CHANGED:
+                case AccessibilityEvent.TYPE_WINDOWS_CHANGED: {
                     mValidWindowCacheTimeStamp = event.getEventTime();
+                    if (event.getWindowChanges() == AccessibilityEvent.WINDOWS_CHANGE_REMOVED) {
+                        mWindowIdToEventSourceClassName.remove(event.getWindowId());
+                    }
                     if (event.getWindowChanges()
                             == AccessibilityEvent.WINDOWS_CHANGE_ACCESSIBILITY_FOCUSED) {
                         // Don't need to clear all cache. Unless the changes are related to
@@ -282,8 +288,15 @@
                         clearWindowCacheLocked();
                         break;
                     }
+                    clear();
+                }
+                break;
                 case AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED: {
                     mValidWindowCacheTimeStamp = event.getEventTime();
+                    if (event.getContentChangeTypes() == 0 && event.getClassName() != null) {
+                        mWindowIdToEventSourceClassName.put(event.getWindowId(),
+                                event.getClassName().toString());
+                    }
                     clear();
                 } break;
             }
@@ -907,6 +920,12 @@
         }
     }
 
+    /** Returns the source class associated with the window with the given id. */
+    @Nullable
+    public String getEventSourceClassName(int windowId) {
+        return mWindowIdToEventSourceClassName.get(windowId);
+    }
+
     // Layer of indirection included to break dependency chain for testing
     public static class AccessibilityNodeRefresher {
         /** Refresh the given AccessibilityNodeInfo object. */
diff --git a/core/java/android/view/accessibility/AccessibilityNodeInfo.java b/core/java/android/view/accessibility/AccessibilityNodeInfo.java
index 90cfcb1..a5ba294 100644
--- a/core/java/android/view/accessibility/AccessibilityNodeInfo.java
+++ b/core/java/android/view/accessibility/AccessibilityNodeInfo.java
@@ -982,7 +982,6 @@
     private long mParentNodeId = UNDEFINED_NODE_ID;
     private long mLabelForId = UNDEFINED_NODE_ID;
     private long mLabeledById = UNDEFINED_NODE_ID;
-    private LongArray mLabeledByIds;
     private long mTraversalBefore = UNDEFINED_NODE_ID;
     private long mTraversalAfter = UNDEFINED_NODE_ID;
 
@@ -3600,131 +3599,6 @@
     }
 
     /**
-     * Adds the view which serves as the label of the view represented by
-     * this info for accessibility purposes. When more than one labels are
-     * added, the content from each label is combined in the order that
-     * they are added.
-     * <p>
-     * If visible text can be used to describe or give meaning to this UI,
-     * this method is preferred. For example, a TextView before an EditText
-     * in the UI usually specifies what information is contained in the
-     * EditText. Hence, the EditText is labelled by the TextView.
-     * </p>
-     *
-     * @param label A view that labels this node's source.
-     */
-    @FlaggedApi(Flags.FLAG_SUPPORT_MULTIPLE_LABELEDBY)
-    public void addLabeledBy(@NonNull View label) {
-        addLabeledBy(label, AccessibilityNodeProvider.HOST_VIEW_ID);
-    }
-
-    /**
-     * Adds the view which serves as the label of the view represented by
-     * this info for accessibility purposes. If <code>virtualDescendantId</code>
-     * is {@link View#NO_ID} the root is set as the label. When more than one
-     * labels are added, the content from each label is combined in the order
-     * that they are added.
-     * <p>
-     * A virtual descendant is an imaginary View that is reported as a part of the view
-     * hierarchy for accessibility purposes. This enables custom views that draw complex
-     * content to report themselves as a tree of virtual views, thus conveying their
-     * logical structure.
-     * </p>
-     * <p>
-     * If visible text can be used to describe or give meaning to this UI,
-     * this method is preferred. For example, a TextView before an EditText
-     * in the UI usually specifies what information is contained in the
-     * EditText. Hence, the EditText is labelled by the TextView.
-     * </p>
-     * <p>
-     *   <strong>Note:</strong> Cannot be called from an
-     *   {@link android.accessibilityservice.AccessibilityService}.
-     *   This class is made immutable before being delivered to an AccessibilityService.
-     * </p>
-     *
-     * @param root A root whose virtual descendant labels this node's source.
-     * @param virtualDescendantId The id of the virtual descendant.
-     */
-    @FlaggedApi(Flags.FLAG_SUPPORT_MULTIPLE_LABELEDBY)
-    public void addLabeledBy(@NonNull View root, int virtualDescendantId) {
-        enforceNotSealed();
-        Preconditions.checkNotNull(root, "%s must not be null", root);
-        if (mLabeledByIds == null) {
-            mLabeledByIds = new LongArray();
-        }
-        mLabeledById = makeNodeId(root.getAccessibilityViewId(), virtualDescendantId);
-        mLabeledByIds.add(mLabeledById);
-    }
-
-    /**
-     * Gets the list of node infos which serve as the labels of the view represented by
-     * this info for accessibility purposes.
-     *
-     * @return The list of labels in the order that they were added.
-     */
-    @FlaggedApi(Flags.FLAG_SUPPORT_MULTIPLE_LABELEDBY)
-    public @NonNull List<AccessibilityNodeInfo> getLabeledByList() {
-        enforceSealed();
-        List<AccessibilityNodeInfo> labels = new ArrayList<>();
-        if (mLabeledByIds == null) {
-            return labels;
-        }
-        for (int i = 0; i < mLabeledByIds.size(); i++) {
-            labels.add(getNodeForAccessibilityId(mConnectionId, mWindowId, mLabeledByIds.get(i)));
-        }
-        return labels;
-    }
-
-    /**
-     * Removes a label. If the label was not previously added to the node,
-     * calling this method has no effect.
-     * <p>
-     * <strong>Note:</strong> Cannot be called from an
-     * {@link android.accessibilityservice.AccessibilityService}.
-     * This class is made immutable before being delivered to an AccessibilityService.
-     * </p>
-     *
-     * @param label The node which serves as this node's label.
-     * @return true if the label was present
-     * @see #addLabeledBy(View)
-     */
-    @FlaggedApi(Flags.FLAG_SUPPORT_MULTIPLE_LABELEDBY)
-    public boolean removeLabeledBy(@NonNull View label) {
-        return removeLabeledBy(label, AccessibilityNodeProvider.HOST_VIEW_ID);
-    }
-
-    /**
-     * Removes a virtual label which is a descendant of the given
-     * <code>root</code>. If the label was not previously added to the node,
-     * calling this method has no effect.
-     *
-     * @param root The root of the virtual subtree.
-     * @param virtualDescendantId The id of the virtual node which serves as this node's label.
-     * @return true if the label was present
-     * @see #addLabeledBy(View, int)
-     */
-    @FlaggedApi(Flags.FLAG_SUPPORT_MULTIPLE_LABELEDBY)
-    public boolean removeLabeledBy(@NonNull View root, int virtualDescendantId) {
-        enforceNotSealed();
-        final LongArray labeledByIds = mLabeledByIds;
-        if (labeledByIds == null) {
-            return false;
-        }
-        final int rootAccessibilityViewId =
-                (root != null) ? root.getAccessibilityViewId() : UNDEFINED_ITEM_ID;
-        final long labeledById = makeNodeId(rootAccessibilityViewId, virtualDescendantId);
-        if (mLabeledById == labeledById) {
-            mLabeledById = UNDEFINED_NODE_ID;
-        }
-        final int index = labeledByIds.indexOf(labeledById);
-        if (index < 0) {
-            return false;
-        }
-        labeledByIds.remove(index);
-        return true;
-    }
-
-    /**
      * Sets the view which serves as the label of the view represented by
      * this info for accessibility purposes.
      *
@@ -3757,17 +3631,7 @@
         enforceNotSealed();
         final int rootAccessibilityViewId = (root != null)
                 ? root.getAccessibilityViewId() : UNDEFINED_ITEM_ID;
-        if (Flags.supportMultipleLabeledby()) {
-            if (mLabeledByIds == null) {
-                mLabeledByIds = new LongArray();
-            } else {
-                mLabeledByIds.clear();
-            }
-        }
         mLabeledById = makeNodeId(rootAccessibilityViewId, virtualDescendantId);
-        if (Flags.supportMultipleLabeledby()) {
-            mLabeledByIds.add(mLabeledById);
-        }
     }
 
     /**
@@ -4378,12 +4242,6 @@
         fieldIndex++;
         if (mLabeledById != DEFAULT.mLabeledById) nonDefaultFields |= bitAt(fieldIndex);
         fieldIndex++;
-        if (Flags.supportMultipleLabeledby()) {
-            if (!LongArray.elementsEqual(mLabeledByIds, DEFAULT.mLabeledByIds)) {
-                nonDefaultFields |= bitAt(fieldIndex);
-            }
-            fieldIndex++;
-        }
         if (mTraversalBefore != DEFAULT.mTraversalBefore) nonDefaultFields |= bitAt(fieldIndex);
         fieldIndex++;
         if (mTraversalAfter != DEFAULT.mTraversalAfter) nonDefaultFields |= bitAt(fieldIndex);
@@ -4525,20 +4383,6 @@
         if (isBitSet(nonDefaultFields, fieldIndex++)) parcel.writeLong(mParentNodeId);
         if (isBitSet(nonDefaultFields, fieldIndex++)) parcel.writeLong(mLabelForId);
         if (isBitSet(nonDefaultFields, fieldIndex++)) parcel.writeLong(mLabeledById);
-        if (Flags.supportMultipleLabeledby()) {
-            if (isBitSet(nonDefaultFields, fieldIndex++)) {
-                final LongArray labeledByIds = mLabeledByIds;
-                if (labeledByIds == null) {
-                    parcel.writeInt(0);
-                } else {
-                    final int labeledByIdsSize = labeledByIds.size();
-                    parcel.writeInt(labeledByIdsSize);
-                    for (int i = 0; i < labeledByIdsSize; i++) {
-                        parcel.writeLong(labeledByIds.get(i));
-                    }
-                }
-            }
-        }
         if (isBitSet(nonDefaultFields, fieldIndex++)) parcel.writeLong(mTraversalBefore);
         if (isBitSet(nonDefaultFields, fieldIndex++)) parcel.writeLong(mTraversalAfter);
         if (isBitSet(nonDefaultFields, fieldIndex++)) {
@@ -4706,9 +4550,6 @@
         mParentNodeId = other.mParentNodeId;
         mLabelForId = other.mLabelForId;
         mLabeledById = other.mLabeledById;
-        if (Flags.supportMultipleLabeledby()) {
-            mLabeledByIds = other.mLabeledByIds;
-        }
         mTraversalBefore = other.mTraversalBefore;
         mTraversalAfter = other.mTraversalAfter;
         mMinDurationBetweenContentChanges = other.mMinDurationBetweenContentChanges;
@@ -4815,20 +4656,6 @@
         if (isBitSet(nonDefaultFields, fieldIndex++)) mParentNodeId = parcel.readLong();
         if (isBitSet(nonDefaultFields, fieldIndex++)) mLabelForId = parcel.readLong();
         if (isBitSet(nonDefaultFields, fieldIndex++)) mLabeledById = parcel.readLong();
-        if (Flags.supportMultipleLabeledby()) {
-            if (isBitSet(nonDefaultFields, fieldIndex++)) {
-                final int labeledByIdsSize = parcel.readInt();
-                if (labeledByIdsSize <= 0) {
-                    mLabeledByIds = null;
-                } else {
-                    mLabeledByIds = new LongArray(labeledByIdsSize);
-                    for (int i = 0; i < labeledByIdsSize; i++) {
-                        final long labeledById = parcel.readLong();
-                        mLabeledByIds.add(labeledById);
-                    }
-                }
-            }
-        }
         if (isBitSet(nonDefaultFields, fieldIndex++)) mTraversalBefore = parcel.readLong();
         if (isBitSet(nonDefaultFields, fieldIndex++)) mTraversalAfter = parcel.readLong();
         if (isBitSet(nonDefaultFields, fieldIndex++)) {
diff --git a/core/java/android/view/animation/Animation.java b/core/java/android/view/animation/Animation.java
index 09306c7..2af935d 100644
--- a/core/java/android/view/animation/Animation.java
+++ b/core/java/android/view/animation/Animation.java
@@ -28,6 +28,7 @@
 import android.os.SystemProperties;
 import android.util.AttributeSet;
 import android.util.TypedValue;
+import android.view.WindowInsets;
 
 import dalvik.system.CloseGuard;
 
@@ -881,12 +882,13 @@
     }
 
     /**
-     * @return if a window animation has outsets applied to it.
+     * @return the edges to which outsets can be applied to
      *
      * @hide
      */
-    public boolean hasExtension() {
-        return false;
+    @WindowInsets.Side.InsetsSide
+    public int getExtensionEdges() {
+        return 0x0;
     }
 
     /**
diff --git a/core/java/android/view/animation/AnimationSet.java b/core/java/android/view/animation/AnimationSet.java
index 5aaa994..bbdc9d0 100644
--- a/core/java/android/view/animation/AnimationSet.java
+++ b/core/java/android/view/animation/AnimationSet.java
@@ -21,6 +21,7 @@
 import android.graphics.RectF;
 import android.os.Build;
 import android.util.AttributeSet;
+import android.view.WindowInsets;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -540,12 +541,12 @@
 
     /** @hide */
     @Override
-    public boolean hasExtension() {
+    @WindowInsets.Side.InsetsSide
+    public int getExtensionEdges() {
+        int edge = 0x0;
         for (Animation animation : mAnimations) {
-            if (animation.hasExtension()) {
-                return true;
-            }
+            edge |= animation.getExtensionEdges();
         }
-        return false;
+        return edge;
     }
 }
diff --git a/core/java/android/view/animation/ExtendAnimation.java b/core/java/android/view/animation/ExtendAnimation.java
index 210eb8a..ed047c7 100644
--- a/core/java/android/view/animation/ExtendAnimation.java
+++ b/core/java/android/view/animation/ExtendAnimation.java
@@ -20,6 +20,7 @@
 import android.content.res.TypedArray;
 import android.graphics.Insets;
 import android.util.AttributeSet;
+import android.view.WindowInsets;
 
 /**
  * An animation that controls the outset of an object.
@@ -50,6 +51,8 @@
     private float mToRightValue;
     private float mToBottomValue;
 
+    private int mExtensionEdges = 0x0;
+
     /**
      * Constructor used when an ExtendAnimation is loaded from a resource.
      *
@@ -151,9 +154,22 @@
 
     /** @hide */
     @Override
-    public boolean hasExtension() {
-        return mFromInsets.left < 0 || mFromInsets.top < 0 || mFromInsets.right < 0
-                || mFromInsets.bottom < 0;
+    @WindowInsets.Side.InsetsSide
+    public int getExtensionEdges() {
+        mExtensionEdges = 0x0;
+        if (mFromLeftValue > 0 || mToLeftValue > 0) {
+            mExtensionEdges |= WindowInsets.Side.LEFT;
+        }
+        if (mFromRightValue > 0 || mToRightValue > 0) {
+            mExtensionEdges |= WindowInsets.Side.RIGHT;
+        }
+        if (mFromTopValue > 0 || mToTopValue > 0) {
+            mExtensionEdges |= WindowInsets.Side.TOP;
+        }
+        if (mFromBottomValue > 0 || mToBottomValue > 0) {
+            mExtensionEdges |= WindowInsets.Side.BOTTOM;
+        }
+        return mExtensionEdges;
     }
 
     @Override
diff --git a/core/java/android/view/autofill/AutofillClientController.java b/core/java/android/view/autofill/AutofillClientController.java
index d505c733..95cae226 100644
--- a/core/java/android/view/autofill/AutofillClientController.java
+++ b/core/java/android/view/autofill/AutofillClientController.java
@@ -582,4 +582,9 @@
             Log.e(TAG, "authenticate() failed for intent:" + intent, e);
         }
     }
+
+    @Override
+    public boolean isActivityResumed() {
+        return mActivity.isResumed();
+    }
 }
diff --git a/core/java/android/view/autofill/AutofillManager.java b/core/java/android/view/autofill/AutofillManager.java
index 02a86c9..79ecfe1e 100644
--- a/core/java/android/view/autofill/AutofillManager.java
+++ b/core/java/android/view/autofill/AutofillManager.java
@@ -114,6 +114,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -774,6 +775,13 @@
     // dataset in responses. Used to avoid request pre-fill request again and again.
     private final ArraySet<AutofillId> mAllTrackedViews = new ArraySet<>();
 
+    // Whether we need to re-attempt fill again. Needed for case of relayout.
+    private boolean mFillReAttemptNeeded = false;
+
+    private Map<Integer, AutofillId> mFingerprintToViewMap = new ArrayMap<>();
+
+    private AutofillStateFingerprint mAutofillStateFingerprint;
+
     /** @hide */
     public interface AutofillClient {
         /**
@@ -909,6 +917,11 @@
          * @return An ID that is unique in the activity.
          */
         @Nullable AutofillId autofillClientGetNextAutofillId();
+
+        /**
+         * @return Whether the activity is resumed or not.
+         */
+        boolean isActivityResumed();
     }
 
     /**
@@ -919,6 +932,7 @@
         mService = service;
         mOptions = context.getAutofillOptions();
         mIsFillRequested = new AtomicBoolean(false);
+        mAutofillStateFingerprint = AutofillStateFingerprint.createInstance();
 
         mIsFillDialogEnabled = AutofillFeatureFlags.isFillDialogEnabled();
         mFillDialogEnabledHints = AutofillFeatureFlags.getFillDialogEnabledHints();
@@ -1357,6 +1371,20 @@
             mOnInvisibleCalled = true;
 
             if (isExpiredResponse) {
+                if (mRelayoutFix && isAuthenticationPending()) {
+                    Log.i(TAG, "onInvisibleForAutofill(): Ignoring expiringResponse due to pending"
+                            + " authentication");
+                    try {
+                        mService.notifyNotExpiringResponseDuringAuth(
+                                mSessionId, mContext.getUserId());
+                    } catch (RemoteException e) {
+                        // The failure could be a consequence of something going wrong on the
+                        // server side. Do nothing here since it's just logging, but it's
+                        // possible follow-up actions may fail.
+                    }
+                    return;
+                }
+                Log.i(TAG, "onInvisibleForAutofill(): expiringResponse");
                 // Notify service the response has expired.
                 updateSessionLocked(/* id= */ null, /* bounds= */ null, /* value= */ null,
                         ACTION_RESPONSE_EXPIRED, /* flags= */ 0);
@@ -1513,14 +1541,29 @@
     }
 
     /**
+     * Called to log notify view entered was ignored due to pending auth
+     * @hide
+     */
+    public void notifyViewEnteredIgnoredDuringAuthCount() {
+        try {
+            mService.notifyViewEnteredIgnoredDuringAuthCount(mSessionId, mContext.getUserId());
+        } catch (RemoteException e) {
+            // The failure could be a consequence of something going wrong on the
+            // server side. Do nothing here since it's just logging, but it's
+            // possible follow-up actions may fail.
+        }
+    }
+
+    /**
      * Called to check if we should retry fill.
      * Useful for knowing whether to attempt refill after relayout.
      *
      * @hide
      */
     public boolean shouldRetryFill() {
-        // TODO: Implement in follow-up cl
-        return false;
+        synchronized (mLock) {
+            return isAuthenticationPending() && mFillReAttemptNeeded;
+        }
     }
 
     /**
@@ -1531,8 +1574,13 @@
      */
     public boolean attemptRefill() {
         Log.i(TAG, "Attempting refill");
-        // TODO: Implement in follow-up cl
-        return false;
+        // Find active autofillable views. Compute their fingerprints
+        List<View> autofillableViews =
+                getClient().autofillClientFindAutofillableViewsByTraversal();
+        if (sDebug) {
+            Log.d(TAG, "Autofillable views count:" + autofillableViews.size());
+        }
+        return mAutofillStateFingerprint.attemptRefill(autofillableViews, this);
     }
 
     /**
@@ -2493,7 +2541,13 @@
 
     /** @hide */
     public void onAuthenticationResult(int authenticationId, Intent data, View focusView) {
+        if (sVerbose) {
+            Log.v(TAG, "onAuthenticationResult(): authId= " + authenticationId + ", data=" + data);
+        }
         if (!hasAutofillFeature()) {
+            if (sVerbose) {
+                Log.v(TAG, "onAuthenticationResult(): autofill not enabled");
+            }
             return;
         }
         // TODO: the result code is being ignored, so this method is not reliably
@@ -2501,10 +2555,6 @@
         // set the EXTRA_AUTHENTICATION_RESULT extra, but it could cause weird results if the
         // service set the extra and returned RESULT_CANCELED...
 
-        if (sDebug) {
-            Log.d(TAG, "onAuthenticationResult(): id= " + authenticationId + ", data=" + data);
-        }
-
         synchronized (mLock) {
             if (!isActiveLocked()) {
                 Log.w(TAG, "onAuthenticationResult(): sessionId=" + mSessionId + " not active");
@@ -2661,6 +2711,7 @@
             mSessionId = receiver.getIntResult();
             if (mSessionId != NO_SESSION) {
                 mState = STATE_ACTIVE;
+                mAutofillStateFingerprint.setSessionId(mSessionId);
             }
             final int extraFlags = receiver.getOptionalExtraIntResult(0);
             if ((extraFlags & RECEIVER_FLAG_SESSION_FOR_AUGMENTED_AUTOFILL_ONLY) != 0) {
@@ -2722,6 +2773,9 @@
         if (resetEnteredIds) {
             mEnteredIds = null;
         }
+        mFillReAttemptNeeded = false;
+        mFingerprintToViewMap.clear();
+        mAutofillStateFingerprint = AutofillStateFingerprint.createInstance();
     }
 
     @GuardedBy("mLock")
@@ -2984,8 +3038,12 @@
             Intent fillInIntent, boolean authenticateInline) {
         synchronized (mLock) {
             if (sessionId == mSessionId) {
-                if (mRelayoutFixDeprecated) {
+                if (mRelayoutFixDeprecated || mRelayoutFix) {
                     mState = STATE_PENDING_AUTHENTICATION;
+                    if (sVerbose) {
+                        Log.v(TAG, "entering STATE_PENDING_AUTHENTICATION : mRelayoutFix:"
+                                + mRelayoutFix);
+                    }
                 }
                 final AutofillClient client = getClient();
                 if (client != null) {
@@ -3191,17 +3249,56 @@
 
     @GuardedBy("mLock")
     private void handleFailedIdsLocked(@NonNull ArrayList<AutofillId> failedIds) {
+        handleFailedIdsLocked(failedIds, null, false, false);
+    }
+
+    @GuardedBy("mLock")
+    private void handleFailedIdsLocked(@NonNull ArrayList<AutofillId> failedIds,
+            ArrayList<AutofillValue> failedAutofillValues, boolean hideHighlight,
+            boolean isRefill) {
         if (!failedIds.isEmpty() && sVerbose) {
             Log.v(TAG, "autofill(): total failed views: " + failedIds);
         }
+
+        if (mRelayoutFix && !failedIds.isEmpty()) {
+            // Activity isn't in resumed state, so it's very possible that relayout could've
+            // occurred, so wait for it to declare proper failure. It's a temporary failure at the
+            // moment. We'll try again later when the activity is resumed.
+
+            // The above doesn't seem to be the correct way. Look for pending auth cases.
+            // TODO(b/238252288): Check whether there was any auth done at all
+            mFillReAttemptNeeded = true;
+            mAutofillStateFingerprint.storeFailedIdsAndValues(
+                    failedIds, failedAutofillValues, hideHighlight);
+        }
         try {
-            mService.setAutofillFailure(mSessionId, failedIds, mContext.getUserId());
+            mService.setAutofillFailure(mSessionId, failedIds, isRefill, mContext.getUserId());
         } catch (RemoteException e) {
             // In theory, we could ignore this error since it's not a big deal, but
             // in reality, we rather crash the app anyways, as the failure could be
             // a consequence of something going wrong on the server side...
             throw e.rethrowFromSystemServer();
         }
+        if (mRelayoutFix && !failedIds.isEmpty()) {
+            if (!getClient().isActivityResumed()) {
+                if (sVerbose) {
+                    Log.v(TAG, "handleFailedIdsLocked(): failed id's exist, but activity not"
+                            + " resumed");
+                }
+            } else {
+                if (isRefill) {
+                    Log.i(TAG, "handleFailedIdsLocked(): Attempted refill, but failed");
+                } else {
+                    // activity has been resumed, try to re-fill
+                    // getClient().isActivityResumed() && !failedIds.isEmpty() && !isRefill
+                    // TODO(b/238252288): Do better state management, and only trigger the following
+                    //  if there was auth previously.
+                    Log.i(TAG, "handleFailedIdsLocked(): Attempting refill");
+                    attemptRefill();
+                    mFillReAttemptNeeded = false;
+                }
+            }
+        }
     }
 
     private void autofill(int sessionId, List<AutofillId> ids, List<AutofillValue> values,
@@ -3216,13 +3313,46 @@
                 return;
             }
 
-            final int itemCount = ids.size();
-            int numApplied = 0;
-            ArrayMap<View, SparseArray<AutofillValue>> virtualValues = null;
             final View[] views = client.autofillClientFindViewsByAutofillIdTraversal(
                     Helper.toArray(ids));
 
+            autofill(views, ids, values, hideHighlight, false);
+        }
+    }
+
+    void autofill(View[] views, List<AutofillId> ids, List<AutofillValue> values,
+            boolean hideHighlight, boolean isRefill) {
+        if (sVerbose) {
+            Log.v(TAG, "autofill() ids:" + ids + " isRefill:" + isRefill);
+        }
+        synchronized (mLock) {
+            final AutofillClient client = getClient();
+            if (client == null) {
+                return;
+            }
+
+            if (ids == null) {
+                Log.i(TAG, "autofill(): No id's to fill");
+                return;
+            }
+
+            if (mRelayoutFix && isRefill) {
+                try {
+                    mService.setAutofillIdsAttemptedForRefill(
+                            mSessionId, ids, mContext.getUserId());
+                } catch (RemoteException e) {
+                    // The failure could be a consequence of something going wrong on the
+                    // server side. Do nothing here since it's just logging, but it's
+                    // possible follow-up actions may fail.
+                }
+            }
+
+            final int itemCount = ids.size();
+            int numApplied = 0;
+            ArrayMap<View, SparseArray<AutofillValue>> virtualValues = null;
+
             ArrayList<AutofillId> failedIds = new ArrayList<>();
+            ArrayList<AutofillValue> failedAutofillValues = new ArrayList<>();
 
             if (mLastAutofilledData == null) {
                 mLastAutofilledData = new ParcelableMap(itemCount);
@@ -3237,7 +3367,9 @@
                     // the service; this is fine, but we need to update the view status in the
                     // server side so it can be triggered again.
                     Log.d(TAG, "autofill(): no View with id " + id);
+                    // Possible relayout scenario
                     failedIds.add(id);
+                    failedAutofillValues.add(value);
                     continue;
                 }
                 // Mark the view as to be autofilled with 'value'
@@ -3268,7 +3400,8 @@
                 }
             }
 
-            handleFailedIdsLocked(failedIds);
+            handleFailedIdsLocked(
+                    failedIds, failedAutofillValues, hideHighlight, isRefill);
 
             if (virtualValues != null) {
                 for (int i = 0; i < virtualValues.size(); i++) {
@@ -3322,7 +3455,7 @@
     private void reportAutofillContentFailure(AutofillId id) {
         try {
             mService.setAutofillFailure(mSessionId, Collections.singletonList(id),
-                    mContext.getUserId());
+                    false /* isRefill */, mContext.getUserId());
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
@@ -3349,20 +3482,22 @@
     }
 
     /**
-     *  Set the tracked views.
+     * Set the tracked views.
      *
-     * @param trackedIds The views to be tracked.
+     * @param trackedIds              The views to be tracked.
      * @param saveOnAllViewsInvisible Finish the session once all tracked views are invisible.
-     * @param saveOnFinish Finish the session once the activity is finished.
-     * @param fillableIds Views that might anchor FillUI.
-     * @param saveTriggerId View that when clicked triggers commit().
+     * @param saveOnFinish            Finish the session once the activity is finished.
+     * @param fillableIds             Views that might anchor FillUI.
+     * @param saveTriggerId           View that when clicked triggers commit().
      */
     private void setTrackedViews(int sessionId, @Nullable AutofillId[] trackedIds,
             boolean saveOnAllViewsInvisible, boolean saveOnFinish,
-            @Nullable AutofillId[] fillableIds, @Nullable AutofillId saveTriggerId) {
+            @Nullable AutofillId[] fillableIds, @Nullable AutofillId saveTriggerId,
+            boolean shouldGrabViewFingerprints) {
         if (saveTriggerId != null) {
             saveTriggerId.resetSessionId();
         }
+        final ArraySet<AutofillId> allFillableIds = new ArraySet<>();
         synchronized (mLock) {
             if (sVerbose) {
                 Log.v(TAG, "setTrackedViews(): sessionId=" + sessionId
@@ -3372,6 +3507,7 @@
                         + ", fillableIds=" + Arrays.toString(fillableIds)
                         + ", saveTrigerId=" + saveTriggerId
                         + ", mFillableIds=" + mFillableIds
+                        + ", shouldGrabViewFingerprints=" + shouldGrabViewFingerprints
                         + ", mEnabled=" + mEnabled
                         + ", mSessionId=" + mSessionId);
             }
@@ -3405,7 +3541,6 @@
                     trackedIds = null;
                 }
 
-                final ArraySet<AutofillId> allFillableIds = new ArraySet<>();
                 if (mFillableIds != null) {
                     allFillableIds.addAll(mFillableIds);
                 }
@@ -3424,6 +3559,12 @@
                     mTrackedViews = null;
                 }
             }
+            if (mRelayoutFix && shouldGrabViewFingerprints) {
+                // For all the views: tracked and others, calculate fingerprints and store them.
+                mAutofillStateFingerprint.setUseRelativePosition(mRelativePositionForRelayout);
+                mAutofillStateFingerprint.storeStatePriorToAuthentication(
+                        getClient(), allFillableIds);
+            }
         }
     }
 
@@ -3845,7 +3986,7 @@
 
     @GuardedBy("mLock")
     private boolean isPendingAuthenticationLocked() {
-        return mRelayoutFixDeprecated && mState == STATE_PENDING_AUTHENTICATION;
+        return (mRelayoutFixDeprecated || mRelayoutFix) && mState == STATE_PENDING_AUTHENTICATION;
     }
 
     @GuardedBy("mLock")
@@ -3858,7 +3999,7 @@
         return mState == STATE_FINISHED;
     }
 
-    private void post(Runnable runnable) {
+    void post(Runnable runnable) {
         final AutofillClient client = getClient();
         if (client == null) {
             if (sVerbose) Log.v(TAG, "ignoring post() because client is null");
@@ -4700,11 +4841,11 @@
         @Override
         public void setTrackedViews(int sessionId, AutofillId[] ids,
                 boolean saveOnAllViewsInvisible, boolean saveOnFinish, AutofillId[] fillableIds,
-                AutofillId saveTriggerId) {
+                AutofillId saveTriggerId, boolean shouldGrabViewFingerprints) {
             final AutofillManager afm = mAfm.get();
             if (afm != null) {
                 afm.post(() -> afm.setTrackedViews(sessionId, ids, saveOnAllViewsInvisible,
-                        saveOnFinish, fillableIds, saveTriggerId));
+                        saveOnFinish, fillableIds, saveTriggerId, shouldGrabViewFingerprints));
             }
         }
 
diff --git a/core/java/android/view/autofill/AutofillStateFingerprint.java b/core/java/android/view/autofill/AutofillStateFingerprint.java
new file mode 100644
index 0000000..2db4285
--- /dev/null
+++ b/core/java/android/view/autofill/AutofillStateFingerprint.java
@@ -0,0 +1,352 @@
+/*
+ * Copyright (C) 2024 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 android.view.autofill;
+
+import static android.view.autofill.Helper.sDebug;
+
+import android.annotation.NonNull;
+import android.util.ArrayMap;
+import android.util.Log;
+import android.util.Slog;
+import android.view.View;
+import android.widget.TextView;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+/**
+ * This class manages and stores the autofillable views fingerprints for use in relayout situations.
+ * @hide
+ */
+@VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+public final class AutofillStateFingerprint {
+
+    ArrayList<AutofillId> mPriorAutofillIds;
+    ArrayList<Integer> mViewHashCodes; // each entry corresponding to mPriorAutofillIds .
+
+    boolean mHideHighlight = false;
+
+    private int mSessionId;
+
+    Map<Integer, AutofillId> mHashToAutofillIdMap = new ArrayMap<>();
+    Map<AutofillId, AutofillId> mOldIdsToCurrentAutofillIdMap = new ArrayMap<>();
+
+    // These failed id's are attempted to be refilled again after relayout.
+    private ArrayList<AutofillId> mFailedIds = new ArrayList<>();
+    private ArrayList<AutofillValue> mFailedAutofillValues = new ArrayList<>();
+
+    // whether to use relative positions for computing hashes.
+    private boolean mUseRelativePosition;
+
+    private static final String TAG = "AutofillStateFingerprint";
+
+    /**
+     * Returns an instance of this class
+     */
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+    public static AutofillStateFingerprint createInstance() {
+        return new AutofillStateFingerprint();
+    }
+
+    private AutofillStateFingerprint() {
+    }
+
+    /**
+     * Set sessionId for the instance
+     */
+    void setSessionId(int sessionId) {
+        mSessionId = sessionId;
+    }
+
+    /**
+     * Sets whether relative position of the views should be used to calculate fingerprints.
+     */
+    void setUseRelativePosition(boolean useRelativePosition) {
+        mUseRelativePosition = useRelativePosition;
+    }
+
+    /**
+     * Store the state of the views prior to the authentication.
+     */
+    void storeStatePriorToAuthentication(
+            AutofillManager.AutofillClient client, Set<AutofillId> autofillIds) {
+        if (mUseRelativePosition) {
+            List<View> autofillableViews = client.autofillClientFindAutofillableViewsByTraversal();
+            if (sDebug) {
+                Log.d(TAG, "Autofillable views count prior to auth:" + autofillableViews.size());
+            }
+//            ArrayList<Integer> hashes = getFingerprintIds(autofillableViews);
+
+            ArrayMap<Integer, View> hashes = getFingerprintIds(autofillableViews);
+            for (Map.Entry<Integer, View> entry : hashes.entrySet()) {
+                View view = entry.getValue();
+                if (view != null) {
+                    mHashToAutofillIdMap.put(entry.getKey(), view.getAutofillId());
+                } else {
+                    if (sDebug) {
+                        Log.d(TAG, "Encountered null view");
+                    }
+                }
+            }
+        } else {
+            // Just use the provided autofillIds and get their hashes
+            if (sDebug) {
+                Log.d(TAG, "Size of autofillId's being stored: " + autofillIds.size()
+                        + " list:" + autofillIds);
+            }
+            AutofillId[] autofillIdsArr = Helper.toArray(autofillIds);
+            View[] views = client.autofillClientFindViewsByAutofillIdTraversal(autofillIdsArr);
+            for (int i = 0; i < autofillIdsArr.length; i++) {
+                View view = views[i];
+                if (view != null) {
+                    int id = getEphemeralFingerprintId(view, 0 /* position irrelevant */);
+                    AutofillId autofillId = view.getAutofillId();
+                    autofillId.setSessionId(mSessionId);
+                    mHashToAutofillIdMap.put(id, autofillId);
+                } else {
+                    if (sDebug) {
+                        Log.d(TAG, "Encountered null view");
+                    }
+                }
+            }
+        }
+    }
+
+    /**
+     * Store failed ids, so that they can be refilled later
+     */
+    void storeFailedIdsAndValues(
+            @NonNull ArrayList<AutofillId> failedIds,
+            ArrayList<AutofillValue> failedAutofillValues,
+            boolean hideHighlight) {
+        for (AutofillId failedId : failedIds) {
+            if (failedId != null) {
+                failedId.setSessionId(mSessionId);
+            } else {
+                if (sDebug) {
+                    Log.d(TAG, "Got null failed ids");
+                }
+            }
+        }
+        mFailedIds = failedIds;
+        mFailedAutofillValues = failedAutofillValues;
+        mHideHighlight = hideHighlight;
+    }
+
+    private void dumpCurrentState() {
+        Log.d(TAG, "FailedId's: " + mFailedIds);
+        Log.d(TAG, "Hashes from map" + mHashToAutofillIdMap);
+    }
+
+    boolean attemptRefill(
+            List<View> currentAutofillableViews, @NonNull AutofillManager autofillManager) {
+        if (sDebug) {
+            dumpCurrentState();
+        }
+        // For the autofillable views, compute their hashes
+        ArrayMap<Integer, View> currentHashes = getFingerprintIds(currentAutofillableViews);
+
+        // For the computed hashes, try to look for the old fingerprints.
+        // If match found, update the new autofill ids of those views
+        Map<AutofillId, View> oldFailedIdsToCurrentViewMap = new HashMap<>();
+        for (Map.Entry<Integer, View> entry : currentHashes.entrySet()) {
+            View view = entry.getValue();
+            int currentHash = entry.getKey();
+            AutofillId currentAutofillId = view.getAutofillId();
+            currentAutofillId.setSessionId(mSessionId);
+            if (mHashToAutofillIdMap.containsKey(currentHash)) {
+                AutofillId oldAutofillId = mHashToAutofillIdMap.get(currentHash);
+                oldAutofillId.setSessionId(mSessionId);
+                mOldIdsToCurrentAutofillIdMap.put(oldAutofillId, currentAutofillId);
+                Log.i(TAG, "Mapping current autofill id: " + view.getAutofillId()
+                        + " to existing autofill id " + oldAutofillId);
+
+                oldFailedIdsToCurrentViewMap.put(oldAutofillId, view);
+            } else {
+                Log.i(TAG, "Couldn't map current autofill id: " + view.getAutofillId()
+                        + " with currentHash:" + currentHash + " for view:" + view);
+            }
+        }
+
+        int viewsCount = 0;
+        View[] views = new View[mFailedIds.size()];
+        for (int i = 0; i < mFailedIds.size(); i++) {
+            AutofillId oldAutofillId = mFailedIds.get(i);
+            AutofillId currentAutofillId = mOldIdsToCurrentAutofillIdMap.get(oldAutofillId);
+            if (currentAutofillId == null) {
+                if (sDebug) {
+                    Log.d(TAG, "currentAutofillId = null");
+                }
+            }
+            mFailedIds.set(i, currentAutofillId);
+            views[i] = oldFailedIdsToCurrentViewMap.get(oldAutofillId);
+            if (views[i] != null) {
+                viewsCount++;
+            }
+        }
+
+        if (sDebug) {
+            dumpCurrentState();
+        }
+
+        // Attempt autofill now
+        Slog.i(TAG, "Attempting refill of views. Found " + viewsCount
+                + " views to refill from previously " + mFailedIds.size()
+                + " failed ids:" + mFailedIds);
+        autofillManager.post(
+                () -> autofillManager.autofill(
+                        views, mFailedIds, mFailedAutofillValues, mHideHighlight,
+                        true /* isRefill */));
+
+        return false;
+    }
+
+    /**
+     * Retrieves fingerprint hashes for the views
+     */
+    ArrayMap<Integer, View> getFingerprintIds(@NonNull List<View> views) {
+        ArrayMap<Integer, View> map = new ArrayMap<>();
+        if (mUseRelativePosition) {
+            Collections.sort(views, (View v1, View v2) -> {
+                int[] posV1 = v1.getLocationOnScreen();
+                int[] posV2 = v2.getLocationOnScreen();
+
+                int compare = posV1[0] - posV2[0]; // x coordinate
+                if (compare != 0) {
+                    return compare;
+                }
+                compare = posV1[1] - posV2[1]; // y coordinate
+                if (compare != 0) {
+                    return compare;
+                }
+                // Sort on vertical
+                compare = compareTop(v1, v2);
+                if (compare != 0) {
+                    return compare;
+                }
+                compare = compareBottom(v1, v2);
+                if (compare != 0) {
+                    return compare;
+                }
+                compare = compareLeft(v1, v2);
+                if (compare != 0) {
+                    return compare;
+                }
+                return compareRight(v1, v2);
+                // Note that if compareRight also returned 0, that means both the views have exact
+                // same location, so just treat them as equal
+            });
+        }
+        for (int i = 0; i < views.size(); i++) {
+            View view = views.get(i);
+            map.put(getEphemeralFingerprintId(view, i), view);
+        }
+        return map;
+    }
+
+    /**
+     * Returns fingerprint hash for the view.
+     */
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PRIVATE)
+    public int getEphemeralFingerprintId(View v, int position) {
+        if (v == null) return -1;
+        int inputType = Integer.MIN_VALUE;
+        int imeOptions = Integer.MIN_VALUE;
+        boolean isSingleLine = false;
+        CharSequence hints = "";
+        if (v instanceof TextView) {
+            TextView tv = (TextView) v;
+            inputType = tv.getInputType();
+            hints = tv.getHint();
+            isSingleLine = tv.isSingleLine();
+            imeOptions = tv.getImeOptions();
+            // TODO(b/238252288): Consider adding more IME related fields.
+        }
+        CharSequence contentDesc = v.getContentDescription();
+        CharSequence tooltip = v.getTooltipText();
+
+        int autofillType = v.getAutofillType();
+        String[] autofillHints = v.getAutofillHints();
+        int visibility = v.getVisibility();
+
+        int paddingLeft = v.getPaddingLeft();
+        int paddingRight = v.getPaddingRight();
+        int paddingTop = v.getPaddingTop();
+        int paddingBottom = v.getPaddingBottom();
+
+        // TODO(b/238252288): Following are making relayout flaky. Do more analysis to figure out
+        //  why.
+        int height = v.getHeight();
+        int width = v.getWidth();
+
+        // Order doesn't matter much here. We can change the order, as long as we use the same
+        // order for storing and fetching fingerprints. The order can be changed in platform
+        // versions.
+        int hash = Objects.hash(visibility, inputType, imeOptions, isSingleLine, hints,
+                contentDesc, tooltip, autofillType, Arrays.deepHashCode(autofillHints),
+                paddingBottom, paddingTop, paddingRight, paddingLeft);
+        if (mUseRelativePosition) {
+            hash = Objects.hash(hash, position);
+        }
+        if (sDebug) {
+            Log.d(TAG, "Hash: " + hash + " for AutofillId:" + v.getAutofillId()
+                    + " visibility:" + visibility
+                    + " inputType:" + inputType
+                    + " imeOptions:" + imeOptions
+                    + " isSingleLine:" + isSingleLine
+                    + " hints:" + hints
+                    + " contentDesc:" + contentDesc
+                    + " tooltipText:" + tooltip
+                    + " autofillType:" + autofillType
+                    + " autofillHints:" + Arrays.toString(autofillHints)
+                    + " height:" + height
+                    + " width:" + width
+                    + " paddingLeft:" + paddingLeft
+                    + " paddingRight:" + paddingRight
+                    + " paddingTop:" + paddingTop
+                    + " paddingBottom:" + paddingBottom
+                    + " mUseRelativePosition" + mUseRelativePosition
+                    + " position:" + position
+            );
+        }
+        return hash;
+    }
+
+    private int compareTop(View v1, View v2) {
+        return v1.getTop() - v2.getTop();
+    }
+
+    private int compareBottom(View v1, View v2) {
+        return v1.getBottom() - v2.getBottom();
+    }
+
+    private int compareLeft(View v1, View v2) {
+        return v1.getLeft() - v2.getLeft();
+    }
+
+    private int compareRight(View v1, View v2) {
+        return v1.getRight() - v2.getRight();
+    }
+}
diff --git a/core/java/android/view/autofill/IAutoFillManager.aidl b/core/java/android/view/autofill/IAutoFillManager.aidl
index 2039b4d..f67405f 100644
--- a/core/java/android/view/autofill/IAutoFillManager.aidl
+++ b/core/java/android/view/autofill/IAutoFillManager.aidl
@@ -48,7 +48,7 @@
         in IResultReceiver result);
     void updateSession(int sessionId, in AutofillId id, in Rect bounds,
         in AutofillValue value, int action, int flags, int userId);
-    void setAutofillFailure(int sessionId, in List<AutofillId> ids, int userId);
+    void setAutofillFailure(int sessionId, in List<AutofillId> ids, boolean isRefill, int userId);
     void setViewAutofilled(int sessionId, in AutofillId id, int userId);
     void finishSession(int sessionId, int userId, int commitReason);
     void cancelSession(int sessionId, int userId);
@@ -67,4 +67,7 @@
     void getDefaultFieldClassificationAlgorithm(in IResultReceiver result);
     void setAugmentedAutofillWhitelist(in List<String> packages, in List<ComponentName> activities,
         in IResultReceiver result);
+    void notifyNotExpiringResponseDuringAuth(int sessionId, int userId);
+    void notifyViewEnteredIgnoredDuringAuthCount(int sessionId, int userId);
+    void setAutofillIdsAttemptedForRefill(int sessionId, in List<AutofillId> ids, int userId);
 }
diff --git a/core/java/android/view/autofill/IAutoFillManagerClient.aidl b/core/java/android/view/autofill/IAutoFillManagerClient.aidl
index 904a7e0..39d71da 100644
--- a/core/java/android/view/autofill/IAutoFillManagerClient.aidl
+++ b/core/java/android/view/autofill/IAutoFillManagerClient.aidl
@@ -72,7 +72,8 @@
       */
     void setTrackedViews(int sessionId, in @nullable AutofillId[] savableIds,
             boolean saveOnAllViewsInvisible, boolean saveOnFinish,
-            in @nullable AutofillId[] fillableIds, in AutofillId saveTriggerId);
+            in @nullable AutofillId[] fillableIds, in AutofillId saveTriggerId,
+            in boolean shouldGrabViewFingerprints);
 
     /**
      * Requests showing the fill UI.
diff --git a/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java b/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java
index 07a9794..a4ca55e 100644
--- a/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java
+++ b/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java
@@ -196,16 +196,20 @@
 
     /**
      * Invokes {@link IInputMethodManager#removeImeSurface()}
+     *
+     * @param displayId display ID from which this request originates
+     * @param exceptionHandler an optional {@link RemoteException} handler
      */
     @AnyThread
     @RequiresPermission(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
-    static void removeImeSurface(@Nullable Consumer<RemoteException> exceptionHandler) {
+    static void removeImeSurface(int displayId,
+            @Nullable Consumer<RemoteException> exceptionHandler) {
         final IInputMethodManager service = getService();
         if (service == null) {
             return;
         }
         try {
-            service.removeImeSurface();
+            service.removeImeSurface(displayId);
         } catch (RemoteException e) {
             handleRemoteExceptionOrRethrow(e, exceptionHandler);
         }
diff --git a/core/java/android/view/inputmethod/InputMethodManagerGlobal.java b/core/java/android/view/inputmethod/InputMethodManagerGlobal.java
index 5df9fd1..244b239 100644
--- a/core/java/android/view/inputmethod/InputMethodManagerGlobal.java
+++ b/core/java/android/view/inputmethod/InputMethodManagerGlobal.java
@@ -95,11 +95,13 @@
     /**
      * Invokes {@link IInputMethodManager#removeImeSurface()}
      *
+     * @param displayId display ID from which this request originates.
      * @param exceptionHandler an optional {@link RemoteException} handler.
      */
     @AnyThread
     @RequiresPermission(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
-    public static void removeImeSurface(@Nullable Consumer<RemoteException> exceptionHandler) {
-        IInputMethodManagerGlobalInvoker.removeImeSurface(exceptionHandler);
+    public static void removeImeSurface(int displayId,
+            @Nullable Consumer<RemoteException> exceptionHandler) {
+        IInputMethodManagerGlobalInvoker.removeImeSurface(displayId, exceptionHandler);
     }
 }
diff --git a/core/java/android/window/TransitionFilter.java b/core/java/android/window/TransitionFilter.java
index ec4e3e9..3cfde87 100644
--- a/core/java/android/window/TransitionFilter.java
+++ b/core/java/android/window/TransitionFilter.java
@@ -30,6 +30,8 @@
 import android.os.Parcelable;
 import android.view.WindowManager;
 
+import com.android.window.flags.Flags;
+
 /**
  * A parcelable filter that can be used for rerouting transitions to a remote. This is a local
  * representation so that the transition system doesn't need to make blocking queries over
@@ -183,6 +185,9 @@
         public ComponentName mTopActivity;
         public IBinder mLaunchCookie;
 
+        /** If non-null, requires the change to specifically have or not-have a custom animation. */
+        public Boolean mCustomAnimation = null;
+
         public Requirement() {
         }
 
@@ -196,6 +201,9 @@
             mOrder = in.readInt();
             mTopActivity = in.readTypedObject(ComponentName.CREATOR);
             mLaunchCookie = in.readStrongBinder();
+            // 0: null, 1: false, 2: true
+            final int customAnimRaw = in.readInt();
+            mCustomAnimation = customAnimRaw == 0 ? null : Boolean.valueOf(customAnimRaw == 2);
         }
 
         /** Go through changes and find if at-least one change matches this filter */
@@ -237,6 +245,23 @@
                 if (!matchesCookie(change.getTaskInfo())) {
                     continue;
                 }
+                if (mCustomAnimation != null
+                        // only applies to activity/task
+                        && (change.getTaskInfo() != null
+                                || change.getActivityComponent() != null)) {
+                    final TransitionInfo.AnimationOptions opts =
+                            Flags.moveAnimationOptionsToChange() ? change.getAnimationOptions()
+                                    : info.getAnimationOptions();
+                    if (opts != null) {
+                        boolean canActuallyOverride = change.getTaskInfo() == null
+                                || opts.getOverrideTaskTransition();
+                        if (mCustomAnimation != canActuallyOverride) {
+                            continue;
+                        }
+                    } else if (mCustomAnimation) {
+                        continue;
+                    }
+                }
                 return true;
             }
             return false;
@@ -286,6 +311,8 @@
             dest.writeInt(mOrder);
             dest.writeTypedObject(mTopActivity, flags);
             dest.writeStrongBinder(mLaunchCookie);
+            int customAnimRaw = mCustomAnimation == null ? 0 : (mCustomAnimation ? 2 : 1);
+            dest.writeInt(customAnimRaw);
         }
 
         @NonNull
@@ -327,6 +354,9 @@
             out.append(" order=" + containerOrderToString(mOrder));
             out.append(" topActivity=").append(mTopActivity);
             out.append(" launchCookie=").append(mLaunchCookie);
+            if (mCustomAnimation != null) {
+                out.append(" customAnim=").append(mCustomAnimation.booleanValue());
+            }
             out.append("}");
             return out.toString();
         }
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index 76989f9..725d496 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -158,17 +158,6 @@
 }
 
 flag {
-  name: "keyguard_appear_transition"
-  namespace: "windowing_frontend"
-  description: "Add transition when keyguard appears"
-  bug: "327970608"
-  is_fixed_read_only: true
-  metadata {
-      purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
   name: "get_dimmer_on_closing"
   namespace: "windowing_frontend"
   description: "Change check for when to ignore a closing task's dim"
diff --git a/core/java/com/android/internal/telephony/IPhoneStateListener.aidl b/core/java/com/android/internal/telephony/IPhoneStateListener.aidl
index 792c223..f177e14 100644
--- a/core/java/com/android/internal/telephony/IPhoneStateListener.aidl
+++ b/core/java/com/android/internal/telephony/IPhoneStateListener.aidl
@@ -82,4 +82,5 @@
     void onCallBackModeStopped(int type, int reason);
     void onSimultaneousCallingStateChanged(in int[] subIds);
     void onCarrierRoamingNtnModeChanged(in boolean active);
+    void onCarrierRoamingNtnEligibleStateChanged(in boolean eligible);
 }
diff --git a/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl b/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
index 04332cd..e500a37a 100644
--- a/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
+++ b/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
@@ -121,4 +121,5 @@
     void notifyCallbackModeStarted(int phoneId, int subId, int type);
     void notifyCallbackModeStopped(int phoneId, int subId, int type, int reason);
     void notifyCarrierRoamingNtnModeChanged(int subId, in boolean active);
+    void notifyCarrierRoamingNtnEligibleStateChanged(int subId, in boolean eligible);
 }
diff --git a/core/java/com/android/internal/view/IInputMethodManager.aidl b/core/java/com/android/internal/view/IInputMethodManager.aidl
index cba27ce..3f7ba0a 100644
--- a/core/java/com/android/internal/view/IInputMethodManager.aidl
+++ b/core/java/com/android/internal/view/IInputMethodManager.aidl
@@ -171,7 +171,7 @@
     @EnforcePermission("INTERNAL_SYSTEM_WINDOW")
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission(value = "
             + "android.Manifest.permission.INTERNAL_SYSTEM_WINDOW)")
-    void removeImeSurface();
+    void removeImeSurface(int displayId);
 
     /** Remove the IME surface. Requires passing the currently focused window. */
     oneway void removeImeSurfaceFromWindowAsync(in IBinder windowToken);
diff --git a/core/java/com/android/internal/widget/remotecompose/core/Operations.java b/core/java/com/android/internal/widget/remotecompose/core/Operations.java
index fc8668e..4b8dbf6 100644
--- a/core/java/com/android/internal/widget/remotecompose/core/Operations.java
+++ b/core/java/com/android/internal/widget/remotecompose/core/Operations.java
@@ -37,6 +37,7 @@
 import com.android.internal.widget.remotecompose.core.operations.FloatConstant;
 import com.android.internal.widget.remotecompose.core.operations.FloatExpression;
 import com.android.internal.widget.remotecompose.core.operations.Header;
+import com.android.internal.widget.remotecompose.core.operations.IntegerExpression;
 import com.android.internal.widget.remotecompose.core.operations.MatrixRestore;
 import com.android.internal.widget.remotecompose.core.operations.MatrixRotate;
 import com.android.internal.widget.remotecompose.core.operations.MatrixSave;
@@ -54,6 +55,8 @@
 import com.android.internal.widget.remotecompose.core.operations.TextMerge;
 import com.android.internal.widget.remotecompose.core.operations.Theme;
 import com.android.internal.widget.remotecompose.core.operations.utilities.IntMap;
+import com.android.internal.widget.remotecompose.core.types.BooleanConstant;
+import com.android.internal.widget.remotecompose.core.types.IntegerConstant;
 
 /**
  * List of operations supported in a RemoteCompose document
@@ -109,6 +112,9 @@
     public static final int TEXT_MERGE = 136;
     public static final int NAMED_VARIABLE = 137;
     public static final int COLOR_CONSTANT = 138;
+    public static final int DATA_INT = 140;
+    public static final int DATA_BOOLEAN = 143;
+    public static final int INTEGER_EXPRESSION = 144;
 
     /////////////////////////////////////////======================
     public static IntMap<CompanionOperation> map = new IntMap<>();
@@ -153,6 +159,9 @@
         map.put(TEXT_MERGE, TextMerge.COMPANION);
         map.put(NAMED_VARIABLE, NamedVariable.COMPANION);
         map.put(COLOR_CONSTANT, ColorConstant.COMPANION);
+        map.put(DATA_INT, IntegerConstant.COMPANION);
+        map.put(INTEGER_EXPRESSION, IntegerExpression.COMPANION);
+        map.put(DATA_BOOLEAN, BooleanConstant.COMPANION);
     }
 
 }
diff --git a/core/java/com/android/internal/widget/remotecompose/core/PaintContext.java b/core/java/com/android/internal/widget/remotecompose/core/PaintContext.java
index ecd0efc..6d8a442 100644
--- a/core/java/com/android/internal/widget/remotecompose/core/PaintContext.java
+++ b/core/java/com/android/internal/widget/remotecompose/core/PaintContext.java
@@ -59,6 +59,16 @@
 
     public abstract void drawRect(float left, float top, float right, float bottom);
 
+    /**
+     * this caches the paint to a paint stack
+     */
+    public abstract void  savePaint();
+
+    /**
+     * This restores the paint form the paint stack
+     */
+    public abstract void  restorePaint();
+
     public abstract void drawRoundRect(float left,
                                        float top,
                                        float right,
@@ -119,6 +129,10 @@
                                        float start,
                                        float stop);
 
+    /**
+     * This applies changes to the current paint
+     * @param mPaintData the list of changes
+     */
     public abstract void applyPaint(PaintBundle mPaintData);
 
     /**
diff --git a/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeBuffer.java b/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeBuffer.java
index d462c7d..f5f155e 100644
--- a/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeBuffer.java
+++ b/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeBuffer.java
@@ -37,6 +37,7 @@
 import com.android.internal.widget.remotecompose.core.operations.FloatConstant;
 import com.android.internal.widget.remotecompose.core.operations.FloatExpression;
 import com.android.internal.widget.remotecompose.core.operations.Header;
+import com.android.internal.widget.remotecompose.core.operations.IntegerExpression;
 import com.android.internal.widget.remotecompose.core.operations.MatrixRestore;
 import com.android.internal.widget.remotecompose.core.operations.MatrixRotate;
 import com.android.internal.widget.remotecompose.core.operations.MatrixSave;
@@ -55,6 +56,7 @@
 import com.android.internal.widget.remotecompose.core.operations.Utils;
 import com.android.internal.widget.remotecompose.core.operations.paint.PaintBundle;
 import com.android.internal.widget.remotecompose.core.operations.utilities.easing.FloatAnimation;
+import com.android.internal.widget.remotecompose.core.types.IntegerConstant;
 
 import java.io.File;
 import java.io.FileInputStream;
@@ -876,6 +878,27 @@
         return Utils.asNan(id);
     }
 
+
+    /**
+     * Add a Integer return an id number pointing to that float.
+     * @param value
+     * @return
+     */
+    public int addInteger(int value) {
+        int id = mRemoteComposeState.cacheInteger(value);
+        IntegerConstant.COMPANION.apply(mBuffer, id, value);
+        return id;
+    }
+
+    /**
+     * Add a IntegerId as float ID.
+     * @param id id to be converted
+     * @return
+     */
+    public float asFloatId(int id) {
+        return Utils.asNan(id);
+    }
+
     /**
      * Add a float that is a computation based on variables
      * @param value A RPN style float operation i.e. "4, 3, ADD" outputs 7
@@ -901,6 +924,18 @@
     }
 
     /**
+     * Add and integer expression
+     * @param mask defines which elements are operators or variables
+     * @param value array of values to calculate maximum 32
+     * @return
+     */
+    public int addIntegerExpression(int mask, int[] value) {
+        int id = mRemoteComposeState.cache(value);
+        IntegerExpression.COMPANION.apply(mBuffer, id, mask, value);
+        return  id;
+    }
+
+    /**
      * Add a simple color
      * @param color
      * @return id that represents that color
@@ -1038,5 +1073,6 @@
         NamedVariable.COMPANION.apply(mBuffer, id,
                 NamedVariable.COLOR_TYPE, name);
     }
+
 }
 
diff --git a/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeState.java b/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeState.java
index bfe67c8..6b06a54 100644
--- a/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeState.java
+++ b/core/java/com/android/internal/widget/remotecompose/core/RemoteComposeState.java
@@ -21,6 +21,8 @@
 import static com.android.internal.widget.remotecompose.core.RemoteContext.ID_WINDOW_HEIGHT;
 import static com.android.internal.widget.remotecompose.core.RemoteContext.ID_WINDOW_WIDTH;
 
+import com.android.internal.widget.remotecompose.core.operations.utilities.IntFloatMap;
+import com.android.internal.widget.remotecompose.core.operations.utilities.IntIntMap;
 import com.android.internal.widget.remotecompose.core.operations.utilities.IntMap;
 
 import java.util.ArrayList;
@@ -37,16 +39,12 @@
     private final IntMap<Object> mIntDataMap = new IntMap<>();
     private final IntMap<Boolean> mIntWrittenMap = new IntMap<>();
     private final HashMap<Object, Integer> mDataIntMap = new HashMap();
-    private final float[] mFloatMap = new float[MAX_FLOATS]; // efficient cache
-    private final int[] mColorMap = new int[MAX_COLORS]; // efficient cache
+    private final IntFloatMap mFloatMap = new IntFloatMap(); // efficient cache
+    private final IntIntMap mIntegerMap = new IntIntMap(); // efficient cache
+    private final IntIntMap mColorMap = new IntIntMap(); // efficient cache
     private final boolean[] mColorOverride = new boolean[MAX_COLORS];
     private int mNextId = START_ID;
 
-    {
-        for (int i = 0; i < mFloatMap.length; i++) {
-            mFloatMap[i] = Float.NaN;
-        }
-    }
 
     /**
      * Get Object based on id. The system will cache things like bitmaps
@@ -113,7 +111,18 @@
      */
     public int cacheFloat(float item) {
         int id = nextId();
-        mFloatMap[id] = item;
+        mFloatMap.put(id, item);
+        mIntegerMap.put(id, (int) item);
+        return id;
+    }
+
+    /**
+     * Insert an item in the cache
+     */
+    public int cacheInteger(int item) {
+        int id = nextId();
+        mIntegerMap.put(id, item);
+        mFloatMap.put(id, item);
         return id;
     }
 
@@ -121,21 +130,43 @@
      * Insert an item in the cache
      */
     public void cacheFloat(int id, float item) {
-        mFloatMap[id] = item;
+        mFloatMap.put(id, item);
     }
 
     /**
-     * Insert an item in the cache
+     * Insert an float item in the cache
      */
     public void updateFloat(int id, float item) {
-        mFloatMap[id] = item;
+        mFloatMap.put(id, item);
+        mIntegerMap.put(id, (int) item);
     }
 
     /**
-     * get float
+     * Insert an integer item in the cache
+     */
+    public void updateInteger(int id, int item) {
+        mFloatMap.put(id, item);
+        mIntegerMap.put(id, item);
+    }
+
+    /**
+     * get a float from the float cache
+     *
+     * @param id of the float value
+     * @return the float value
      */
     public float getFloat(int id) {
-        return mFloatMap[id];
+        return mFloatMap.get(id);
+    }
+
+    /**
+     * get an integer from the cache
+     *
+     * @param id of the integer value
+     * @return the integer
+     */
+    public int getInteger(int id) {
+        return mIntegerMap.get(id);
     }
 
     /**
@@ -145,11 +176,12 @@
      * @return
      */
     public int getColor(int id) {
-        return mColorMap[id];
+        return mColorMap.get(id);
     }
 
     /**
      * Modify the color at id.
+     *
      * @param id
      * @param color
      */
@@ -157,7 +189,7 @@
         if (mColorOverride[id]) {
             return;
         }
-        mColorMap[id] = color;
+        mColorMap.put(id, color);
     }
 
     /**
@@ -169,7 +201,7 @@
      */
     public void overrideColor(int id, int color) {
         mColorOverride[id] = true;
-        mColorMap[id] = color;
+        mColorMap.put(id, color);
     }
 
     /**
@@ -205,6 +237,7 @@
 
     /**
      * Get the next available id
+     *
      * @return
      */
     public int nextId() {
@@ -213,6 +246,7 @@
 
     /**
      * Set the next id
+     *
      * @param id
      */
     public void setNextId(int id) {
@@ -234,6 +268,7 @@
 
     /**
      * Commands that listen to variables add themselves.
+     *
      * @param id
      * @param variableSupport
      */
@@ -243,6 +278,7 @@
 
     /**
      * List of Commands that need to be updated
+     *
      * @param context
      * @return
      */
@@ -264,6 +300,7 @@
 
     /**
      * Set the width of the overall document on screen.
+     *
      * @param width
      */
     public void setWindowWidth(float width) {
@@ -272,6 +309,7 @@
 
     /**
      * Set the width of the overall document on screen.
+     *
      * @param height
      */
     public void setWindowHeight(float height) {
diff --git a/core/java/com/android/internal/widget/remotecompose/core/RemoteContext.java b/core/java/com/android/internal/widget/remotecompose/core/RemoteContext.java
index 32027d8..41eeb5b 100644
--- a/core/java/com/android/internal/widget/remotecompose/core/RemoteContext.java
+++ b/core/java/com/android/internal/widget/remotecompose/core/RemoteContext.java
@@ -39,6 +39,7 @@
 
     public float mWidth = 0f;
     public float mHeight = 0f;
+    private float mAnimationTime;
 
     /**
      * Load a path under an id.
@@ -65,11 +66,20 @@
     public abstract void loadColor(int id, int color);
 
     /**
+     * Set the animation time allowing the creator to control animation rates
+     * @param time
+     */
+    public void setAnimationTime(float time) {
+        mAnimationTime = time;
+    }
+
+    /**
      * gets the time animation clock as float in seconds
      * @return a monotonic time in seconds (arbitrary zero point)
      */
     public float getAnimationTime() {
-        return (System.nanoTime() - mStart) * 1E-9f;
+        mAnimationTime = (System.nanoTime() - mStart) * 1E-9f; // Eliminate
+        return mAnimationTime;
     }
 
     /**
@@ -213,6 +223,13 @@
     public abstract void loadFloat(int id, float value);
 
     /**
+     * Load a float
+     * @param id
+     * @param value
+     */
+    public abstract void loadInteger(int id, int value);
+
+    /**
      * Load an animated float associated with an id
      * Todo: Remove?
      * @param id
@@ -235,6 +252,13 @@
     public abstract float getFloat(int id);
 
     /**
+     * Get a float given an id
+     * @param id
+     * @return
+     */
+    public abstract int getInteger(int id);
+
+    /**
      * Get the color given and ID
      * @param id
      * @return
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/DrawBase3.java b/core/java/com/android/internal/widget/remotecompose/core/operations/DrawBase3.java
index 56b2f1f..5a4a9f3 100644
--- a/core/java/com/android/internal/widget/remotecompose/core/operations/DrawBase3.java
+++ b/core/java/com/android/internal/widget/remotecompose/core/operations/DrawBase3.java
@@ -63,23 +63,23 @@
 
     @Override
     public void updateVariables(RemoteContext context) {
-        mV1 = (Float.isNaN(mValue1))
+        mV1 = (Utils.isVariable(mValue1))
                 ? context.getFloat(Utils.idFromNan(mValue1)) : mValue1;
-        mV2 = (Float.isNaN(mValue2))
+        mV2 = (Utils.isVariable(mValue2))
                 ? context.getFloat(Utils.idFromNan(mValue2)) : mValue2;
-        mV3 = (Float.isNaN(mValue3))
+        mV3 = (Utils.isVariable(mValue3))
                 ? context.getFloat(Utils.idFromNan(mValue3)) : mValue3;
     }
 
     @Override
     public void registerListening(RemoteContext context) {
-        if (Float.isNaN(mValue1)) {
+        if (Utils.isVariable(mValue1)) {
             context.listensTo(Utils.idFromNan(mValue1), this);
         }
-        if (Float.isNaN(mValue2)) {
+        if (Utils.isVariable(mValue2)) {
             context.listensTo(Utils.idFromNan(mValue2), this);
         }
-        if (Float.isNaN(mValue3)) {
+        if (Utils.isVariable(mValue3)) {
             context.listensTo(Utils.idFromNan(mValue3), this);
         }
     }
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/DrawTextRun.java b/core/java/com/android/internal/widget/remotecompose/core/operations/DrawTextRun.java
deleted file mode 100644
index a099252..0000000
--- a/core/java/com/android/internal/widget/remotecompose/core/operations/DrawTextRun.java
+++ /dev/null
@@ -1,121 +0,0 @@
-/*
- * Copyright (C) 2024 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.internal.widget.remotecompose.core.operations;
-
-import com.android.internal.widget.remotecompose.core.CompanionOperation;
-import com.android.internal.widget.remotecompose.core.Operation;
-import com.android.internal.widget.remotecompose.core.Operations;
-import com.android.internal.widget.remotecompose.core.PaintContext;
-import com.android.internal.widget.remotecompose.core.PaintOperation;
-import com.android.internal.widget.remotecompose.core.WireBuffer;
-
-import java.util.List;
-
-public class DrawTextRun extends PaintOperation {
-    public static final Companion COMPANION = new Companion();
-    int mTextID;
-    int mStart = 0;
-    int mEnd = 0;
-    int mContextStart = 0;
-    int mContextEnd = 0;
-    float mX = 0f;
-    float mY = 0f;
-    boolean mRtl = false;
-
-    public DrawTextRun(int textID,
-                       int start,
-                       int end,
-                       int contextStart,
-                       int contextEnd,
-                       float x,
-                       float y,
-                       boolean rtl) {
-        mTextID = textID;
-        mStart = start;
-        mEnd = end;
-        mContextStart = contextStart;
-        mContextEnd = contextEnd;
-        mX = x;
-        mY = y;
-        mRtl = rtl;
-    }
-
-    @Override
-    public void write(WireBuffer buffer) {
-        COMPANION.apply(buffer, mTextID, mStart, mEnd, mContextStart, mContextEnd, mX, mY, mRtl);
-
-    }
-
-    @Override
-    public String toString() {
-        return "";
-    }
-
-    public static class Companion implements CompanionOperation {
-        private Companion() {
-        }
-
-        @Override
-        public void read(WireBuffer buffer, List<Operation> operations) {
-            int text = buffer.readInt();
-            int start = buffer.readInt();
-            int end = buffer.readInt();
-            int contextStart = buffer.readInt();
-            int contextEnd = buffer.readInt();
-            float x = buffer.readFloat();
-            float y = buffer.readFloat();
-            boolean rtl = buffer.readBoolean();
-            DrawTextRun op = new DrawTextRun(text, start, end, contextStart, contextEnd, x, y, rtl);
-
-            operations.add(op);
-        }
-
-        @Override
-        public String name() {
-            return "";
-        }
-
-        @Override
-        public int id() {
-            return 0;
-        }
-
-        public void apply(WireBuffer buffer,
-                          int textID,
-                          int start,
-                          int end,
-                          int contextStart,
-                          int contextEnd,
-                          float x,
-                          float y,
-                          boolean rtl) {
-            buffer.start(Operations.DRAW_TEXT_RUN);
-            buffer.writeInt(textID);
-            buffer.writeInt(start);
-            buffer.writeInt(end);
-            buffer.writeInt(contextStart);
-            buffer.writeInt(contextEnd);
-            buffer.writeFloat(x);
-            buffer.writeFloat(y);
-            buffer.writeBoolean(rtl);
-        }
-    }
-
-    @Override
-    public void paint(PaintContext context) {
-        context.drawTextRun(mTextID, mStart, mEnd, mContextStart, mContextEnd, mX, mY, mRtl);
-    }
-}
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/IntegerExpression.java b/core/java/com/android/internal/widget/remotecompose/core/operations/IntegerExpression.java
new file mode 100644
index 0000000..d52df5d
--- /dev/null
+++ b/core/java/com/android/internal/widget/remotecompose/core/operations/IntegerExpression.java
@@ -0,0 +1,177 @@
+/*
+ * Copyright (C) 2024 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.internal.widget.remotecompose.core.operations;
+
+import com.android.internal.widget.remotecompose.core.CompanionOperation;
+import com.android.internal.widget.remotecompose.core.Operation;
+import com.android.internal.widget.remotecompose.core.Operations;
+import com.android.internal.widget.remotecompose.core.RemoteContext;
+import com.android.internal.widget.remotecompose.core.VariableSupport;
+import com.android.internal.widget.remotecompose.core.WireBuffer;
+import com.android.internal.widget.remotecompose.core.operations.utilities.IntegerExpressionEvaluator;
+
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * Operation to deal with AnimatedFloats
+ * This is designed to be an optimized calculation for things like
+ * injecting the width of the component int draw rect
+ * As well as supporting generalized animation floats.
+ * The floats represent a RPN style calculator
+ */
+public class IntegerExpression implements Operation, VariableSupport {
+    public int mId;
+    private int mMask;
+    private int mPreMask;
+    public int[] mSrcValue;
+    public int[] mPreCalcValue;
+    private float mLastChange = Float.NaN;
+    public static final Companion COMPANION = new Companion();
+    public static final int MAX_STRING_SIZE = 4000;
+    IntegerExpressionEvaluator mExp = new IntegerExpressionEvaluator();
+
+    public IntegerExpression(int id, int mask, int[] value) {
+        this.mId = id;
+        this.mMask = mask;
+        this.mSrcValue = value;
+    }
+
+    @Override
+    public void updateVariables(RemoteContext context) {
+        if (mPreCalcValue == null || mPreCalcValue.length != mSrcValue.length) {
+            mPreCalcValue = new int[mSrcValue.length];
+        }
+        mPreMask = mMask;
+        for (int i = 0; i < mSrcValue.length; i++) {
+            if (isId(mMask, i, mSrcValue[i])) {
+                mPreMask &= ~(0x1 << i);
+                mPreCalcValue[i] = context.getInteger(mSrcValue[i]);
+            } else {
+                mPreCalcValue[i] = mSrcValue[i];
+            }
+        }
+    }
+
+
+    @Override
+    public void registerListening(RemoteContext context) {
+        for (int i = 0; i < mSrcValue.length; i++) {
+            if (isId(mMask, i, mSrcValue[i])) {
+                context.listensTo(mSrcValue[i], this);
+            }
+        }
+    }
+
+    @Override
+    public void apply(RemoteContext context) {
+        updateVariables(context);
+        float t = context.getAnimationTime();
+        if (Float.isNaN(mLastChange)) {
+            mLastChange = t;
+        }
+        int v = mExp.eval(mPreMask, Arrays.copyOf(mPreCalcValue, mPreCalcValue.length));
+        context.loadInteger(mId, v);
+    }
+
+    @Override
+    public void write(WireBuffer buffer) {
+        COMPANION.apply(buffer, mId, mMask, mSrcValue);
+    }
+
+    @Override
+    public String toString() {
+        StringBuilder s = new StringBuilder();
+        for (int i = 0; i < mPreCalcValue.length; i++) {
+            if (i != 0) {
+                s.append(" ");
+            }
+            if (IntegerExpressionEvaluator.isOperation(mMask, i)) {
+                if (isId(mMask, i, mSrcValue[i])) {
+                    s.append("[" + mSrcValue[i] + "]");
+                } else {
+                    s.append(IntegerExpressionEvaluator.toMathName(mPreCalcValue[i]));
+                }
+            } else {
+                s.append(mSrcValue[i]);
+            }
+        }
+        return "IntegerExpression[" + mId + "] = (" + s + ")";
+    }
+
+    public static class Companion implements CompanionOperation {
+        private Companion() {
+        }
+
+        @Override
+        public String name() {
+            return "FloatExpression";
+        }
+
+        @Override
+        public int id() {
+            return Operations.INTEGER_EXPRESSION;
+        }
+
+        /**
+         * Writes out the operation to the buffer
+         *
+         * @param buffer
+         * @param id
+         * @param mask
+         * @param value
+         */
+        public void apply(WireBuffer buffer, int id, int mask, int[] value) {
+            buffer.start(Operations.INTEGER_EXPRESSION);
+            buffer.writeInt(id);
+            buffer.writeInt(mask);
+            buffer.writeInt(value.length);
+            for (int i = 0; i < value.length; i++) {
+                buffer.writeInt(value[i]);
+            }
+        }
+
+        @Override
+        public void read(WireBuffer buffer, List<Operation> operations) {
+            int id = buffer.readInt();
+            int mask = buffer.readInt();
+            int len = buffer.readInt();
+
+            int[] values = new int[len];
+            for (int i = 0; i < values.length; i++) {
+                values[i] = buffer.readInt();
+            }
+
+            operations.add(new IntegerExpression(id, mask, values));
+        }
+    }
+
+    @Override
+    public String deepToString(String indent) {
+        return indent + toString();
+    }
+
+    /**
+     * given the "i" position in the mask is this an ID
+     * @param mask 32 bit mask used for defining numbers vs other
+     * @param i the bit in question
+     * @param value the value
+     * @return true if this is an ID
+     */
+    public static boolean isId(int mask, int i, int value) {
+        return ((1 << i) & mask) != 0 && value < IntegerExpressionEvaluator.OFFSET;
+    }
+}
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/Utils.java b/core/java/com/android/internal/widget/remotecompose/core/operations/Utils.java
index fcb3bfa..e9b0c3b 100644
--- a/core/java/com/android/internal/widget/remotecompose/core/operations/Utils.java
+++ b/core/java/com/android/internal/widget/remotecompose/core/operations/Utils.java
@@ -40,7 +40,7 @@
      * @param n
      * @return
      */
-    static String trimString(String str, int n) {
+    public static String trimString(String str, int n) {
         if (str.length() > n) {
             str = str.substring(0, n - 3) + "...";
         }
@@ -55,6 +55,9 @@
      */
     public static String floatToString(float idvalue, float value) {
         if (Float.isNaN(idvalue)) {
+            if (idFromNan(value) == 0) {
+                return "NaN";
+            }
             return "[" + idFromNan(idvalue) + "]" + floatToString(value);
         }
         return floatToString(value);
@@ -67,6 +70,9 @@
      */
     public static String floatToString(float value) {
         if (Float.isNaN(value)) {
+            if (idFromNan(value) == 0) {
+                return "NaN";
+            }
             return "[" + idFromNan(value) + "]";
         }
         return Float.toString(value);
@@ -107,6 +113,7 @@
     public static boolean isVariable(float v) {
         if (Float.isNaN(v)) {
             int id = idFromNan(v);
+            if (id == 0) return false;
             return id > 40 || id < 10;
         }
         return false;
@@ -222,6 +229,4 @@
         }
         return 0;
     }
-
-
 }
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/paint/Painter.java b/core/java/com/android/internal/widget/remotecompose/core/operations/paint/Painter.java
new file mode 100644
index 0000000..ada3757
--- /dev/null
+++ b/core/java/com/android/internal/widget/remotecompose/core/operations/paint/Painter.java
@@ -0,0 +1,260 @@
+/*
+ * Copyright (C) 2023 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.internal.widget.remotecompose.core.operations.paint;
+
+
+/**
+ * Provides a Builder pattern for a PaintBundle
+ */
+class Painter {
+    PaintBundle mPaint;
+
+    /**
+     * Write the paint to the buffer
+     */
+    public PaintBundle commit() {
+        return mPaint;
+    }
+
+    public Painter setAntiAlias(boolean aa) {
+        mPaint.setAntiAlias(aa);
+        return this;
+    }
+
+    public Painter setColor(int color) {
+        mPaint.setColor(color);
+        return this;
+    }
+
+    public Painter setColorId(int colorId) {
+        mPaint.setColorId(colorId);
+        return this;
+    }
+
+    /**
+     * Set the paint's Join.
+     *
+     * @param join set the paint's Join, used whenever the paint's style is
+     *             Stroke or StrokeAndFill.
+     */
+    public Painter setStrokeJoin(int join) {
+        mPaint.setStrokeJoin(join);
+        return this;
+    }
+
+    /**
+     * Set the width for stroking. Pass 0 to stroke in hairline mode.
+     * Hairlines always draws a single
+     * pixel independent of the canvas's matrix.
+     *
+     * @param width set the paint's stroke width, used whenever the paint's
+     *             style is Stroke or StrokeAndFill.
+     */
+    public Painter setStrokeWidth(float width) {
+        mPaint.setStrokeWidth(width);
+        return this;
+    }
+
+    /**
+     * Set the paint's style, used for controlling how primitives' geometries
+     * are interpreted (except for drawBitmap, which always assumes Fill).
+     *
+     * @param style The new style to set in the paint
+     */
+    public Painter setStyle(int style) {
+        mPaint.setStyle(style);
+        return this;
+    }
+
+    /**
+     * Set the paint's Cap.
+     *
+     * @param cap set the paint's line cap style, used whenever the paint's
+     *           style is Stroke or StrokeAndFill.
+     */
+    public Painter setStrokeCap(int cap) {
+        mPaint.setStrokeCap(cap);
+        return this;
+    }
+
+    /**
+     * Set the paint's stroke miter value. This is used to control the behavior
+     * of miter joins when the joins angle is sharp. This value must be >= 0.
+     *
+     * @param miter set the miter limit on the paint, used whenever the paint's
+     *             style is Stroke or StrokeAndFill.
+     */
+    public Painter setStrokeMiter(float miter) {
+        mPaint.setStrokeMiter(miter);
+        return this;
+    }
+
+    /**
+     * Helper to setColor(), that only assigns the color's alpha value,
+     * leaving its r,g,b values unchanged. Results are undefined if the alpha
+     * value is outside of the range [0..1.0]
+     *
+     * @param alpha set the alpha component [0..1.0] of the paint's color.
+     */
+    public Painter setAlpha(float alpha) {
+        mPaint.setAlpha((alpha > 2) ? alpha / 255f : alpha);
+        return this;
+    }
+
+    /**
+     * Create a color filter that uses the specified color and Porter-Duff mode.
+     *
+     * @param color The ARGB source color used with the specified Porter-Duff
+     *             mode
+     * @param mode  The porter-duff mode that is applied
+     */
+    public Painter setPorterDuffColorFilter(int color, int mode) {
+        mPaint.setColorFilter(color, mode);
+        return this;
+    }
+
+    /**
+     * sets a shader that draws a linear gradient along a line.
+     *
+     * @param startX    The x-coordinate for the start of the gradient line
+     * @param startY    The y-coordinate for the start of the gradient line
+     * @param endX      The x-coordinate for the end of the gradient line
+     * @param endY      The y-coordinate for the end of the gradient line
+     * @param colors    The sRGB colors to be distributed along the gradient
+     *                  line
+     * @param positions May be null. The relative positions [0..1] of each
+     *                 corresponding color in the colors array. If this is null,
+     *                 the colors are distributed evenly along the gradient
+     *                 line.
+     * @param tileMode  The Shader tiling mode
+     */
+    public Painter setLinearGradient(
+            float startX,
+            float startY,
+            float endX,
+            float endY,
+            int[] colors,
+            float[] positions,
+            int tileMode
+    ) {
+        mPaint.setLinearGradient(colors, positions, startX,
+                startY, endX, endY, tileMode);
+        return this;
+    }
+
+    /**
+     * Sets a shader that draws a radial gradient given the center and radius.
+     *
+     * @param centerX   The x-coordinate of the center of the radius
+     * @param centerY   The y-coordinate of the center of the radius
+     * @param radius    Must be positive. The radius of the circle for this
+     *                  gradient.
+     * @param colors    The sRGB colors to be distributed between the center
+     *                  and edge of the circle
+     * @param positions May be <code>null</code>. Valid values are between
+     *                  <code>0.0f</code> and
+     *                  <code>1.0f</code>. The relative position of each
+     *                  corresponding color in the colors array. If
+     *                  <code>null</code>, colors are distributed evenly
+     *                  between the center and edge of the circle.
+     * @param tileMode  The Shader tiling mode
+     */
+    public Painter setRadialGradient(
+            float centerX,
+            float centerY,
+            float radius,
+            int[] colors,
+            float[] positions,
+            int tileMode
+    ) {
+        mPaint.setRadialGradient(colors, positions, centerX,
+                centerY, radius, tileMode);
+        return this;
+    }
+
+    /**
+     * Set a shader that draws a sweep gradient around a center point.
+     *
+     * @param centerX   The x-coordinate of the center
+     * @param centerY   The y-coordinate of the center
+     * @param colors    The sRGB colors to be distributed between around the
+     *                  center. There must be at least 2 colors in the array.
+     * @param positions May be NULL. The relative position of each corresponding
+     *                 color in the colors array, beginning with 0 and ending
+     *                 with 1.0. If the values are not monotonic, the drawing
+     *                  may produce unexpected results. If positions is NULL,
+     *                  then the colors are automatically spaced evenly.
+     */
+    public Painter setSweepGradient(
+            float centerX,
+            float centerY,
+            int[] colors,
+            float[] positions
+    ) {
+        mPaint.setSweepGradient(colors, positions, centerX, centerY);
+        return this;
+    }
+
+    /**
+     * Set the paint's text size. This value must be > 0
+     *
+     * @param size set the paint's text size in pixel units.
+     */
+    public Painter setTextSize(float size) {
+        mPaint.setTextSize(size);
+        return this;
+    }
+
+    /**
+     * sets a typeface object that best matches the specified existing
+     * typeface and the specified weight and italic style
+     *
+     * <p>Below are numerical values and corresponding common weight names.</p>
+     * <table> <thead>
+     * <tr><th>Value</th><th>Common weight name</th></tr> </thead> <tbody>
+     * <tr><td>100</td><td>Thin</td></tr>
+     * <tr><td>200</td><td>Extra Light</td></tr>
+     * <tr><td>300</td><td>Light</td></tr>
+     * <tr><td>400</td><td>Normal</td></tr>
+     * <tr><td>500</td><td>Medium</td></tr>
+     * <tr><td>600</td><td>Semi Bold</td></tr>
+     * <tr><td>700</td><td>Bold</td></tr>
+     * <tr><td>800</td><td>Extra Bold</td></tr>
+     * <tr><td>900</td><td>Black</td></tr> </tbody> </table>
+     *
+     * @param fontType 0 = default 1 = sans serif 2 = serif 3 = monospace
+     * @param weight   The desired weight to be drawn.
+     * @param italic   {@code true} if italic style is desired to be drawn.
+     *                            Otherwise, {@code false}
+     */
+    public Painter setTypeface(int fontType, int weight, boolean italic) {
+        mPaint.setTextStyle(fontType, weight, italic);
+        return this;
+    }
+
+
+    public Painter setFilterBitmap(boolean filter) {
+        mPaint.setFilterBitmap(filter);
+        return this;
+    }
+
+
+    public Painter setShader(int id) {
+        mPaint.setShader(id);
+        return this;
+    }
+
+}
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntFloatMap.java b/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntFloatMap.java
new file mode 100644
index 0000000..23c3ec5
--- /dev/null
+++ b/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntFloatMap.java
@@ -0,0 +1,140 @@
+/*
+ * Copyright (C) 2023 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.internal.widget.remotecompose.core.operations.utilities;
+
+import java.util.Arrays;
+
+public class IntFloatMap {
+
+    private static final int DEFAULT_CAPACITY = 16;
+    private static final float LOAD_FACTOR = 0.75f;
+    private static final int NOT_PRESENT = Integer.MIN_VALUE;
+    private int[] mKeys;
+    private float[] mValues;
+    int mSize;
+
+    public IntFloatMap() {
+        mKeys = new int[DEFAULT_CAPACITY];
+        Arrays.fill(mKeys, NOT_PRESENT);
+        mValues = new float[DEFAULT_CAPACITY];
+    }
+
+    /**
+     * clear the map
+     */
+    public void clear() {
+        Arrays.fill(mKeys, NOT_PRESENT);
+        Arrays.fill(mValues, Float.NaN); // not strictly necessary but defensive
+        mSize = 0;
+    }
+
+    /**
+     * is the key contained in map
+     *
+     * @param key the key to check
+     * @return true if the map contains the key
+     */
+    public boolean contains(int key) {
+        return findKey(key) != -1;
+    }
+
+    /**
+     * Put a item in the map
+     *
+     * @param key   item'values key
+     * @param value item's value
+     * @return old value if exist
+     */
+    public float put(int key, float value) {
+        if (key == NOT_PRESENT) {
+            throw new IllegalArgumentException("Key cannot be NOT_PRESENT");
+        }
+        if (mSize > mKeys.length * LOAD_FACTOR) {
+            resize();
+        }
+        return insert(key, value);
+    }
+
+    /**
+     * get an element given the key
+     *
+     * @param key the key to fetch
+     * @return the value
+     */
+    public float get(int key) {
+        int index = findKey(key);
+        if (index == -1) {
+            return 0;
+        } else
+            return mValues[index];
+    }
+
+    /**
+     * how many elements in the map
+     *
+     * @return number of elements
+     */
+    public int size() {
+        return mSize;
+    }
+
+    private float insert(int key, float value) {
+        int index = hash(key) % mKeys.length;
+        while (mKeys[index] != NOT_PRESENT && mKeys[index] != key) {
+            index = (index + 1) % mKeys.length;
+        }
+        float oldValue = 0;
+        if (mKeys[index] == NOT_PRESENT) {
+            mSize++;
+        } else {
+            oldValue = mValues[index];
+        }
+        mKeys[index] = key;
+        mValues[index] = value;
+        return oldValue;
+    }
+
+    private int findKey(int key) {
+        int index = hash(key) % mKeys.length;
+        while (mKeys[index] != NOT_PRESENT) {
+            if (mKeys[index] == key) {
+                return index;
+            }
+            index = (index + 1) % mKeys.length;
+        }
+        return -1;
+    }
+
+    private int hash(int key) {
+        return key;
+    }
+
+    private void resize() {
+        int[] oldKeys = mKeys;
+        float[] oldValues = mValues;
+        mKeys = new int[(oldKeys.length * 2)];
+        for (int i = 0; i < mKeys.length; i++) {
+            mKeys[i] = NOT_PRESENT;
+        }
+        mValues = new float[oldKeys.length * 2];
+        mSize = 0;
+        for (int i = 0; i < oldKeys.length; i++) {
+            if (oldKeys[i] != NOT_PRESENT) {
+                put(oldKeys[i], oldValues[i]);
+            }
+        }
+    }
+}
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntIntMap.java b/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntIntMap.java
new file mode 100644
index 0000000..221014c
--- /dev/null
+++ b/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntIntMap.java
@@ -0,0 +1,139 @@
+/*
+ * Copyright (C) 2023 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.internal.widget.remotecompose.core.operations.utilities;
+
+import java.util.Arrays;
+
+public class IntIntMap {
+    private static final int DEFAULT_CAPACITY = 16;
+    private static final float LOAD_FACTOR = 0.75f;
+    private static final int NOT_PRESENT = Integer.MIN_VALUE;
+    private int[] mKeys;
+    private int[] mValues;
+    int mSize;
+
+    public IntIntMap() {
+        mKeys = new int[DEFAULT_CAPACITY];
+        Arrays.fill(mKeys, NOT_PRESENT);
+        mValues = new int[DEFAULT_CAPACITY];
+    }
+
+    /**
+     * clear the map
+     */
+    public void clear() {
+        Arrays.fill(mKeys, NOT_PRESENT);
+        Arrays.fill(mValues, 0);
+        mSize = 0;
+    }
+
+    /**
+     * is the key contained in map
+     *
+     * @param key the key to check
+     * @return true if the map contains the key
+     */
+    public boolean contains(int key) {
+        return findKey(key) != -1;
+    }
+
+    /**
+     * Put a item in the map
+     *
+     * @param key   item'values key
+     * @param value item's value
+     * @return old value if exist
+     */
+    public int put(int key, int value) {
+        if (key == NOT_PRESENT) {
+            throw new IllegalArgumentException("Key cannot be NOT_PRESENT");
+        }
+        if (mSize > mKeys.length * LOAD_FACTOR) {
+            resize();
+        }
+        return insert(key, value);
+    }
+
+    /**
+     * get an element given the key
+     *
+     * @param key the key to fetch
+     * @return the value
+     */
+    public int get(int key) {
+        int index = findKey(key);
+        if (index == -1) {
+            return 0;
+        } else
+            return mValues[index];
+    }
+
+    /**
+     * how many elements in the map
+     *
+     * @return number of elements
+     */
+    public int size() {
+        return mSize;
+    }
+
+    private int insert(int key, int value) {
+        int index = hash(key) % mKeys.length;
+        while (mKeys[index] != NOT_PRESENT && mKeys[index] != key) {
+            index = (index + 1) % mKeys.length;
+        }
+        int oldValue = 0;
+        if (mKeys[index] == NOT_PRESENT) {
+            mSize++;
+        } else {
+            oldValue = mValues[index];
+        }
+        mKeys[index] = key;
+        mValues[index] = value;
+        return oldValue;
+    }
+
+    private int findKey(int key) {
+        int index = hash(key) % mKeys.length;
+        while (mKeys[index] != NOT_PRESENT) {
+            if (mKeys[index] == key) {
+                return index;
+            }
+            index = (index + 1) % mKeys.length;
+        }
+        return -1;
+    }
+
+    private int hash(int key) {
+        return key;
+    }
+
+    private void resize() {
+        int[] oldKeys = mKeys;
+        int[] oldValues = mValues;
+        mKeys = new int[(oldKeys.length * 2)];
+        for (int i = 0; i < mKeys.length; i++) {
+            mKeys[i] = NOT_PRESENT;
+        }
+        mValues = new int[oldKeys.length * 2];
+        mSize = 0;
+        for (int i = 0; i < oldKeys.length; i++) {
+            if (oldKeys[i] != NOT_PRESENT) {
+                put(oldKeys[i], oldValues[i]);
+            }
+        }
+    }
+}
diff --git a/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntegerExpressionEvaluator.java b/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntegerExpressionEvaluator.java
new file mode 100644
index 0000000..4c1389c
--- /dev/null
+++ b/core/java/com/android/internal/widget/remotecompose/core/operations/utilities/IntegerExpressionEvaluator.java
@@ -0,0 +1,429 @@
+/*
+ * Copyright (C) 2024 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.internal.widget.remotecompose.core.operations.utilities;
+
+/**
+ * High performance Integer expression evaluator
+ */
+public class IntegerExpressionEvaluator {
+    static IntMap<String> sNames = new IntMap<>();
+    public static final int OFFSET = 0x10000;
+    // add, sub, mul,div,mod,min,max, shl, shr, ushr, OR, AND , XOR, COPY_SIGN
+    public static final int I_ADD = OFFSET + 1;
+    public static final int I_SUB = OFFSET + 2;
+    public static final int I_MUL = OFFSET + 3;
+    public static final int I_DIV = OFFSET + 4;
+    public static final int I_MOD = OFFSET + 5;
+    public static final int I_SHL = OFFSET + 6;
+    public static final int I_SHR = OFFSET + 7;
+    public static final int I_USHR = OFFSET + 8;
+    public static final int I_OR = OFFSET + 9;
+    public static final int I_AND = OFFSET + 10;
+    public static final int I_XOR = OFFSET + 11;
+    public static final int I_COPY_SIGN = OFFSET + 12;
+    public static final int I_MIN = OFFSET + 13;
+    public static final int I_MAX = OFFSET + 14;
+
+    public static final int I_NEG = OFFSET + 15;
+    public static final int I_ABS = OFFSET + 16;
+    public static final int I_INCR = OFFSET + 17;
+    public static final int I_DECR = OFFSET + 18;
+    public static final int I_NOT = OFFSET + 19;
+    public static final int I_SIGN = OFFSET + 20;
+
+    public static final int I_CLAMP = OFFSET + 21;
+    public static final int I_IFELSE = OFFSET + 22;
+    public static final int I_MAD = OFFSET + 23;
+
+    public static final float LAST_OP = 24;
+
+    public static final int I_VAR1 = OFFSET + 24;
+    public static final int I_VAR2 = OFFSET + 24;
+
+
+    int[] mStack;
+    int[] mLocalStack = new int[128];
+    int[] mVar;
+
+
+    interface Op {
+        int eval(int sp);
+    }
+
+    /**
+     * Evaluate a float expression
+     *
+     * @param exp
+     * @param var
+     * @return
+     */
+    public int eval(int mask, int[] exp, int... var) {
+        mStack = exp;
+        mVar = var;
+        int sp = -1;
+        for (int i = 0; i < mStack.length; i++) {
+            int v = mStack[i];
+            if (((1 << i) & mask) != 0) {
+                sp = mOps[v - OFFSET].eval(sp);
+            } else {
+                mStack[++sp] = v;
+            }
+        }
+        return mStack[sp];
+    }
+
+    /**
+     * Evaluate a int expression
+     *
+     * @param exp
+     * @param len
+     * @param var
+     * @return
+     */
+    public int eval(int mask, int[] exp, int len, int... var) {
+        System.arraycopy(exp, 0, mLocalStack, 0, len);
+        mStack = mLocalStack;
+        mVar = var;
+        int sp = -1;
+        for (int i = 0; i < len; i++) {
+            int v = mStack[i];
+            if (((1 << i) & mask) != 0) {
+                sp = mOps[v - OFFSET].eval(sp);
+            } else {
+                mStack[++sp] = v;
+            }
+        }
+        return mStack[sp];
+    }
+
+    /**
+     * Evaluate a int expression
+     *
+     * @param exp
+     * @param var
+     * @return
+     */
+    public int evalDB(int mask, int[] exp, int... var) {
+        mStack = exp;
+        mVar = var;
+        int sp = -1;
+        for (int i = 0; i < exp.length; i++) {
+            int v = mStack[i];
+            if (((1 << i) & mask) != 0) {
+                System.out.print(" " + sNames.get((v - OFFSET)));
+                sp = mOps[v - OFFSET].eval(sp);
+            } else {
+                System.out.print(" " + v);
+                mStack[++sp] = v;
+            }
+        }
+        return mStack[sp];
+    }
+
+    Op[] mOps = {
+            null,
+            (sp) -> { // ADD
+                mStack[sp - 1] = mStack[sp - 1] + mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // SUB
+                mStack[sp - 1] = mStack[sp - 1] - mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // MUL
+                mStack[sp - 1] = mStack[sp - 1] * mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> {  // DIV
+                mStack[sp - 1] = mStack[sp - 1] / mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> {  // MOD
+                mStack[sp - 1] = mStack[sp - 1] % mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // SHL shift left
+                mStack[sp - 1] = mStack[sp - 1] << mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // SHR shift right
+                mStack[sp - 1] = mStack[sp - 1] >> mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // USHR unsigned shift right
+                mStack[sp - 1] = mStack[sp - 1] >>> mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // OR operator
+                mStack[sp - 1] = mStack[sp - 1] | mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // AND operator
+                mStack[sp - 1] = mStack[sp - 1] & mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // XOR xor operator
+                mStack[sp - 1] = mStack[sp - 1] ^ mStack[sp];
+                return sp - 1;
+            },
+            (sp) -> { // COPY_SIGN copy the sing of (using bit magic)
+                mStack[sp - 1] = (mStack[sp - 1] ^ (mStack[sp] >> 31))
+                        - (mStack[sp] >> 31);
+                return sp - 1;
+            },
+            (sp) -> { // MIN
+                mStack[sp - 1] = Math.min(mStack[sp - 1], mStack[sp]);
+                return sp - 1;
+            },
+            (sp) -> { // MAX
+                mStack[sp - 1] = Math.max(mStack[sp - 1], mStack[sp]);
+                return sp - 1;
+            },
+            (sp) -> { // NEG
+                mStack[sp] = -mStack[sp];
+                return sp;
+            },
+            (sp) -> { // ABS
+                mStack[sp] = Math.abs(mStack[sp]);
+                return sp;
+            },
+            (sp) -> { // INCR increment
+                mStack[sp] = mStack[sp] + 1;
+                return sp;
+            },
+            (sp) -> { // DECR decrement
+                mStack[sp] = mStack[sp] - 1;
+                return sp;
+            },
+            (sp) -> { // NOT Bit invert
+                mStack[sp] = ~mStack[sp];
+                return sp;
+            },
+            (sp) -> { // SIGN x<0 = -1,x==0 =  0 , x>0 = 1
+                mStack[sp] = (mStack[sp] >> 31) | (-mStack[sp] >>> 31);
+                return sp;
+            },
+
+            (sp) -> { // CLAMP(min,max, val)
+                mStack[sp - 2] = Math.min(Math.max(mStack[sp - 2], mStack[sp]),
+                        mStack[sp - 1]);
+                return sp - 2;
+            },
+            (sp) -> { // Ternary conditional
+                mStack[sp - 2] = (mStack[sp] > 0)
+                        ? mStack[sp - 1] : mStack[sp - 2];
+                return sp - 2;
+            },
+            (sp) -> { // MAD
+                mStack[sp - 2] = mStack[sp] + mStack[sp - 1] * mStack[sp - 2];
+                return sp - 2;
+            },
+
+            (sp) -> { // first var =
+                mStack[sp] = mVar[0];
+                return sp;
+            },
+            (sp) -> { // second var y?
+                mStack[sp] = mVar[1];
+                return sp;
+            },
+            (sp) -> { // 3rd var z?
+                mStack[sp] = mVar[2];
+                return sp;
+            },
+    };
+
+    static {
+        int k = 0;
+        sNames.put(k++, "NOP");
+        sNames.put(k++, "+");
+        sNames.put(k++, "-");
+        sNames.put(k++, "*");
+        sNames.put(k++, "/");
+        sNames.put(k++, "%");
+        sNames.put(k++, "<<");
+        sNames.put(k++, ">>");
+        sNames.put(k++, ">>>");
+        sNames.put(k++, "|");
+        sNames.put(k++, "&");
+        sNames.put(k++, "^");
+        sNames.put(k++, "copySign");
+        sNames.put(k++, "min");
+        sNames.put(k++, "max");
+        sNames.put(k++, "neg");
+        sNames.put(k++, "abs");
+        sNames.put(k++, "incr");
+        sNames.put(k++, "decr");
+        sNames.put(k++, "not");
+        sNames.put(k++, "sign");
+        sNames.put(k++, "clamp");
+        sNames.put(k++, "ifElse");
+        sNames.put(k++, "mad");
+        sNames.put(k++, "ceil");
+        sNames.put(k++, "a[0]");
+        sNames.put(k++, "a[1]");
+        sNames.put(k++, "a[2]");
+    }
+
+    /**
+     * given a int command return its math name (e.g sin, cos etc.)
+     *
+     * @param f
+     * @return
+     */
+    public static String toMathName(int f) {
+        int id = f - OFFSET;
+        return sNames.get(id);
+    }
+
+    /**
+     * Convert an expression encoded as an array of ints int ot a string
+     *
+     * @param exp
+     * @param labels
+     * @return
+     */
+    public static String toString(int mask, int[] exp, String[] labels) {
+        StringBuilder s = new StringBuilder();
+        for (int i = 0; i < exp.length; i++) {
+            int v = exp[i];
+
+            if (((1 << i) & mask) != 0) {
+                if (v < OFFSET) {
+                    s.append(toMathName(v));
+                } else {
+                    s.append("[");
+                    s.append(v);
+                    s.append("]");
+                }
+            } else {
+                if (labels[i] != null) {
+                    s.append(labels[i]);
+                }
+                s.append(v);
+            }
+            s.append(" ");
+        }
+        return s.toString();
+    }
+
+    /**
+     * Convert an expression encoded as an array of ints int ot a string
+     *
+     * @param mask bit mask of operators vs commands
+     * @param exp
+     * @return
+     */
+    public static String toString(int mask, int[] exp) {
+        StringBuilder s = new StringBuilder();
+        s.append(Integer.toBinaryString(mask));
+        s.append(" : ");
+        for (int i = 0; i < exp.length; i++) {
+            int v = exp[i];
+
+            if (((1 << i) & mask) != 0) {
+                if (v > OFFSET) {
+                    s.append(" ");
+                    s.append(toMathName(v));
+                    s.append(" ");
+
+                } else {
+                    s.append("[");
+                    s.append(v);
+                    s.append("]");
+                }
+            }
+            s.append(" " + v);
+        }
+        return s.toString();
+    }
+
+    /**
+     * This creates an infix string expression
+     * @param mask The bits that are operators
+     * @param exp the array of expressions
+     * @return infix string
+     */
+    public static String toStringInfix(int mask, int[] exp) {
+        return toString(mask, exp, exp.length - 1);
+    }
+
+    static String toString(int mask, int[] exp, int sp) {
+        String[] str = new String[exp.length];
+        if (((1 << sp) & mask) != 0) {
+            int id = exp[sp] - OFFSET;
+            switch (NO_OF_OPS[id]) {
+                case -1:
+                    return "nop";
+                case 1:
+                    return sNames.get(id) + "(" + toString(mask, exp, sp - 1) + ") ";
+                case 2:
+                    if (infix(id)) {
+                        return "(" + toString(mask, exp, sp - 2)
+                                + " " + sNames.get(id) + " "
+                                + toString(mask, exp, sp - 1) + ") ";
+                    } else {
+                        return sNames.get(id) + "("
+                                + toString(mask, exp, sp - 2) + ", "
+                                + toString(mask, exp, sp - 1) + ")";
+                    }
+                case 3:
+                    if (infix(id)) {
+                        return "((" + toString(mask, exp, sp + 3) + ") ? "
+                                + toString(mask, exp, sp - 2) + ":"
+                                + toString(mask, exp, sp - 1) + ")";
+                    } else {
+                        return sNames.get(id)
+                                + "(" + toString(mask, exp, sp - 3)
+                                + ", " + toString(mask, exp, sp - 2)
+                                + ", " + toString(mask, exp, sp - 1) + ")";
+                    }
+            }
+        }
+        return Integer.toString(exp[sp]);
+    }
+
+    static final int[] NO_OF_OPS = {
+            -1, // no op
+            2, 2, 2, 2, 2, // + - * / %
+            2, 2, 2, 2, 2, 2, 2, 2, 2, //<<, >> , >>> , | , &, ^, min max
+            1, 1, 1, 1, 1, 1,   // neg, abs, ++, -- , not , sign
+
+            3, 3, 3, // clamp, ifElse, mad,
+            0, 0, 0 // mad, ?:,
+            // a[0],a[1],a[2]
+    };
+
+    /**
+     * to be used by parser to determine if command is infix
+     *
+     * @param n the operator (minus the offset)
+     * @return true if the operator is infix
+     */
+    static boolean infix(int n) {
+        return ((n < 12));
+    }
+
+    /**
+     * is it an id or operation
+     * @param mask the bits that mark elements as an operation
+     * @param i the bit to check
+     * @return true if the bit is 1
+     */
+    public static boolean isOperation(int mask, int i) {
+        return ((1 << i) & mask) != 0;
+    }
+}
diff --git a/core/java/com/android/internal/widget/remotecompose/core/types/BooleanConstant.java b/core/java/com/android/internal/widget/remotecompose/core/types/BooleanConstant.java
new file mode 100644
index 0000000..1051192
--- /dev/null
+++ b/core/java/com/android/internal/widget/remotecompose/core/types/BooleanConstant.java
@@ -0,0 +1,96 @@
+/*
+ * Copyright (C) 2024 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.internal.widget.remotecompose.core.types;
+
+import com.android.internal.widget.remotecompose.core.CompanionOperation;
+import com.android.internal.widget.remotecompose.core.Operation;
+import com.android.internal.widget.remotecompose.core.Operations;
+import com.android.internal.widget.remotecompose.core.RemoteContext;
+import com.android.internal.widget.remotecompose.core.WireBuffer;
+
+import java.util.List;
+
+/**
+ * Used to represent a boolean
+ */
+public class BooleanConstant implements Operation {
+    boolean mValue = false;
+    private int mId;
+
+    public BooleanConstant(int id, boolean value) {
+        mId = id;
+        mValue = value;
+    }
+
+    @Override
+    public void write(WireBuffer buffer) {
+        COMPANION.apply(buffer, mId, mValue);
+    }
+
+    @Override
+    public void apply(RemoteContext context) {
+
+    }
+
+    @Override
+    public String deepToString(String indent) {
+        return toString();
+    }
+
+    @Override
+    public String toString() {
+        return "BooleanConstant[" + mId + "] = " + mValue + "";
+    }
+
+    public static final Companion COMPANION = new Companion();
+
+    public static class Companion implements CompanionOperation {
+        private Companion() {
+        }
+
+        @Override
+        public String name() {
+            return "OrigamiBoolean";
+        }
+
+        @Override
+        public int id() {
+            return Operations.DATA_BOOLEAN;
+        }
+
+        /**
+         * Writes out the operation to the buffer
+         *
+         * @param buffer
+         * @param id
+         * @param value
+         */
+        public void apply(WireBuffer buffer, int id, boolean value) {
+            buffer.start(Operations.DATA_BOOLEAN);
+            buffer.writeInt(id);
+            buffer.writeBoolean(value);
+        }
+
+        @Override
+        public void read(WireBuffer buffer, List<Operation> operations) {
+            int id = buffer.readInt();
+
+            boolean value = buffer.readBoolean();
+            operations.add(new BooleanConstant(id, value));
+        }
+    }
+
+}
diff --git a/core/java/com/android/internal/widget/remotecompose/core/types/IntegerConstant.java b/core/java/com/android/internal/widget/remotecompose/core/types/IntegerConstant.java
new file mode 100644
index 0000000..ceb3236
--- /dev/null
+++ b/core/java/com/android/internal/widget/remotecompose/core/types/IntegerConstant.java
@@ -0,0 +1,96 @@
+/*
+ * Copyright (C) 2024 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.internal.widget.remotecompose.core.types;
+
+import com.android.internal.widget.remotecompose.core.CompanionOperation;
+import com.android.internal.widget.remotecompose.core.Operation;
+import com.android.internal.widget.remotecompose.core.Operations;
+import com.android.internal.widget.remotecompose.core.RemoteContext;
+import com.android.internal.widget.remotecompose.core.WireBuffer;
+
+import java.util.List;
+
+/**
+ * Represents a single integer typically used for states
+ * or named for input into the system
+ */
+public class IntegerConstant implements Operation {
+    private int mValue = 0;
+    private int mId;
+
+    IntegerConstant(int id, int value) {
+        mId = id;
+        mValue = value;
+    }
+
+    @Override
+    public void write(WireBuffer buffer) {
+        COMPANION.apply(buffer, mId, mValue);
+    }
+
+    @Override
+    public void apply(RemoteContext context) {
+        context.loadInteger(mId, mValue);
+    }
+
+    @Override
+    public String deepToString(String indent) {
+        return toString();
+    }
+
+    @Override
+    public String toString() {
+        return "IntegerConstant[" + mId + "] = " + mValue + "";
+    }
+
+    public static final Companion COMPANION = new Companion();
+
+    public static class Companion implements CompanionOperation {
+        private Companion() {
+        }
+
+        @Override
+        public String name() {
+            return "IntegerConstant";
+        }
+
+        @Override
+        public int id() {
+            return Operations.DATA_INT;
+        }
+
+        /**
+         * Writes out the operation to the buffer
+         *
+         * @param buffer
+         * @param textId
+         * @param value
+         */
+        public void apply(WireBuffer buffer, int textId, int value) {
+            buffer.start(Operations.DATA_INT);
+            buffer.writeInt(textId);
+            buffer.writeInt(value);
+        }
+
+        @Override
+        public void read(WireBuffer buffer, List<Operation> operations) {
+            int id = buffer.readInt();
+
+            int value = buffer.readInt();
+            operations.add(new IntegerConstant(id, value));
+        }
+    }
+}
diff --git a/core/java/com/android/internal/widget/remotecompose/player/RemoteComposePlayer.java b/core/java/com/android/internal/widget/remotecompose/player/RemoteComposePlayer.java
index 73e94fa..b2406bf 100644
--- a/core/java/com/android/internal/widget/remotecompose/player/RemoteComposePlayer.java
+++ b/core/java/com/android/internal/widget/remotecompose/player/RemoteComposePlayer.java
@@ -410,4 +410,3 @@
         }
     }
 }
-
diff --git a/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidPaintContext.java b/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidPaintContext.java
index ecb68bb..39a770a 100644
--- a/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidPaintContext.java
+++ b/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidPaintContext.java
@@ -39,12 +39,16 @@
 import com.android.internal.widget.remotecompose.core.operations.paint.PaintBundle;
 import com.android.internal.widget.remotecompose.core.operations.paint.PaintChanges;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * An implementation of PaintContext for the Android Canvas.
  * This is used to play the RemoteCompose operations on Android.
  */
 public class AndroidPaintContext extends PaintContext {
     Paint mPaint = new Paint();
+    List<Paint> mPaintList = new ArrayList<>();
     Canvas mCanvas;
     Rect mTmpRect = new Rect(); // use in calculation of bounds
 
@@ -162,6 +166,16 @@
     }
 
     @Override
+    public void savePaint() {
+        mPaintList.add(new Paint(mPaint));
+    }
+
+    @Override
+    public void restorePaint() {
+        mPaint = mPaintList.remove(mPaintList.size() - 1);
+    }
+
+    @Override
     public void drawRoundRect(float left,
                               float top,
                               float right,
@@ -335,6 +349,11 @@
         return null;
     }
 
+    /**
+     * This applies paint changes to the current paint
+     *
+     * @param mPaintData the list change to the paint
+     */
     @Override
     public void applyPaint(PaintBundle mPaintData) {
         mPaintData.applyPaintChange((PaintContext) this, new PaintChanges() {
diff --git a/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidRemoteContext.java b/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidRemoteContext.java
index dd43bd5..5a87c70 100644
--- a/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidRemoteContext.java
+++ b/core/java/com/android/internal/widget/remotecompose/player/platform/AndroidRemoteContext.java
@@ -120,6 +120,11 @@
         mRemoteComposeState.updateFloat(id, value);
     }
 
+    @Override
+    public void loadInteger(int id, int value) {
+        mRemoteComposeState.updateInteger(id, value);
+    }
+
 
     @Override
     public void loadColor(int id, int color) {
@@ -142,6 +147,11 @@
     }
 
     @Override
+    public int getInteger(int id) {
+        return  mRemoteComposeState.getInteger(id);
+    }
+
+    @Override
     public int getColor(int id) {
         return mRemoteComposeState.getColor(id);
     }
diff --git a/core/java/com/android/internal/widget/remotecompose/player/platform/FloatsToPath.java b/core/java/com/android/internal/widget/remotecompose/player/platform/FloatsToPath.java
index 2d766f8..7a85427 100644
--- a/core/java/com/android/internal/widget/remotecompose/player/platform/FloatsToPath.java
+++ b/core/java/com/android/internal/widget/remotecompose/player/platform/FloatsToPath.java
@@ -58,7 +58,7 @@
                 break;
                 case PathData.CONIC: {
                     i += 3;
-                    if (Build.VERSION.SDK_INT >= 34) {
+                    if (Build.VERSION.SDK_INT >= 34) { // REMOVE IN PLATFORM
                         path.conicTo(
                                 floatPath[i + 0], floatPath[i + 1],
                                 floatPath[i + 2], floatPath[i + 3],
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 9ce7658..0f53164 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -22,6 +22,7 @@
 #include <android/graphics/properties.h>
 #include <android/graphics/region.h>
 #include <android/gui/BnWindowInfosReportedListener.h>
+#include <android/gui/EdgeExtensionParameters.h>
 #include <android/gui/JankData.h>
 #include <android/hardware/display/IDeviceProductInfoConstants.h>
 #include <android/os/IInputConstants.h>
@@ -799,6 +800,20 @@
     transaction->setStretchEffect(ctrl, stretch);
 }
 
+static void nativeSetEdgeExtensionEffect(JNIEnv* env, jclass clazz, jlong transactionObj,
+                                         jlong nativeObj, jboolean leftEdge, jboolean rightEdge,
+                                         jboolean topEdge, jboolean bottomEdge) {
+    auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
+    auto* const ctrl = reinterpret_cast<SurfaceControl*>(nativeObj);
+
+    auto effect = gui::EdgeExtensionParameters();
+    effect.extendLeft = leftEdge;
+    effect.extendRight = rightEdge;
+    effect.extendTop = topEdge;
+    effect.extendBottom = bottomEdge;
+    transaction->setEdgeExtensionEffect(ctrl, effect);
+}
+
 static void nativeSetFlags(JNIEnv* env, jclass clazz, jlong transactionObj,
         jlong nativeObject, jint flags, jint mask) {
     auto transaction = reinterpret_cast<SurfaceComposerClient::Transaction*>(transactionObj);
@@ -2340,6 +2355,8 @@
             (void*)nativeSetBlurRegions },
     {"nativeSetStretchEffect", "(JJFFFFFFFFFF)V",
             (void*) nativeSetStretchEffect },
+    {"nativeSetEdgeExtensionEffect", "(JJZZZZ)V",
+            (void*) nativeSetEdgeExtensionEffect },
     {"nativeSetShadowRadius", "(JJF)V",
             (void*)nativeSetShadowRadius },
     {"nativeSetFrameRate", "(JJFII)V",
diff --git a/core/proto/android/app/appexitinfo.proto b/core/proto/android/app/appexitinfo.proto
index e560a94..8ee7962 100644
--- a/core/proto/android/app/appexitinfo.proto
+++ b/core/proto/android/app/appexitinfo.proto
@@ -20,7 +20,7 @@
 package android.app;
 
 import "frameworks/base/core/proto/android/privacy.proto";
-import "frameworks/proto_logging/stats/enums/app/app_enums.proto";
+import "frameworks/proto_logging/stats/enums/app_shared/app_enums.proto";
 
 /**
  * An android.app.ApplicationExitInfo object.
diff --git a/core/proto/android/app/appstartinfo.proto b/core/proto/android/app/appstartinfo.proto
index c137533..8de5458 100644
--- a/core/proto/android/app/appstartinfo.proto
+++ b/core/proto/android/app/appstartinfo.proto
@@ -20,7 +20,7 @@
 package android.app;
 
 import "frameworks/base/core/proto/android/privacy.proto";
-import "frameworks/proto_logging/stats/enums/app/app_enums.proto";
+import "frameworks/proto_logging/stats/enums/app_shared/app_enums.proto";
 
 /**
  * An android.app.ApplicationStartInfo object.
diff --git a/core/proto/android/server/activitymanagerservice.proto b/core/proto/android/server/activitymanagerservice.proto
index 90069f1..58f39a9 100644
--- a/core/proto/android/server/activitymanagerservice.proto
+++ b/core/proto/android/server/activitymanagerservice.proto
@@ -35,7 +35,7 @@
 import "frameworks/base/core/proto/android/server/windowmanagerservice.proto";
 import "frameworks/base/core/proto/android/util/common.proto";
 import "frameworks/base/core/proto/android/privacy.proto";
-import "frameworks/proto_logging/stats/enums/app/app_enums.proto";
+import "frameworks/proto_logging/stats/enums/app_shared/app_enums.proto";
 
 option java_multiple_files = true;
 
diff --git a/core/proto/android/server/powermanagerservice.proto b/core/proto/android/server/powermanagerservice.proto
index 593bbc6..8fd5d71 100644
--- a/core/proto/android/server/powermanagerservice.proto
+++ b/core/proto/android/server/powermanagerservice.proto
@@ -26,7 +26,7 @@
 import "frameworks/base/core/proto/android/providers/settings.proto";
 import "frameworks/base/core/proto/android/server/wirelesschargerdetector.proto";
 import "frameworks/base/core/proto/android/privacy.proto";
-import "frameworks/proto_logging/stats/enums/app/app_enums.proto";
+import "frameworks/proto_logging/stats/enums/app_shared/app_enums.proto";
 import "frameworks/proto_logging/stats/enums/os/enums.proto";
 import "frameworks/proto_logging/stats/enums/view/enums.proto";
 
diff --git a/core/tests/coretests/src/android/accessibilityservice/AccessibilityShortcutInfoTest.java b/core/tests/coretests/src/android/accessibilityservice/AccessibilityShortcutInfoTest.java
index eebc578..ef7df59 100644
--- a/core/tests/coretests/src/android/accessibilityservice/AccessibilityShortcutInfoTest.java
+++ b/core/tests/coretests/src/android/accessibilityservice/AccessibilityShortcutInfoTest.java
@@ -29,9 +29,9 @@
 import android.view.accessibility.AccessibilityManager;
 import android.view.accessibility.AccessibilityTestActivity;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 
diff --git a/core/tests/coretests/src/android/app/activity/ServiceTest.java b/core/tests/coretests/src/android/app/activity/ServiceTest.java
index 3f3d6a3..4e3b2af 100644
--- a/core/tests/coretests/src/android/app/activity/ServiceTest.java
+++ b/core/tests/coretests/src/android/app/activity/ServiceTest.java
@@ -157,6 +157,21 @@
         assertThat(mCurrentConnection.takePid(), is(NOT_STARTED));
     }
 
+    @Test
+    public void testRestart_stickyStartedService_unbindHappenedAfterRestart_restarted() {
+        final int servicePid = startService(Service.START_STICKY);
+        assertThat(servicePid, not(NOT_STARTED));
+        assertThat(bindService(0 /* flags */), is(servicePid));
+
+        final int restartedServicePid = waitForServiceStarted(
+                () -> {
+                    Process.killProcess(servicePid);
+                    mContext.unbindService(mCurrentConnection);
+                    mCurrentConnection = null;
+                });
+        assertThat(restartedServicePid, not(NOT_STARTED));
+    }
+
     /** @return The pid of the started service. */
     private int startService(int code) {
         return waitForServiceStarted(
diff --git a/core/tests/coretests/src/android/content/AbstractCrossUserContentResolverTest.java b/core/tests/coretests/src/android/content/AbstractCrossUserContentResolverTest.java
index 92b1c04..1d360cc 100644
--- a/core/tests/coretests/src/android/content/AbstractCrossUserContentResolverTest.java
+++ b/core/tests/coretests/src/android/content/AbstractCrossUserContentResolverTest.java
@@ -18,7 +18,6 @@
 
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
-import org.junit.Ignore;
 
 import android.app.ActivityManager;
 import android.app.activity.LocalProvider;
@@ -33,13 +32,14 @@
 import android.os.UserManager;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.compatibility.common.util.SystemUtil;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
diff --git a/core/tests/coretests/src/android/content/ApexEnvironmentTest.java b/core/tests/coretests/src/android/content/ApexEnvironmentTest.java
index 438c5ae..f3803d2 100644
--- a/core/tests/coretests/src/android/content/ApexEnvironmentTest.java
+++ b/core/tests/coretests/src/android/content/ApexEnvironmentTest.java
@@ -20,8 +20,8 @@
 
 import android.os.UserHandle;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/content/BroadcastReceiverTests.java b/core/tests/coretests/src/android/content/BroadcastReceiverTests.java
index 407c6c3..a9e781c 100644
--- a/core/tests/coretests/src/android/content/BroadcastReceiverTests.java
+++ b/core/tests/coretests/src/android/content/BroadcastReceiverTests.java
@@ -19,8 +19,8 @@
 import static org.junit.Assert.fail;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/content/ContentProviderTest.java b/core/tests/coretests/src/android/content/ContentProviderTest.java
index c9a6d22..4ecd2da 100644
--- a/core/tests/coretests/src/android/content/ContentProviderTest.java
+++ b/core/tests/coretests/src/android/content/ContentProviderTest.java
@@ -26,7 +26,7 @@
 import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/content/ContentResolverTest.java b/core/tests/coretests/src/android/content/ContentResolverTest.java
index c8015d4..dfde0bc 100644
--- a/core/tests/coretests/src/android/content/ContentResolverTest.java
+++ b/core/tests/coretests/src/android/content/ContentResolverTest.java
@@ -39,7 +39,7 @@
 import android.util.Size;
 
 import androidx.test.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/content/PermissionCheckerTest.java b/core/tests/coretests/src/android/content/PermissionCheckerTest.java
index cb04a74..65d8e2b 100644
--- a/core/tests/coretests/src/android/content/PermissionCheckerTest.java
+++ b/core/tests/coretests/src/android/content/PermissionCheckerTest.java
@@ -23,7 +23,7 @@
 import android.os.Binder;
 
 import androidx.test.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/content/pm/ModuleInfoTest.java b/core/tests/coretests/src/android/content/pm/ModuleInfoTest.java
index 4366e02c..7799185 100644
--- a/core/tests/coretests/src/android/content/pm/ModuleInfoTest.java
+++ b/core/tests/coretests/src/android/content/pm/ModuleInfoTest.java
@@ -22,7 +22,7 @@
 import android.platform.test.annotations.AppModeFull;
 import android.text.TextUtils;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/content/pm/PackageManagerPropertyTests.java b/core/tests/coretests/src/android/content/pm/PackageManagerPropertyTests.java
index 86e95832..ad3a16b 100644
--- a/core/tests/coretests/src/android/content/pm/PackageManagerPropertyTests.java
+++ b/core/tests/coretests/src/android/content/pm/PackageManagerPropertyTests.java
@@ -19,15 +19,14 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.assertTrue;
 
 import android.content.pm.PackageManager.Property;
 import android.os.Bundle;
 import android.platform.test.annotations.Presubmit;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/content/pm/PackageManagerTest.java b/core/tests/coretests/src/android/content/pm/PackageManagerTest.java
index 20421d1..b60d614 100644
--- a/core/tests/coretests/src/android/content/pm/PackageManagerTest.java
+++ b/core/tests/coretests/src/android/content/pm/PackageManagerTest.java
@@ -18,8 +18,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/content/pm/PackageParserCacheHelperTest.java b/core/tests/coretests/src/android/content/pm/PackageParserCacheHelperTest.java
index 61a3a11..dbd6c2b 100644
--- a/core/tests/coretests/src/android/content/pm/PackageParserCacheHelperTest.java
+++ b/core/tests/coretests/src/android/content/pm/PackageParserCacheHelperTest.java
@@ -24,8 +24,8 @@
 import android.os.Parcel;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/content/pm/PackagePartitionsTest.java b/core/tests/coretests/src/android/content/pm/PackagePartitionsTest.java
index 2986d61..6c23ea3 100644
--- a/core/tests/coretests/src/android/content/pm/PackagePartitionsTest.java
+++ b/core/tests/coretests/src/android/content/pm/PackagePartitionsTest.java
@@ -24,7 +24,7 @@
 import android.os.SystemProperties;
 import android.platform.test.annotations.Presubmit;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/content/pm/ShortcutQueryWrapperTest.java b/core/tests/coretests/src/android/content/pm/ShortcutQueryWrapperTest.java
index 8f8488f..ec5e205 100644
--- a/core/tests/coretests/src/android/content/pm/ShortcutQueryWrapperTest.java
+++ b/core/tests/coretests/src/android/content/pm/ShortcutQueryWrapperTest.java
@@ -23,7 +23,7 @@
 import android.os.Parcel;
 import android.platform.test.annotations.Presubmit;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.google.android.collect.Lists;
 
diff --git a/core/tests/coretests/src/android/content/pm/UserInfoTest.java b/core/tests/coretests/src/android/content/pm/UserInfoTest.java
index af36dbb..edeea6d 100644
--- a/core/tests/coretests/src/android/content/pm/UserInfoTest.java
+++ b/core/tests/coretests/src/android/content/pm/UserInfoTest.java
@@ -20,8 +20,8 @@
 
 import android.os.UserHandle;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/content/res/FontResourcesParserTest.java b/core/tests/coretests/src/android/content/res/FontResourcesParserTest.java
index 9aef2ca..85f5d69 100644
--- a/core/tests/coretests/src/android/content/res/FontResourcesParserTest.java
+++ b/core/tests/coretests/src/android/content/res/FontResourcesParserTest.java
@@ -30,8 +30,8 @@
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 
diff --git a/core/tests/coretests/src/android/content/res/ResourcesDrawableTest.java b/core/tests/coretests/src/android/content/res/ResourcesDrawableTest.java
index f7f9569..ac69a0f 100644
--- a/core/tests/coretests/src/android/content/res/ResourcesDrawableTest.java
+++ b/core/tests/coretests/src/android/content/res/ResourcesDrawableTest.java
@@ -27,8 +27,8 @@
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 
diff --git a/core/tests/coretests/src/android/database/CursorWindowTest.java b/core/tests/coretests/src/android/database/CursorWindowTest.java
index 255020a..64e4d3b 100644
--- a/core/tests/coretests/src/android/database/CursorWindowTest.java
+++ b/core/tests/coretests/src/android/database/CursorWindowTest.java
@@ -23,8 +23,8 @@
 import android.database.sqlite.SQLiteException;
 import android.platform.test.ravenwood.RavenwoodRule;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Rule;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/database/DatabaseUtilsTest.java b/core/tests/coretests/src/android/database/DatabaseUtilsTest.java
index e25fdf9..c00c171 100644
--- a/core/tests/coretests/src/android/database/DatabaseUtilsTest.java
+++ b/core/tests/coretests/src/android/database/DatabaseUtilsTest.java
@@ -16,20 +16,19 @@
 
 package android.database;
 
-import static android.database.DatabaseUtils.bindSelection;
-import static android.database.DatabaseUtils.getSqlStatementType;
-import static android.database.DatabaseUtils.getSqlStatementTypeExtended;
-import static android.database.DatabaseUtils.STATEMENT_COMMENT;
 import static android.database.DatabaseUtils.STATEMENT_CREATE;
 import static android.database.DatabaseUtils.STATEMENT_DDL;
 import static android.database.DatabaseUtils.STATEMENT_OTHER;
 import static android.database.DatabaseUtils.STATEMENT_SELECT;
 import static android.database.DatabaseUtils.STATEMENT_UPDATE;
 import static android.database.DatabaseUtils.STATEMENT_WITH;
+import static android.database.DatabaseUtils.bindSelection;
+import static android.database.DatabaseUtils.getSqlStatementType;
+import static android.database.DatabaseUtils.getSqlStatementTypeExtended;
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/database/RedactingCursorTest.java b/core/tests/coretests/src/android/database/RedactingCursorTest.java
index e2d2bae..470d4a9 100644
--- a/core/tests/coretests/src/android/database/RedactingCursorTest.java
+++ b/core/tests/coretests/src/android/database/RedactingCursorTest.java
@@ -24,7 +24,7 @@
 import android.net.Uri;
 
 import androidx.test.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/database/SQLiteOpenHelperTest.java b/core/tests/coretests/src/android/database/SQLiteOpenHelperTest.java
index 9bad6d2..efa9b7a 100644
--- a/core/tests/coretests/src/android/database/SQLiteOpenHelperTest.java
+++ b/core/tests/coretests/src/android/database/SQLiteOpenHelperTest.java
@@ -29,8 +29,8 @@
 import android.util.Log;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/database/sqlite/SQLiteCantOpenDatabaseExceptionTest.java b/core/tests/coretests/src/android/database/sqlite/SQLiteCantOpenDatabaseExceptionTest.java
index 09c380a..e20a806 100644
--- a/core/tests/coretests/src/android/database/sqlite/SQLiteCantOpenDatabaseExceptionTest.java
+++ b/core/tests/coretests/src/android/database/sqlite/SQLiteCantOpenDatabaseExceptionTest.java
@@ -20,14 +20,12 @@
 import static org.junit.Assert.fail;
 
 import android.content.Context;
-import android.database.sqlite.SQLiteCantOpenDatabaseException;
-import android.database.sqlite.SQLiteDatabase;
 import android.database.sqlite.SQLiteDatabase.OpenParams;
 import android.util.Log;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/database/sqlite/SQLiteCompatibilityWalFlagsTest.java b/core/tests/coretests/src/android/database/sqlite/SQLiteCompatibilityWalFlagsTest.java
index 82bd588..a4dedc5 100644
--- a/core/tests/coretests/src/android/database/sqlite/SQLiteCompatibilityWalFlagsTest.java
+++ b/core/tests/coretests/src/android/database/sqlite/SQLiteCompatibilityWalFlagsTest.java
@@ -25,8 +25,8 @@
 import android.database.DatabaseUtils;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/database/sqlite/SQLiteConnectionPoolTest.java b/core/tests/coretests/src/android/database/sqlite/SQLiteConnectionPoolTest.java
index 2b663bd..dbe7a9a 100644
--- a/core/tests/coretests/src/android/database/sqlite/SQLiteConnectionPoolTest.java
+++ b/core/tests/coretests/src/android/database/sqlite/SQLiteConnectionPoolTest.java
@@ -24,8 +24,8 @@
 import android.util.Log;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/database/sqlite/SQLiteDatabaseTest.java b/core/tests/coretests/src/android/database/sqlite/SQLiteDatabaseTest.java
index c4695d9..bd9c4b8 100644
--- a/core/tests/coretests/src/android/database/sqlite/SQLiteDatabaseTest.java
+++ b/core/tests/coretests/src/android/database/sqlite/SQLiteDatabaseTest.java
@@ -25,16 +25,13 @@
 import android.content.Context;
 import android.database.Cursor;
 import android.database.DatabaseUtils;
-import android.os.SystemClock;
 import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
-import android.test.AndroidTestCase;
-import android.util.Log;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
@@ -46,11 +43,9 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Phaser;
-import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 
 @RunWith(AndroidJUnit4.class)
diff --git a/core/tests/coretests/src/android/database/sqlite/SQLiteRawStatementTest.java b/core/tests/coretests/src/android/database/sqlite/SQLiteRawStatementTest.java
index 24d27c4..832ebe5 100644
--- a/core/tests/coretests/src/android/database/sqlite/SQLiteRawStatementTest.java
+++ b/core/tests/coretests/src/android/database/sqlite/SQLiteRawStatementTest.java
@@ -29,9 +29,9 @@
 import android.os.SystemClock;
 import android.util.Log;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/database/sqlite/SQLiteTokenizerTest.java b/core/tests/coretests/src/android/database/sqlite/SQLiteTokenizerTest.java
index a9d1482..ced1846 100644
--- a/core/tests/coretests/src/android/database/sqlite/SQLiteTokenizerTest.java
+++ b/core/tests/coretests/src/android/database/sqlite/SQLiteTokenizerTest.java
@@ -20,7 +20,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/security/CredentialManagementAppTest.java b/core/tests/coretests/src/android/security/CredentialManagementAppTest.java
index fa824b1..1690741 100644
--- a/core/tests/coretests/src/android/security/CredentialManagementAppTest.java
+++ b/core/tests/coretests/src/android/security/CredentialManagementAppTest.java
@@ -23,8 +23,8 @@
 import android.net.Uri;
 import android.util.Xml;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/security/keystore/recovery/KeyChainProtectionParamsTest.java b/core/tests/coretests/src/android/security/keystore/recovery/KeyChainProtectionParamsTest.java
index ce0bf30..938147c 100644
--- a/core/tests/coretests/src/android/security/keystore/recovery/KeyChainProtectionParamsTest.java
+++ b/core/tests/coretests/src/android/security/keystore/recovery/KeyChainProtectionParamsTest.java
@@ -21,8 +21,8 @@
 
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/security/keystore/recovery/KeyChainSnapshotTest.java b/core/tests/coretests/src/android/security/keystore/recovery/KeyChainSnapshotTest.java
index 37fe22f..242a273 100644
--- a/core/tests/coretests/src/android/security/keystore/recovery/KeyChainSnapshotTest.java
+++ b/core/tests/coretests/src/android/security/keystore/recovery/KeyChainSnapshotTest.java
@@ -21,8 +21,8 @@
 
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.google.common.collect.Lists;
 
diff --git a/core/tests/coretests/src/android/security/keystore/recovery/KeyDerivationParamsTest.java b/core/tests/coretests/src/android/security/keystore/recovery/KeyDerivationParamsTest.java
index 2b37b52..d310b76 100644
--- a/core/tests/coretests/src/android/security/keystore/recovery/KeyDerivationParamsTest.java
+++ b/core/tests/coretests/src/android/security/keystore/recovery/KeyDerivationParamsTest.java
@@ -21,8 +21,8 @@
 
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/security/keystore/recovery/RecoveryCertPathTest.java b/core/tests/coretests/src/android/security/keystore/recovery/RecoveryCertPathTest.java
index 3b8f715..29001ae 100644
--- a/core/tests/coretests/src/android/security/keystore/recovery/RecoveryCertPathTest.java
+++ b/core/tests/coretests/src/android/security/keystore/recovery/RecoveryCertPathTest.java
@@ -21,8 +21,8 @@
 
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/security/keystore/recovery/TrustedRootCertificatesTest.java b/core/tests/coretests/src/android/security/keystore/recovery/TrustedRootCertificatesTest.java
index 2b15d73..7677c3c 100644
--- a/core/tests/coretests/src/android/security/keystore/recovery/TrustedRootCertificatesTest.java
+++ b/core/tests/coretests/src/android/security/keystore/recovery/TrustedRootCertificatesTest.java
@@ -20,8 +20,8 @@
 
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/security/keystore/recovery/WrappedApplicationKeyTest.java b/core/tests/coretests/src/android/security/keystore/recovery/WrappedApplicationKeyTest.java
index aec54e1..07353f8 100644
--- a/core/tests/coretests/src/android/security/keystore/recovery/WrappedApplicationKeyTest.java
+++ b/core/tests/coretests/src/android/security/keystore/recovery/WrappedApplicationKeyTest.java
@@ -21,8 +21,8 @@
 
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/security/keystore/recovery/X509CertificateParsingUtilsTest.java b/core/tests/coretests/src/android/security/keystore/recovery/X509CertificateParsingUtilsTest.java
index ba5e74a..7f91d03 100644
--- a/core/tests/coretests/src/android/security/keystore/recovery/X509CertificateParsingUtilsTest.java
+++ b/core/tests/coretests/src/android/security/keystore/recovery/X509CertificateParsingUtilsTest.java
@@ -21,8 +21,8 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/BidiFormatterTest.java b/core/tests/coretests/src/android/text/BidiFormatterTest.java
index 312fb68..307c95b 100644
--- a/core/tests/coretests/src/android/text/BidiFormatterTest.java
+++ b/core/tests/coretests/src/android/text/BidiFormatterTest.java
@@ -20,8 +20,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/DynamicLayoutBlocksTest.java b/core/tests/coretests/src/android/text/DynamicLayoutBlocksTest.java
index cca1ad3..81c4982 100644
--- a/core/tests/coretests/src/android/text/DynamicLayoutBlocksTest.java
+++ b/core/tests/coretests/src/android/text/DynamicLayoutBlocksTest.java
@@ -23,8 +23,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/DynamicLayoutOffsetMappingTest.java b/core/tests/coretests/src/android/text/DynamicLayoutOffsetMappingTest.java
index 5939c06..5ff659b 100644
--- a/core/tests/coretests/src/android/text/DynamicLayoutOffsetMappingTest.java
+++ b/core/tests/coretests/src/android/text/DynamicLayoutOffsetMappingTest.java
@@ -23,8 +23,8 @@
 import android.platform.test.annotations.Presubmit;
 import android.text.method.OffsetMapping;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/DynamicLayoutTest.java b/core/tests/coretests/src/android/text/DynamicLayoutTest.java
index 699243b..1036928 100644
--- a/core/tests/coretests/src/android/text/DynamicLayoutTest.java
+++ b/core/tests/coretests/src/android/text/DynamicLayoutTest.java
@@ -30,8 +30,8 @@
 import android.text.style.ReplacementSpan;
 import android.util.ArraySet;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/EmojiConsistencyTest.java b/core/tests/coretests/src/android/text/EmojiConsistencyTest.java
index c6e9e9c..72d09c8 100644
--- a/core/tests/coretests/src/android/text/EmojiConsistencyTest.java
+++ b/core/tests/coretests/src/android/text/EmojiConsistencyTest.java
@@ -18,8 +18,8 @@
 
 import static junit.framework.Assert.assertEquals;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/EmojiTest.java b/core/tests/coretests/src/android/text/EmojiTest.java
index 0aeeb74..21f346e 100644
--- a/core/tests/coretests/src/android/text/EmojiTest.java
+++ b/core/tests/coretests/src/android/text/EmojiTest.java
@@ -22,8 +22,8 @@
 import android.icu.lang.UCharacterDirection;
 import android.icu.text.Bidi;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/LayoutBidiCursorPathTest.java b/core/tests/coretests/src/android/text/LayoutBidiCursorPathTest.java
index 96e7fb9..7728866 100644
--- a/core/tests/coretests/src/android/text/LayoutBidiCursorPathTest.java
+++ b/core/tests/coretests/src/android/text/LayoutBidiCursorPathTest.java
@@ -26,8 +26,8 @@
 import android.view.KeyEvent;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/LayoutTest.java b/core/tests/coretests/src/android/text/LayoutTest.java
index 98f8b7f..25f9cb7 100644
--- a/core/tests/coretests/src/android/text/LayoutTest.java
+++ b/core/tests/coretests/src/android/text/LayoutTest.java
@@ -42,8 +42,8 @@
 import android.text.style.ForegroundColorSpan;
 import android.text.style.StrikethroughSpan;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.google.common.truth.Expect;
 
diff --git a/core/tests/coretests/src/android/text/MeasuredParagraphTest.java b/core/tests/coretests/src/android/text/MeasuredParagraphTest.java
index 02b67e2..921a6bd 100644
--- a/core/tests/coretests/src/android/text/MeasuredParagraphTest.java
+++ b/core/tests/coretests/src/android/text/MeasuredParagraphTest.java
@@ -26,8 +26,8 @@
 import android.graphics.text.MeasuredText;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/PackedIntVectorTest.java b/core/tests/coretests/src/android/text/PackedIntVectorTest.java
index ba15b92..e8d706d 100644
--- a/core/tests/coretests/src/android/text/PackedIntVectorTest.java
+++ b/core/tests/coretests/src/android/text/PackedIntVectorTest.java
@@ -20,8 +20,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/SpanColorsTest.java b/core/tests/coretests/src/android/text/SpanColorsTest.java
index 3d8d8f9..d2cb8c1 100644
--- a/core/tests/coretests/src/android/text/SpanColorsTest.java
+++ b/core/tests/coretests/src/android/text/SpanColorsTest.java
@@ -25,8 +25,8 @@
 import android.text.style.ImageSpan;
 import android.text.style.UnderlineSpan;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/SpannableStringBuilderTest.java b/core/tests/coretests/src/android/text/SpannableStringBuilderTest.java
index 91b8c6a..b725133 100644
--- a/core/tests/coretests/src/android/text/SpannableStringBuilderTest.java
+++ b/core/tests/coretests/src/android/text/SpannableStringBuilderTest.java
@@ -25,8 +25,8 @@
 import android.text.style.SubscriptSpan;
 import android.text.style.UnderlineSpan;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/SpannableStringNoCopyTest.java b/core/tests/coretests/src/android/text/SpannableStringNoCopyTest.java
index 9149f7b..a2952f6 100644
--- a/core/tests/coretests/src/android/text/SpannableStringNoCopyTest.java
+++ b/core/tests/coretests/src/android/text/SpannableStringNoCopyTest.java
@@ -24,8 +24,8 @@
 import android.text.style.QuoteSpan;
 import android.text.style.UnderlineSpan;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/SpannableTest.java b/core/tests/coretests/src/android/text/SpannableTest.java
index d248a1f..a3e6a78 100644
--- a/core/tests/coretests/src/android/text/SpannableTest.java
+++ b/core/tests/coretests/src/android/text/SpannableTest.java
@@ -21,8 +21,8 @@
 import android.platform.test.annotations.Presubmit;
 import android.test.MoreAsserts;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/SpannedStringNoCopyTest.java b/core/tests/coretests/src/android/text/SpannedStringNoCopyTest.java
index ca43733..3e2516f 100644
--- a/core/tests/coretests/src/android/text/SpannedStringNoCopyTest.java
+++ b/core/tests/coretests/src/android/text/SpannedStringNoCopyTest.java
@@ -24,8 +24,8 @@
 import android.text.style.QuoteSpan;
 import android.text.style.UnderlineSpan;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/SpannedTest.java b/core/tests/coretests/src/android/text/SpannedTest.java
index 3ab0755..e9a357c 100644
--- a/core/tests/coretests/src/android/text/SpannedTest.java
+++ b/core/tests/coretests/src/android/text/SpannedTest.java
@@ -26,8 +26,8 @@
 import android.text.style.TextAppearanceSpan;
 import android.text.style.TypefaceSpan;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/StaticLayoutBidiTest.java b/core/tests/coretests/src/android/text/StaticLayoutBidiTest.java
index 32370b3e..3deda8c 100644
--- a/core/tests/coretests/src/android/text/StaticLayoutBidiTest.java
+++ b/core/tests/coretests/src/android/text/StaticLayoutBidiTest.java
@@ -21,8 +21,8 @@
 import android.platform.test.annotations.Presubmit;
 import android.util.Log;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/StaticLayoutDirectionsTest.java b/core/tests/coretests/src/android/text/StaticLayoutDirectionsTest.java
index 4221ac2..bc7efe4 100644
--- a/core/tests/coretests/src/android/text/StaticLayoutDirectionsTest.java
+++ b/core/tests/coretests/src/android/text/StaticLayoutDirectionsTest.java
@@ -22,8 +22,8 @@
 import android.text.Layout.Directions;
 import android.text.StaticLayoutTest.LayoutBuilder;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/StaticLayoutTest.java b/core/tests/coretests/src/android/text/StaticLayoutTest.java
index 0ebf03f..3541900 100644
--- a/core/tests/coretests/src/android/text/StaticLayoutTest.java
+++ b/core/tests/coretests/src/android/text/StaticLayoutTest.java
@@ -31,8 +31,8 @@
 import android.text.style.LocaleSpan;
 import android.util.Log;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/StaticLayoutTextMeasuringTest.java b/core/tests/coretests/src/android/text/StaticLayoutTextMeasuringTest.java
index 0d42326..b32e94a 100644
--- a/core/tests/coretests/src/android/text/StaticLayoutTextMeasuringTest.java
+++ b/core/tests/coretests/src/android/text/StaticLayoutTextMeasuringTest.java
@@ -21,8 +21,8 @@
 
 import android.text.Layout.Alignment;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/TextLayoutTest.java b/core/tests/coretests/src/android/text/TextLayoutTest.java
index 15fbc9e..1584bc3 100644
--- a/core/tests/coretests/src/android/text/TextLayoutTest.java
+++ b/core/tests/coretests/src/android/text/TextLayoutTest.java
@@ -18,8 +18,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/TextLineTest.java b/core/tests/coretests/src/android/text/TextLineTest.java
index 8ae5669..2997853 100644
--- a/core/tests/coretests/src/android/text/TextLineTest.java
+++ b/core/tests/coretests/src/android/text/TextLineTest.java
@@ -30,10 +30,10 @@
 import android.text.style.ReplacementSpan;
 import android.text.style.TabStopSpan;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.filters.Suppress;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/TextShaperTest.java b/core/tests/coretests/src/android/text/TextShaperTest.java
index 77b14e6..84112ae 100644
--- a/core/tests/coretests/src/android/text/TextShaperTest.java
+++ b/core/tests/coretests/src/android/text/TextShaperTest.java
@@ -21,8 +21,8 @@
 import android.graphics.fonts.Font;
 import android.graphics.fonts.FontFileUtil;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/TextUtilsTest.java b/core/tests/coretests/src/android/text/TextUtilsTest.java
index c4bcfd4..f552265 100644
--- a/core/tests/coretests/src/android/text/TextUtilsTest.java
+++ b/core/tests/coretests/src/android/text/TextUtilsTest.java
@@ -34,9 +34,9 @@
 import android.text.util.Rfc822Tokenizer;
 import android.view.View;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.google.android.collect.Lists;
 
diff --git a/core/tests/coretests/src/android/text/VariationParserTest.java b/core/tests/coretests/src/android/text/VariationParserTest.java
index 0afe811..8e93dd4 100644
--- a/core/tests/coretests/src/android/text/VariationParserTest.java
+++ b/core/tests/coretests/src/android/text/VariationParserTest.java
@@ -22,8 +22,8 @@
 import android.graphics.fonts.FontVariationAxis;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/format/DateFormatTest.java b/core/tests/coretests/src/android/text/format/DateFormatTest.java
index 212cc44..59af6dd 100644
--- a/core/tests/coretests/src/android/text/format/DateFormatTest.java
+++ b/core/tests/coretests/src/android/text/format/DateFormatTest.java
@@ -25,8 +25,8 @@
 import android.icu.text.DateFormatSymbols;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
 import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
diff --git a/core/tests/coretests/src/android/text/format/DateIntervalFormatTest.java b/core/tests/coretests/src/android/text/format/DateIntervalFormatTest.java
index 9750de3..a07d399 100644
--- a/core/tests/coretests/src/android/text/format/DateIntervalFormatTest.java
+++ b/core/tests/coretests/src/android/text/format/DateIntervalFormatTest.java
@@ -42,8 +42,8 @@
 import android.icu.util.ULocale;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/format/DateUtilsTest.java b/core/tests/coretests/src/android/text/format/DateUtilsTest.java
index 381c051..47be893 100644
--- a/core/tests/coretests/src/android/text/format/DateUtilsTest.java
+++ b/core/tests/coretests/src/android/text/format/DateUtilsTest.java
@@ -23,8 +23,8 @@
 import android.os.LocaleList;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/text/format/FormatterTest.java b/core/tests/coretests/src/android/text/format/FormatterTest.java
index 986cee5..555292e 100644
--- a/core/tests/coretests/src/android/text/format/FormatterTest.java
+++ b/core/tests/coretests/src/android/text/format/FormatterTest.java
@@ -28,8 +28,8 @@
 import android.text.format.Formatter.BytesResult;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/text/format/RelativeDateTimeFormatterTest.java b/core/tests/coretests/src/android/text/format/RelativeDateTimeFormatterTest.java
index 2337802..cd31950 100644
--- a/core/tests/coretests/src/android/text/format/RelativeDateTimeFormatterTest.java
+++ b/core/tests/coretests/src/android/text/format/RelativeDateTimeFormatterTest.java
@@ -36,8 +36,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/format/TimeMigrationUtilsTest.java b/core/tests/coretests/src/android/text/format/TimeMigrationUtilsTest.java
index b605520..c8cb5f3 100644
--- a/core/tests/coretests/src/android/text/format/TimeMigrationUtilsTest.java
+++ b/core/tests/coretests/src/android/text/format/TimeMigrationUtilsTest.java
@@ -20,8 +20,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/text/format/TimeTest.java b/core/tests/coretests/src/android/text/format/TimeTest.java
index ac00411..6138ea1 100644
--- a/core/tests/coretests/src/android/text/format/TimeTest.java
+++ b/core/tests/coretests/src/android/text/format/TimeTest.java
@@ -24,9 +24,9 @@
 import android.util.Log;
 import android.util.TimeFormatException;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.filters.Suppress;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/method/BackspaceTest.java b/core/tests/coretests/src/android/text/method/BackspaceTest.java
index 19c2c61..a7ff244 100644
--- a/core/tests/coretests/src/android/text/method/BackspaceTest.java
+++ b/core/tests/coretests/src/android/text/method/BackspaceTest.java
@@ -24,8 +24,8 @@
 import android.widget.TextView.BufferType;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/method/ForwardDeleteTest.java b/core/tests/coretests/src/android/text/method/ForwardDeleteTest.java
index 652622d..1e4024d 100644
--- a/core/tests/coretests/src/android/text/method/ForwardDeleteTest.java
+++ b/core/tests/coretests/src/android/text/method/ForwardDeleteTest.java
@@ -24,8 +24,8 @@
 import android.widget.TextView.BufferType;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/method/InsertModeTransformationMethodTest.java b/core/tests/coretests/src/android/text/method/InsertModeTransformationMethodTest.java
index 9ef137b..2f336ab 100644
--- a/core/tests/coretests/src/android/text/method/InsertModeTransformationMethodTest.java
+++ b/core/tests/coretests/src/android/text/method/InsertModeTransformationMethodTest.java
@@ -27,9 +27,8 @@
 import android.view.View;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
-
 
 import org.junit.BeforeClass;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/text/method/WordIteratorTest.java b/core/tests/coretests/src/android/text/method/WordIteratorTest.java
index cc345f5..046496a 100644
--- a/core/tests/coretests/src/android/text/method/WordIteratorTest.java
+++ b/core/tests/coretests/src/android/text/method/WordIteratorTest.java
@@ -23,8 +23,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/style/UnderlineSpanTest.java b/core/tests/coretests/src/android/text/style/UnderlineSpanTest.java
index a0d2f85..043960d0 100644
--- a/core/tests/coretests/src/android/text/style/UnderlineSpanTest.java
+++ b/core/tests/coretests/src/android/text/style/UnderlineSpanTest.java
@@ -24,8 +24,8 @@
 import android.text.StaticLayout;
 import android.text.TextPaint;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/text/util/LinkifyTest.java b/core/tests/coretests/src/android/text/util/LinkifyTest.java
index 107ecd7..52f3b2e 100644
--- a/core/tests/coretests/src/android/text/util/LinkifyTest.java
+++ b/core/tests/coretests/src/android/text/util/LinkifyTest.java
@@ -31,8 +31,8 @@
 import android.widget.TextView;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/util/ArrayMapTest.java b/core/tests/coretests/src/android/util/ArrayMapTest.java
index 1e444ad..711ff94 100644
--- a/core/tests/coretests/src/android/util/ArrayMapTest.java
+++ b/core/tests/coretests/src/android/util/ArrayMapTest.java
@@ -19,8 +19,8 @@
 import android.platform.test.annotations.IgnoreUnderRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Rule;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/ArraySetTest.java b/core/tests/coretests/src/android/util/ArraySetTest.java
index 51de634..8888991 100644
--- a/core/tests/coretests/src/android/util/ArraySetTest.java
+++ b/core/tests/coretests/src/android/util/ArraySetTest.java
@@ -18,8 +18,8 @@
 
 import static org.junit.Assert.fail;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/Base64Test.java b/core/tests/coretests/src/android/util/Base64Test.java
index b648266..3b322c2 100644
--- a/core/tests/coretests/src/android/util/Base64Test.java
+++ b/core/tests/coretests/src/android/util/Base64Test.java
@@ -20,8 +20,8 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Ignore;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/BinaryXmlTest.java b/core/tests/coretests/src/android/util/BinaryXmlTest.java
index da29828..96c79013 100644
--- a/core/tests/coretests/src/android/util/BinaryXmlTest.java
+++ b/core/tests/coretests/src/android/util/BinaryXmlTest.java
@@ -30,7 +30,7 @@
 
 import android.os.PersistableBundle;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
diff --git a/core/tests/coretests/src/android/util/CharsetUtilsTest.java b/core/tests/coretests/src/android/util/CharsetUtilsTest.java
index fbbe311..33936e9 100644
--- a/core/tests/coretests/src/android/util/CharsetUtilsTest.java
+++ b/core/tests/coretests/src/android/util/CharsetUtilsTest.java
@@ -21,7 +21,7 @@
 import android.platform.test.annotations.IgnoreUnderRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.util.HexDump;
 
diff --git a/core/tests/coretests/src/android/util/DayOfMonthCursorTest.java b/core/tests/coretests/src/android/util/DayOfMonthCursorTest.java
index 72bd578..4587f5b 100644
--- a/core/tests/coretests/src/android/util/DayOfMonthCursorTest.java
+++ b/core/tests/coretests/src/android/util/DayOfMonthCursorTest.java
@@ -21,8 +21,8 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/FloatMathTest.java b/core/tests/coretests/src/android/util/FloatMathTest.java
index f748acd..a52c9ac 100644
--- a/core/tests/coretests/src/android/util/FloatMathTest.java
+++ b/core/tests/coretests/src/android/util/FloatMathTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/KeyValueListParserTest.java b/core/tests/coretests/src/android/util/KeyValueListParserTest.java
index f65c4c7..eaf8d19 100644
--- a/core/tests/coretests/src/android/util/KeyValueListParserTest.java
+++ b/core/tests/coretests/src/android/util/KeyValueListParserTest.java
@@ -20,8 +20,8 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/LogNullabilityTest.java b/core/tests/coretests/src/android/util/LogNullabilityTest.java
index 475e347..5aa2626 100644
--- a/core/tests/coretests/src/android/util/LogNullabilityTest.java
+++ b/core/tests/coretests/src/android/util/LogNullabilityTest.java
@@ -21,8 +21,8 @@
 
 import android.platform.test.ravenwood.RavenwoodRule;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/LogTest.java b/core/tests/coretests/src/android/util/LogTest.java
index 15caac9..0e44e68 100644
--- a/core/tests/coretests/src/android/util/LogTest.java
+++ b/core/tests/coretests/src/android/util/LogTest.java
@@ -19,8 +19,8 @@
 import android.os.SystemProperties;
 import android.test.PerformanceTestCase;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.Suppress;
-import androidx.test.runner.AndroidJUnit4;
 
 import junit.framework.TestCase;
 
diff --git a/core/tests/coretests/src/android/util/LogWriterTest.java b/core/tests/coretests/src/android/util/LogWriterTest.java
index 890a401..739fbcd 100644
--- a/core/tests/coretests/src/android/util/LogWriterTest.java
+++ b/core/tests/coretests/src/android/util/LogWriterTest.java
@@ -16,7 +16,7 @@
 
 package android.util;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/LongArrayQueueTest.java b/core/tests/coretests/src/android/util/LongArrayQueueTest.java
index 77e8d60..f696f70 100644
--- a/core/tests/coretests/src/android/util/LongArrayQueueTest.java
+++ b/core/tests/coretests/src/android/util/LongArrayQueueTest.java
@@ -20,8 +20,8 @@
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.fail;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/LongSparseLongArrayTest.java b/core/tests/coretests/src/android/util/LongSparseLongArrayTest.java
index 9dbaae0d..fb0b2e3 100644
--- a/core/tests/coretests/src/android/util/LongSparseLongArrayTest.java
+++ b/core/tests/coretests/src/android/util/LongSparseLongArrayTest.java
@@ -18,8 +18,8 @@
 
 import static org.junit.Assert.assertEquals;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/LruCacheTest.java b/core/tests/coretests/src/android/util/LruCacheTest.java
index 10e8308..1c6dcdf 100644
--- a/core/tests/coretests/src/android/util/LruCacheTest.java
+++ b/core/tests/coretests/src/android/util/LruCacheTest.java
@@ -20,7 +20,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.fail;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/MonthDisplayHelperTest.java b/core/tests/coretests/src/android/util/MonthDisplayHelperTest.java
index 06f970f..cb34c98 100644
--- a/core/tests/coretests/src/android/util/MonthDisplayHelperTest.java
+++ b/core/tests/coretests/src/android/util/MonthDisplayHelperTest.java
@@ -20,8 +20,8 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/MutableTest.java b/core/tests/coretests/src/android/util/MutableTest.java
index dfdff4d..1f73c16 100644
--- a/core/tests/coretests/src/android/util/MutableTest.java
+++ b/core/tests/coretests/src/android/util/MutableTest.java
@@ -20,7 +20,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/PatternsTest.java b/core/tests/coretests/src/android/util/PatternsTest.java
index a180ec3..8d0785f 100644
--- a/core/tests/coretests/src/android/util/PatternsTest.java
+++ b/core/tests/coretests/src/android/util/PatternsTest.java
@@ -20,8 +20,8 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/PoolsTest.java b/core/tests/coretests/src/android/util/PoolsTest.java
index bdbc9b1..e31ab78 100644
--- a/core/tests/coretests/src/android/util/PoolsTest.java
+++ b/core/tests/coretests/src/android/util/PoolsTest.java
@@ -19,7 +19,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/PrefixPrinterTest.java b/core/tests/coretests/src/android/util/PrefixPrinterTest.java
index a8d48ee..8199155 100644
--- a/core/tests/coretests/src/android/util/PrefixPrinterTest.java
+++ b/core/tests/coretests/src/android/util/PrefixPrinterTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/RecurrenceRuleTest.java b/core/tests/coretests/src/android/util/RecurrenceRuleTest.java
index 32548b4..8b2068c 100644
--- a/core/tests/coretests/src/android/util/RecurrenceRuleTest.java
+++ b/core/tests/coretests/src/android/util/RecurrenceRuleTest.java
@@ -20,8 +20,8 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/android/util/SequenceUtilsTest.java b/core/tests/coretests/src/android/util/SequenceUtilsTest.java
index 6ca1751..e5ee04c 100644
--- a/core/tests/coretests/src/android/util/SequenceUtilsTest.java
+++ b/core/tests/coretests/src/android/util/SequenceUtilsTest.java
@@ -29,8 +29,8 @@
 import android.platform.test.annotations.Presubmit;
 import android.platform.test.ravenwood.RavenwoodRule;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Rule;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/SingletonTest.java b/core/tests/coretests/src/android/util/SingletonTest.java
index 8c5a963..31ae650 100644
--- a/core/tests/coretests/src/android/util/SingletonTest.java
+++ b/core/tests/coretests/src/android/util/SingletonTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertTrue;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/SparseDoubleArrayTest.java b/core/tests/coretests/src/android/util/SparseDoubleArrayTest.java
index ba9c8d9..cc5ffbe 100644
--- a/core/tests/coretests/src/android/util/SparseDoubleArrayTest.java
+++ b/core/tests/coretests/src/android/util/SparseDoubleArrayTest.java
@@ -18,8 +18,8 @@
 
 import static org.junit.Assert.assertEquals;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/SparseLongArrayTest.java b/core/tests/coretests/src/android/util/SparseLongArrayTest.java
index b29b6f1..4038d88 100644
--- a/core/tests/coretests/src/android/util/SparseLongArrayTest.java
+++ b/core/tests/coretests/src/android/util/SparseLongArrayTest.java
@@ -21,8 +21,8 @@
 
 import android.annotation.NonNull;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/SparseSetArrayTest.java b/core/tests/coretests/src/android/util/SparseSetArrayTest.java
index a8dce70..48ea3a1 100644
--- a/core/tests/coretests/src/android/util/SparseSetArrayTest.java
+++ b/core/tests/coretests/src/android/util/SparseSetArrayTest.java
@@ -19,8 +19,8 @@
 
 import android.platform.test.ravenwood.RavenwoodRule;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Rule;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/StateSetTest.java b/core/tests/coretests/src/android/util/StateSetTest.java
index dfd1523..14e4e20 100644
--- a/core/tests/coretests/src/android/util/StateSetTest.java
+++ b/core/tests/coretests/src/android/util/StateSetTest.java
@@ -22,8 +22,8 @@
 import android.platform.test.annotations.IgnoreUnderRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Rule;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/util/TeeWriterTest.java b/core/tests/coretests/src/android/util/TeeWriterTest.java
index c78376a..cc1c091 100644
--- a/core/tests/coretests/src/android/util/TeeWriterTest.java
+++ b/core/tests/coretests/src/android/util/TeeWriterTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/util/XmlTest.java b/core/tests/coretests/src/android/util/XmlTest.java
index 91ebc2a..540f118 100644
--- a/core/tests/coretests/src/android/util/XmlTest.java
+++ b/core/tests/coretests/src/android/util/XmlTest.java
@@ -26,7 +26,7 @@
 
 import android.os.PersistableBundle;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.TypedXmlPullParser;
diff --git a/core/tests/coretests/src/android/view/CompositionSamplingListenerTest.java b/core/tests/coretests/src/android/view/CompositionSamplingListenerTest.java
index 729a555..e74027e 100644
--- a/core/tests/coretests/src/android/view/CompositionSamplingListenerTest.java
+++ b/core/tests/coretests/src/android/view/CompositionSamplingListenerTest.java
@@ -21,8 +21,8 @@
 import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/CutoutSpecificationTest.java b/core/tests/coretests/src/android/view/CutoutSpecificationTest.java
index 7872810..0fdc239 100644
--- a/core/tests/coretests/src/android/view/CutoutSpecificationTest.java
+++ b/core/tests/coretests/src/android/view/CutoutSpecificationTest.java
@@ -23,8 +23,8 @@
 import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/view/DisplayCutoutTest.java b/core/tests/coretests/src/android/view/DisplayCutoutTest.java
index 0d1dde3..2c66330 100644
--- a/core/tests/coretests/src/android/view/DisplayCutoutTest.java
+++ b/core/tests/coretests/src/android/view/DisplayCutoutTest.java
@@ -44,8 +44,8 @@
 import android.platform.test.annotations.Presubmit;
 import android.view.DisplayCutout.ParcelableWrapper;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/DisplayShapeTest.java b/core/tests/coretests/src/android/view/DisplayShapeTest.java
index 77dd8bd..7778ba1 100644
--- a/core/tests/coretests/src/android/view/DisplayShapeTest.java
+++ b/core/tests/coretests/src/android/view/DisplayShapeTest.java
@@ -27,8 +27,8 @@
 import android.graphics.RectF;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
index 1682135..668487d 100644
--- a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
+++ b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
@@ -42,7 +42,7 @@
 import android.view.SyncRtSurfaceTransactionApplier.SurfaceParams;
 import android.view.animation.LinearInterpolator;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/view/InsetsSourceTest.java b/core/tests/coretests/src/android/view/InsetsSourceTest.java
index 61e05da..c3bd065 100644
--- a/core/tests/coretests/src/android/view/InsetsSourceTest.java
+++ b/core/tests/coretests/src/android/view/InsetsSourceTest.java
@@ -32,7 +32,7 @@
 import android.platform.test.annotations.Presubmit;
 import android.util.SparseArray;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/view/InsetsStateTest.java b/core/tests/coretests/src/android/view/InsetsStateTest.java
index 75749c7..1144ee1 100644
--- a/core/tests/coretests/src/android/view/InsetsStateTest.java
+++ b/core/tests/coretests/src/android/view/InsetsStateTest.java
@@ -61,7 +61,7 @@
 import android.util.SparseIntArray;
 import android.view.WindowInsets.Type;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.window.flags.Flags;
 
diff --git a/core/tests/coretests/src/android/view/MotionEventTest.java b/core/tests/coretests/src/android/view/MotionEventTest.java
index bad0485..d0f9a38 100644
--- a/core/tests/coretests/src/android/view/MotionEventTest.java
+++ b/core/tests/coretests/src/android/view/MotionEventTest.java
@@ -33,8 +33,8 @@
 import android.view.MotionEvent.PointerCoords;
 import android.view.MotionEvent.PointerProperties;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/PendingInsetsControllerTest.java b/core/tests/coretests/src/android/view/PendingInsetsControllerTest.java
index b5b2d0c..8ac9292 100644
--- a/core/tests/coretests/src/android/view/PendingInsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/PendingInsetsControllerTest.java
@@ -37,7 +37,7 @@
 import android.view.WindowInsetsController.OnControllableInsetsChangedListener;
 import android.view.animation.LinearInterpolator;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/view/RoundedCornerTest.java b/core/tests/coretests/src/android/view/RoundedCornerTest.java
index 4349021..3992aa1 100644
--- a/core/tests/coretests/src/android/view/RoundedCornerTest.java
+++ b/core/tests/coretests/src/android/view/RoundedCornerTest.java
@@ -23,8 +23,8 @@
 import android.graphics.Point;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/RoundedCornersTest.java b/core/tests/coretests/src/android/view/RoundedCornersTest.java
index ec665ad..c26d945 100644
--- a/core/tests/coretests/src/android/view/RoundedCornersTest.java
+++ b/core/tests/coretests/src/android/view/RoundedCornersTest.java
@@ -42,8 +42,8 @@
 import android.util.DisplayMetrics;
 import android.util.Pair;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.R;
 
diff --git a/core/tests/coretests/src/android/view/ScrollCaptureConnectionTest.java b/core/tests/coretests/src/android/view/ScrollCaptureConnectionTest.java
index 5f258949..bee5dc4 100644
--- a/core/tests/coretests/src/android/view/ScrollCaptureConnectionTest.java
+++ b/core/tests/coretests/src/android/view/ScrollCaptureConnectionTest.java
@@ -38,8 +38,8 @@
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/view/ScrollCaptureSearchResultsTest.java b/core/tests/coretests/src/android/view/ScrollCaptureSearchResultsTest.java
index dc43204..726ee85 100644
--- a/core/tests/coretests/src/android/view/ScrollCaptureSearchResultsTest.java
+++ b/core/tests/coretests/src/android/view/ScrollCaptureSearchResultsTest.java
@@ -35,8 +35,8 @@
 import android.platform.test.annotations.Presubmit;
 
 import androidx.annotation.NonNull;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/view/SurfaceControlRegistryTests.java b/core/tests/coretests/src/android/view/SurfaceControlRegistryTests.java
index 71bdce4..5a5510c 100644
--- a/core/tests/coretests/src/android/view/SurfaceControlRegistryTests.java
+++ b/core/tests/coretests/src/android/view/SurfaceControlRegistryTests.java
@@ -22,8 +22,8 @@
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.eq;
@@ -33,7 +33,7 @@
 import android.content.Context;
 import android.platform.test.annotations.Presubmit;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
diff --git a/core/tests/coretests/src/android/view/TunnelModeEnabledListenerTest.java b/core/tests/coretests/src/android/view/TunnelModeEnabledListenerTest.java
index 65dd34f..9ab14af 100644
--- a/core/tests/coretests/src/android/view/TunnelModeEnabledListenerTest.java
+++ b/core/tests/coretests/src/android/view/TunnelModeEnabledListenerTest.java
@@ -21,16 +21,16 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.MediumTest;
-import androidx.test.runner.AndroidJUnit4;
-
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 @RunWith(AndroidJUnit4.class)
 @MediumTest
 @Presubmit
diff --git a/core/tests/coretests/src/android/view/ViewCaptureTest.java b/core/tests/coretests/src/android/view/ViewCaptureTest.java
index 218047c..9144cd6 100644
--- a/core/tests/coretests/src/android/view/ViewCaptureTest.java
+++ b/core/tests/coretests/src/android/view/ViewCaptureTest.java
@@ -26,9 +26,9 @@
 import android.view.ViewDebug.HardwareCanvasProvider;
 import android.view.ViewDebug.SoftwareCanvasProvider;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.rule.ActivityTestRule;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 
diff --git a/core/tests/coretests/src/android/view/ViewFrameRateTest.java b/core/tests/coretests/src/android/view/ViewFrameRateTest.java
index 62291d4..b8ff595 100644
--- a/core/tests/coretests/src/android/view/ViewFrameRateTest.java
+++ b/core/tests/coretests/src/android/view/ViewFrameRateTest.java
@@ -50,10 +50,10 @@
 import android.widget.ProgressBar;
 
 import androidx.test.annotation.UiThreadTest;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
 import androidx.test.rule.ActivityTestRule;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 
@@ -278,7 +278,7 @@
     @RequiresFlagsEnabled({FLAG_VIEW_VELOCITY_API,
             FLAG_TOOLKIT_FRAME_RATE_VELOCITY_MAPPING_READ_ONLY,
             FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
-    public void lowVelocity80() throws Throwable {
+    public void lowVelocity60() throws Throwable {
         if (!ViewProperties.vrr_enabled().orElse(true)) {
             return;
         }
@@ -292,6 +292,31 @@
         mActivityRule.runOnUiThread(() -> {
             mMovingView.setFrameContentVelocity(1f);
             mMovingView.invalidate();
+            runAfterDraw(() -> assertEquals(60f, mViewRoot.getLastPreferredFrameRate(), 0f));
+        });
+        waitForAfterDraw();
+    }
+
+    @Test
+    @RequiresFlagsEnabled({FLAG_VIEW_VELOCITY_API,
+            FLAG_TOOLKIT_FRAME_RATE_VELOCITY_MAPPING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
+    public void midVelocity80() throws Throwable {
+        if (!ViewProperties.vrr_enabled().orElse(true)) {
+            return;
+        }
+        mActivityRule.runOnUiThread(() -> {
+            ViewGroup.LayoutParams layoutParams = mMovingView.getLayoutParams();
+            layoutParams.width = ViewGroup.LayoutParams.MATCH_PARENT;
+            layoutParams.height = ViewGroup.LayoutParams.MATCH_PARENT;
+            mMovingView.setLayoutParams(layoutParams);
+        });
+        waitForFrameRateCategoryToSettle();
+        mActivityRule.runOnUiThread(() -> {
+            float midSpeed =
+                    200f * mMovingView.getContext().getResources().getDisplayMetrics().density;
+            mMovingView.setFrameContentVelocity(midSpeed);
+            mMovingView.invalidate();
             runAfterDraw(() -> assertEquals(80f, mViewRoot.getLastPreferredFrameRate(), 0f));
         });
         waitForAfterDraw();
@@ -321,7 +346,7 @@
             frameLayout.setFrameContentVelocity(1f);
             mMovingView.offsetTopAndBottom(100);
             frameLayout.invalidate();
-            runAfterDraw(() -> assertEquals(80f, mViewRoot.getLastPreferredFrameRate(), 0f));
+            runAfterDraw(() -> assertEquals(60f, mViewRoot.getLastPreferredFrameRate(), 0f));
         });
         waitForAfterDraw();
     }
@@ -590,7 +615,7 @@
             runAfterDraw(() -> {
                 assertEquals(FRAME_RATE_CATEGORY_LOW,
                         mViewRoot.getLastPreferredFrameRateCategory());
-                assertEquals(80f, mViewRoot.getLastPreferredFrameRate());
+                assertEquals(60f, mViewRoot.getLastPreferredFrameRate());
             });
         });
         waitForAfterDraw();
diff --git a/core/tests/coretests/src/android/view/ViewGroupTransientViewTest.java b/core/tests/coretests/src/android/view/ViewGroupTransientViewTest.java
index 54524b2..f5c71f8 100644
--- a/core/tests/coretests/src/android/view/ViewGroupTransientViewTest.java
+++ b/core/tests/coretests/src/android/view/ViewGroupTransientViewTest.java
@@ -25,9 +25,9 @@
 import android.widget.FrameLayout;
 
 import androidx.test.annotation.UiThreadTest;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.MediumTest;
 import androidx.test.rule.ActivityTestRule;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Rule;
diff --git a/core/tests/coretests/src/android/view/ViewInvalidateTest.java b/core/tests/coretests/src/android/view/ViewInvalidateTest.java
index c25a2deb..d4181d3 100644
--- a/core/tests/coretests/src/android/view/ViewInvalidateTest.java
+++ b/core/tests/coretests/src/android/view/ViewInvalidateTest.java
@@ -31,9 +31,9 @@
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.annotation.UiThreadTest;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 import androidx.test.rule.ActivityTestRule;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Assert;
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java
index dd8cc6e..e5ad561 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityCacheTest.java
@@ -42,8 +42,8 @@
 import android.view.Display;
 import android.view.View;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.google.common.base.Throwables;
 
@@ -1053,6 +1053,28 @@
         assertFalse(mAccessibilityCache.isNodeInCache(childInfo));
     }
 
+    @Test
+    public void getEventSourceClassName_windowStateChangedThenRemoved() {
+        final String sourceActivityClassName = "com.example.SomeActivity";
+        final AccessibilityEvent windowStateChangedEvent = new AccessibilityEvent(
+                AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED);
+        final View mockView = getMockViewWithA11yAndWindowIds(PARENT_VIEW_ID, WINDOW_ID_1);
+        windowStateChangedEvent.setSource(mockView);
+        windowStateChangedEvent.setClassName(sourceActivityClassName);
+
+        mAccessibilityCache.onAccessibilityEvent(windowStateChangedEvent);
+        assertEquals(mAccessibilityCache.getEventSourceClassName(WINDOW_ID_1),
+                sourceActivityClassName);
+
+        final AccessibilityEvent windowRemovedEvent = new AccessibilityEvent(
+                AccessibilityEvent.TYPE_WINDOWS_CHANGED);
+        windowRemovedEvent.setWindowChanges(AccessibilityEvent.WINDOWS_CHANGE_REMOVED);
+        windowRemovedEvent.setSource(mockView);
+
+        mAccessibilityCache.onAccessibilityEvent(windowRemovedEvent);
+        assertNull(mAccessibilityCache.getEventSourceClassName(WINDOW_ID_1));
+    }
+
     private AccessibilityWindowInfo obtainAccessibilityWindowInfo(int windowId, int layer) {
         AccessibilityWindowInfo windowInfo = AccessibilityWindowInfo.obtain();
         windowInfo.setId(windowId);
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityEventTest.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityEventTest.java
index ddc27aa..3b8f66a 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityEventTest.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityEventTest.java
@@ -23,7 +23,7 @@
 import android.os.Parcel;
 import android.view.Display;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityInteractionClientTest.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityInteractionClientTest.java
index 3e061d2..eb482f2e 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityInteractionClientTest.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityInteractionClientTest.java
@@ -26,8 +26,8 @@
 import android.os.Bundle;
 import android.os.RemoteException;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import libcore.util.EmptyArray;
 
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityManagerTest.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityManagerTest.java
index ce36ee0..82e3427 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityManagerTest.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityManagerTest.java
@@ -55,7 +55,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.test.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.R;
 import com.android.internal.accessibility.common.ShortcutConstants;
diff --git a/core/tests/coretests/src/android/view/accessibility/AccessibilityNodeInfoTest.java b/core/tests/coretests/src/android/view/accessibility/AccessibilityNodeInfoTest.java
index 2d82d23..a5137bdf 100644
--- a/core/tests/coretests/src/android/view/accessibility/AccessibilityNodeInfoTest.java
+++ b/core/tests/coretests/src/android/view/accessibility/AccessibilityNodeInfoTest.java
@@ -28,8 +28,8 @@
 import android.util.ArraySet;
 import android.view.accessibility.AccessibilityNodeInfo.AccessibilityAction;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.util.CollectionUtils;
 
@@ -46,7 +46,7 @@
     // The number of fields tested in the corresponding CTS AccessibilityNodeInfoTest:
     // See fullyPopulateAccessibilityNodeInfo, assertEqualsAccessibilityNodeInfo,
     // and assertAccessibilityNodeInfoCleared in that class.
-    private static final int NUM_MARSHALLED_PROPERTIES = 44;
+    private static final int NUM_MARSHALLED_PROPERTIES = 43;
 
     /**
      * The number of properties that are purposely not marshalled
diff --git a/core/tests/coretests/src/android/view/autofill/AutofillStateFingerprintTest.java b/core/tests/coretests/src/android/view/autofill/AutofillStateFingerprintTest.java
new file mode 100644
index 0000000..7cbfc40
--- /dev/null
+++ b/core/tests/coretests/src/android/view/autofill/AutofillStateFingerprintTest.java
@@ -0,0 +1,155 @@
+/*
+ * Copyright (C) 2024 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 android.view.autofill;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+import android.content.Context;
+import android.text.InputType;
+import android.view.View;
+import android.view.inputmethod.EditorInfo;
+import android.widget.TextView;
+
+import androidx.test.core.app.ApplicationProvider;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class AutofillStateFingerprintTest {
+
+    private static final Context sContext = ApplicationProvider.getApplicationContext();
+
+    private static final int MAGIC_AUTOFILL_NUMBER = 1000;
+
+    private AutofillStateFingerprint mAutofillStateFingerprint =
+            AutofillStateFingerprint.createInstance();
+
+    @Test
+    public void testSameFingerprintsForTextView() throws Exception {
+        TextView tv = new TextView(sContext);
+        tv.setHint("Password");
+        tv.setSingleLine(true);
+        tv.setInputType(InputType.TYPE_TEXT_VARIATION_PASSWORD);
+        tv.setImeOptions(EditorInfo.IME_FLAG_NAVIGATE_NEXT);
+        fillViewProperties(tv);
+
+        // Create a copy Text View, and compare both id's
+        View tvCopy = copySelectiveViewAttributes(tv);
+        assertIdsEqual(tv, tvCopy);
+    }
+
+    @Test
+    public void testDifferentFingerprintsForTextViewWithDifferentHint() throws Exception {
+        TextView tv = new TextView(sContext);
+        tv.setHint("Password");
+        tv.setSingleLine(true);
+        tv.setInputType(InputType.TYPE_TEXT_VARIATION_PASSWORD);
+        tv.setImeOptions(EditorInfo.IME_FLAG_NAVIGATE_NEXT);
+        fillViewProperties(tv);
+
+        TextView tvCopy = (TextView) copySelectiveViewAttributes(tv);
+        tvCopy.setHint("what a useless different hint");
+        assertIdsNotEqual(tv, tvCopy);
+    }
+
+    @Test
+    public void testSameFingerprintsForNonTextView() throws Exception {
+        View v = new View(sContext);
+        fillViewProperties(v);
+
+        // Create a copy Text View, and compare both id's
+        View copy = copySelectiveViewAttributes(v);
+        assertIdsEqual(v, copy);
+    }
+
+    @Test
+    public void testDifferentFingerprintsForNonTextViewWithDifferentVisibility() throws Exception {
+        View v = new View(sContext);
+        fillViewProperties(v);
+
+        View copy = copySelectiveViewAttributes(v);
+        copy.setVisibility(View.GONE);
+        assertIdsNotEqual(v, copy);
+    }
+
+    private void assertIdsEqual(View v1, View v2) {
+        assertEquals(mAutofillStateFingerprint.getEphemeralFingerprintId(v1, 0),
+                mAutofillStateFingerprint.getEphemeralFingerprintId(v2, 0));
+    }
+
+    private void assertIdsNotEqual(View v1, View v2) {
+        assertNotEquals(mAutofillStateFingerprint.getEphemeralFingerprintId(v1, 0),
+                mAutofillStateFingerprint.getEphemeralFingerprintId(v2, 0));
+    }
+
+    private void fillViewProperties(View view) {
+        // Fill in relevant view properties
+        view.setContentDescription("ContentDesc");
+        view.setTooltipText("TooltipText");
+        view.setAutofillHints(new String[] {"password"});
+        view.setVisibility(View.VISIBLE);
+        view.setLeft(20);
+        view.setRight(200);
+        view.setTop(20);
+        view.setBottom(200);
+        view.setPadding(0, 1, 2, 3);
+    }
+
+    // Only copy interesting view attributes, particularly the view attributes that are critical
+    // for calculating fingerprint. Keep Autofill Id different.
+    private View copySelectiveViewAttributes(View view) {
+        View copy;
+        if (view instanceof TextView) {
+            copy = new TextView(sContext);
+            copySelectiveTextViewAttributes((TextView) view, (TextView) copy);
+        } else {
+            copy = new View(sContext) {
+                public @AutofillType int getAutofillType() {
+                    return view.getAutofillType();
+                }
+            };
+        }
+        // Copy over interested view properties.
+        // Keep the order same as with the tested code for easier clarity.
+        copy.setVisibility(view.getVisibility());
+        copy.setAutofillHints(view.getAutofillHints());
+        copy.setContentDescription(view.getContentDescription());
+        copy.setTooltip(view.getTooltipText());
+
+        copy.setRight(view.getRight());
+        copy.setLeft(view.getLeft());
+        copy.setTop(view.getTop());
+        copy.setBottom(view.getBottom());
+        copy.setPadding(view.getPaddingLeft(), view.getPaddingTop(),
+                view.getPaddingRight(), view.getPaddingBottom());
+
+        // DO not copy over autofill id
+        AutofillId newId = new AutofillId(view.getAutofillId().getViewId() + MAGIC_AUTOFILL_NUMBER);
+        copy.setAutofillId(newId);
+        return copy;
+    }
+
+    private void copySelectiveTextViewAttributes(TextView fromView, TextView toView) {
+        toView.setInputType(fromView.getInputType());
+        toView.setHint(fromView.getHint());
+        toView.setSingleLine(fromView.isSingleLine());
+        toView.setImeOptions(fromView.getImeOptions());
+    }
+}
diff --git a/core/tests/coretests/src/android/view/textclassifier/ConversationActionTest.java b/core/tests/coretests/src/android/view/textclassifier/ConversationActionTest.java
index e1b403f..c5e0ebd 100644
--- a/core/tests/coretests/src/android/view/textclassifier/ConversationActionTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/ConversationActionTest.java
@@ -26,8 +26,8 @@
 import android.os.Bundle;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/SelectionEventTest.java b/core/tests/coretests/src/android/view/textclassifier/SelectionEventTest.java
index 46e3a4c..1c26da7 100644
--- a/core/tests/coretests/src/android/view/textclassifier/SelectionEventTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/SelectionEventTest.java
@@ -20,8 +20,8 @@
 
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/SystemTextClassifierMetadataTest.java b/core/tests/coretests/src/android/view/textclassifier/SystemTextClassifierMetadataTest.java
index e4cfc53..4635e9b 100644
--- a/core/tests/coretests/src/android/view/textclassifier/SystemTextClassifierMetadataTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/SystemTextClassifierMetadataTest.java
@@ -20,11 +20,10 @@
 
 import static org.testng.Assert.assertThrows;
 
-
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextClassificationConstantsTest.java b/core/tests/coretests/src/android/view/textclassifier/TextClassificationConstantsTest.java
index 20a8768..07fe12f 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextClassificationConstantsTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextClassificationConstantsTest.java
@@ -20,8 +20,8 @@
 
 import android.provider.DeviceConfig;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextClassificationTest.java b/core/tests/coretests/src/android/view/textclassifier/TextClassificationTest.java
index 8225afc..2ed016c 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextClassificationTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextClassificationTest.java
@@ -38,8 +38,8 @@
 import android.view.View;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextClassifierEventTest.java b/core/tests/coretests/src/android/view/textclassifier/TextClassifierEventTest.java
index 11eb567..d92da6e 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextClassifierEventTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextClassifierEventTest.java
@@ -19,8 +19,8 @@
 
 import android.annotation.Nullable;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextClassifierUtilsTest.java b/core/tests/coretests/src/android/view/textclassifier/TextClassifierUtilsTest.java
index 011866d..ec46426 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextClassifierUtilsTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextClassifierUtilsTest.java
@@ -20,8 +20,8 @@
 
 import static org.testng.Assert.assertThrows;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextLanguageTest.java b/core/tests/coretests/src/android/view/textclassifier/TextLanguageTest.java
index 31f8029..de6f1d2 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextLanguageTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextLanguageTest.java
@@ -24,8 +24,8 @@
 import android.os.Bundle;
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextLinksTest.java b/core/tests/coretests/src/android/view/textclassifier/TextLinksTest.java
index 4f0b44b..bd1f7e1 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextLinksTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextLinksTest.java
@@ -25,8 +25,8 @@
 import android.os.Parcel;
 import android.util.ArrayMap;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java b/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java
index 14c077c..61e6738 100644
--- a/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java
+++ b/core/tests/coretests/src/android/view/textclassifier/TextSelectionTest.java
@@ -36,8 +36,8 @@
 import android.os.Parcel;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/accessibility/AccessibilityShortcutChooserActivityTest.java b/core/tests/coretests/src/com/android/internal/accessibility/AccessibilityShortcutChooserActivityTest.java
index 9f5ed29..37625e2 100644
--- a/core/tests/coretests/src/com/android/internal/accessibility/AccessibilityShortcutChooserActivityTest.java
+++ b/core/tests/coretests/src/com/android/internal/accessibility/AccessibilityShortcutChooserActivityTest.java
@@ -67,8 +67,8 @@
 
 import androidx.lifecycle.Lifecycle;
 import androidx.test.core.app.ActivityScenario;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.R;
 import com.android.internal.accessibility.dialog.AccessibilityShortcutChooserActivity;
diff --git a/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityTargetTest.java b/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityTargetTest.java
index f01ac6f..8608f6c 100644
--- a/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityTargetTest.java
+++ b/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityTargetTest.java
@@ -27,7 +27,7 @@
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Flags;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.accessibility.common.ShortcutConstants;
 
diff --git a/core/tests/coretests/src/com/android/internal/accessibility/dialog/InvisibleToggleAccessibilityServiceTargetTest.java b/core/tests/coretests/src/com/android/internal/accessibility/dialog/InvisibleToggleAccessibilityServiceTargetTest.java
index 9cac312..5339d91 100644
--- a/core/tests/coretests/src/com/android/internal/accessibility/dialog/InvisibleToggleAccessibilityServiceTargetTest.java
+++ b/core/tests/coretests/src/com/android/internal/accessibility/dialog/InvisibleToggleAccessibilityServiceTargetTest.java
@@ -41,8 +41,8 @@
 import android.view.accessibility.Flags;
 import android.view.accessibility.IAccessibilityManager;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.accessibility.TestUtils;
 import com.android.internal.accessibility.common.ShortcutConstants;
diff --git a/core/tests/coretests/src/com/android/internal/accessibility/util/AccessibilityUtilsTest.java b/core/tests/coretests/src/com/android/internal/accessibility/util/AccessibilityUtilsTest.java
index 58ab92a..f37ec9b 100644
--- a/core/tests/coretests/src/com/android/internal/accessibility/util/AccessibilityUtilsTest.java
+++ b/core/tests/coretests/src/com/android/internal/accessibility/util/AccessibilityUtilsTest.java
@@ -33,7 +33,7 @@
 import android.text.SpannableString;
 import android.text.style.LocaleSpan;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigParserTest.java b/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigParserTest.java
index 4eccbe5..a466caf 100644
--- a/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigParserTest.java
+++ b/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigParserTest.java
@@ -22,7 +22,7 @@
 
 import android.platform.test.annotations.Presubmit;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.content.om.OverlayConfigParser;
 
diff --git a/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigTest.java b/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigTest.java
index a0e9947..43cff8d 100644
--- a/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigTest.java
+++ b/core/tests/coretests/src/com/android/internal/content/res/OverlayConfigTest.java
@@ -28,7 +28,7 @@
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 import com.android.internal.content.om.OverlayConfig;
diff --git a/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java b/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
index 7f4e9ad..2f3b7f9 100644
--- a/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
+++ b/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
@@ -36,9 +36,9 @@
 import android.view.WindowManagerImpl;
 
 import androidx.test.core.app.ApplicationProvider;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.rule.ActivityTestRule;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Rule;
diff --git a/core/tests/coretests/src/com/android/internal/policy/PhoneWindowTest.java b/core/tests/coretests/src/com/android/internal/policy/PhoneWindowTest.java
index 4921e4a..e037f2a 100644
--- a/core/tests/coretests/src/com/android/internal/policy/PhoneWindowTest.java
+++ b/core/tests/coretests/src/com/android/internal/policy/PhoneWindowTest.java
@@ -44,8 +44,8 @@
 import android.view.WindowManager;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 
diff --git a/core/tests/coretests/src/com/android/internal/security/VerityUtilsTest.java b/core/tests/coretests/src/com/android/internal/security/VerityUtilsTest.java
index a978e3b..7b9ea55 100644
--- a/core/tests/coretests/src/com/android/internal/security/VerityUtilsTest.java
+++ b/core/tests/coretests/src/com/android/internal/security/VerityUtilsTest.java
@@ -21,9 +21,9 @@
 
 import android.platform.test.annotations.Presubmit;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.frameworks.coretests.R;
 
diff --git a/core/tests/coretests/src/com/android/internal/util/BitUtilsTest.java b/core/tests/coretests/src/com/android/internal/util/BitUtilsTest.java
index fdde36a..fdba811 100644
--- a/core/tests/coretests/src/com/android/internal/util/BitUtilsTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/BitUtilsTest.java
@@ -30,8 +30,8 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/CollectionUtilsTest.java b/core/tests/coretests/src/com/android/internal/util/CollectionUtilsTest.java
index ac954d6a..32c2d63 100644
--- a/core/tests/coretests/src/com/android/internal/util/CollectionUtilsTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/CollectionUtilsTest.java
@@ -22,7 +22,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/ContrastColorUtilTest.java b/core/tests/coretests/src/com/android/internal/util/ContrastColorUtilTest.java
index e6ebfef..aa59afe 100644
--- a/core/tests/coretests/src/com/android/internal/util/ContrastColorUtilTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/ContrastColorUtilTest.java
@@ -33,8 +33,8 @@
 import android.text.style.TextAppearanceSpan;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.R;
 
diff --git a/core/tests/coretests/src/com/android/internal/util/DumpUtilsTest.java b/core/tests/coretests/src/com/android/internal/util/DumpUtilsTest.java
index d2d3c13..7bd062a 100644
--- a/core/tests/coretests/src/com/android/internal/util/DumpUtilsTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/DumpUtilsTest.java
@@ -31,7 +31,7 @@
 import android.content.ComponentName;
 import android.util.SparseArray;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/DumpableContainerImplTest.java b/core/tests/coretests/src/com/android/internal/util/DumpableContainerImplTest.java
index 61d4e3d..9259181f 100644
--- a/core/tests/coretests/src/com/android/internal/util/DumpableContainerImplTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/DumpableContainerImplTest.java
@@ -19,7 +19,7 @@
 
 import android.util.Dumpable;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.util.dump.DumpableContainerImpl;
 
diff --git a/core/tests/coretests/src/com/android/internal/util/FakeLatencyTrackerTest.java b/core/tests/coretests/src/com/android/internal/util/FakeLatencyTrackerTest.java
index 6bd67ea..aee352b 100644
--- a/core/tests/coretests/src/com/android/internal/util/FakeLatencyTrackerTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/FakeLatencyTrackerTest.java
@@ -28,7 +28,7 @@
 import android.platform.test.ravenwood.RavenwoodRule;
 import android.provider.DeviceConfig;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/core/tests/coretests/src/com/android/internal/util/FastMathTest.java b/core/tests/coretests/src/com/android/internal/util/FastMathTest.java
index dd26334..bedcf4c 100644
--- a/core/tests/coretests/src/com/android/internal/util/FastMathTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/FastMathTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/GrowingArrayUtilsTest.java b/core/tests/coretests/src/com/android/internal/util/GrowingArrayUtilsTest.java
index 8456161..a0eb058 100644
--- a/core/tests/coretests/src/com/android/internal/util/GrowingArrayUtilsTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/GrowingArrayUtilsTest.java
@@ -23,7 +23,7 @@
 
 import android.util.EmptyArray;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/HexDumpTest.java b/core/tests/coretests/src/com/android/internal/util/HexDumpTest.java
index dcffa1c..9adf607 100644
--- a/core/tests/coretests/src/com/android/internal/util/HexDumpTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/HexDumpTest.java
@@ -22,7 +22,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/IntPairTest.java b/core/tests/coretests/src/com/android/internal/util/IntPairTest.java
index af6503f..527be8f 100644
--- a/core/tests/coretests/src/com/android/internal/util/IntPairTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/IntPairTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/LatencyTrackerTest.java b/core/tests/coretests/src/com/android/internal/util/LatencyTrackerTest.java
index 010f724..ce265a3 100644
--- a/core/tests/coretests/src/com/android/internal/util/LatencyTrackerTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/LatencyTrackerTest.java
@@ -29,7 +29,7 @@
 import android.platform.test.ravenwood.RavenwoodRule;
 import android.provider.DeviceConfig;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.internal.util.LatencyTracker.ActionProperties;
 
diff --git a/core/tests/coretests/src/com/android/internal/util/LineBreakBufferedWriterTest.java b/core/tests/coretests/src/com/android/internal/util/LineBreakBufferedWriterTest.java
index e6418fa..93262f0 100644
--- a/core/tests/coretests/src/com/android/internal/util/LineBreakBufferedWriterTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/LineBreakBufferedWriterTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/core/tests/coretests/src/com/android/internal/util/ParseUtilsTest.java b/core/tests/coretests/src/com/android/internal/util/ParseUtilsTest.java
index d24cbfe..b22014e 100644
--- a/core/tests/coretests/src/com/android/internal/util/ParseUtilsTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/ParseUtilsTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/ProgressReporterTest.java b/core/tests/coretests/src/com/android/internal/util/ProgressReporterTest.java
index 0d21335..e0d5499 100644
--- a/core/tests/coretests/src/com/android/internal/util/ProgressReporterTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/ProgressReporterTest.java
@@ -21,7 +21,7 @@
 import android.platform.test.annotations.IgnoreUnderRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Before;
 import org.junit.Rule;
diff --git a/core/tests/coretests/src/com/android/internal/util/RingBufferTest.java b/core/tests/coretests/src/com/android/internal/util/RingBufferTest.java
index d7a100a..4497770 100644
--- a/core/tests/coretests/src/com/android/internal/util/RingBufferTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/RingBufferTest.java
@@ -20,8 +20,8 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.fail;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.runner.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/SizedInputStreamTest.java b/core/tests/coretests/src/com/android/internal/util/SizedInputStreamTest.java
index efef7ff..dbd6fc1 100644
--- a/core/tests/coretests/src/com/android/internal/util/SizedInputStreamTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/SizedInputStreamTest.java
@@ -18,7 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/core/tests/coretests/src/com/android/internal/util/TokenBucketTest.java b/core/tests/coretests/src/com/android/internal/util/TokenBucketTest.java
index ef579fe..43ee3c5 100644
--- a/core/tests/coretests/src/com/android/internal/util/TokenBucketTest.java
+++ b/core/tests/coretests/src/com/android/internal/util/TokenBucketTest.java
@@ -24,7 +24,7 @@
 import android.os.SystemClock;
 import android.text.format.DateUtils;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/util/AcceptOnceConsumer.java b/libs/WindowManager/Jetpack/src/androidx/window/common/AcceptOnceConsumer.java
similarity index 98%
rename from libs/WindowManager/Jetpack/src/androidx/window/util/AcceptOnceConsumer.java
rename to libs/WindowManager/Jetpack/src/androidx/window/common/AcceptOnceConsumer.java
index 63828ab..c2f827a 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/util/AcceptOnceConsumer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/common/AcceptOnceConsumer.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package androidx.window.util;
+package androidx.window.common;
 
 import android.annotation.NonNull;
 
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/util/BaseDataProducer.java b/libs/WindowManager/Jetpack/src/androidx/window/common/BaseDataProducer.java
similarity index 98%
rename from libs/WindowManager/Jetpack/src/androidx/window/util/BaseDataProducer.java
rename to libs/WindowManager/Jetpack/src/androidx/window/common/BaseDataProducer.java
index cd26efd..e7099dc 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/util/BaseDataProducer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/common/BaseDataProducer.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package androidx.window.util;
+package androidx.window.common;
 
 import androidx.annotation.GuardedBy;
 import androidx.annotation.NonNull;
@@ -125,4 +125,4 @@
             mCallbacksToRemove.add(callback);
         }
     }
-}
\ No newline at end of file
+}
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/common/CommonFoldingFeature.java b/libs/WindowManager/Jetpack/src/androidx/window/common/CommonFoldingFeature.java
index e37dea4..b95bca1 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/common/CommonFoldingFeature.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/common/CommonFoldingFeature.java
@@ -16,7 +16,7 @@
 
 package androidx.window.common;
 
-import static androidx.window.util.ExtensionHelper.isZero;
+import static androidx.window.common.ExtensionHelper.isZero;
 
 import android.annotation.IntDef;
 import android.annotation.Nullable;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/common/DeviceStateManagerFoldingFeatureProducer.java b/libs/WindowManager/Jetpack/src/androidx/window/common/DeviceStateManagerFoldingFeatureProducer.java
index 98935e9..b2bc3de 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/common/DeviceStateManagerFoldingFeatureProducer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/common/DeviceStateManagerFoldingFeatureProducer.java
@@ -31,9 +31,6 @@
 import android.util.Log;
 import android.util.SparseIntArray;
 
-import androidx.window.util.AcceptOnceConsumer;
-import androidx.window.util.BaseDataProducer;
-
 import com.android.internal.R;
 
 import java.util.ArrayList;
@@ -44,7 +41,7 @@
 import java.util.function.Consumer;
 
 /**
- * An implementation of {@link androidx.window.util.BaseDataProducer} that returns
+ * An implementation of {@link BaseDataProducer} that returns
  * the device's posture by mapping the state returned from {@link DeviceStateManager} to
  * values provided in the resources' config at {@link R.array#config_device_state_postures}.
  */
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/util/ExtensionHelper.java b/libs/WindowManager/Jetpack/src/androidx/window/common/ExtensionHelper.java
similarity index 99%
rename from libs/WindowManager/Jetpack/src/androidx/window/util/ExtensionHelper.java
rename to libs/WindowManager/Jetpack/src/androidx/window/common/ExtensionHelper.java
index a08db79..f466d60 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/util/ExtensionHelper.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/common/ExtensionHelper.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package androidx.window.util;
+package androidx.window.common;
 
 import static android.view.Surface.ROTATION_270;
 import static android.view.Surface.ROTATION_90;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/common/RawFoldingFeatureProducer.java b/libs/WindowManager/Jetpack/src/androidx/window/common/RawFoldingFeatureProducer.java
index 88264f3..6d758f1 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/common/RawFoldingFeatureProducer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/common/RawFoldingFeatureProducer.java
@@ -26,15 +26,13 @@
 import android.provider.Settings;
 import android.text.TextUtils;
 
-import androidx.window.util.BaseDataProducer;
-
 import com.android.internal.R;
 
 import java.util.Optional;
 import java.util.function.Consumer;
 
 /**
- * Implementation of {@link androidx.window.util.BaseDataProducer} that produces a
+ * Implementation of {@link BaseDataProducer} that produces a
  * {@link String} that can be parsed to a {@link CommonFoldingFeature}.
  * {@link RawFoldingFeatureProducer} searches for the value in two places. The first check is in
  * settings where the {@link String} property is saved with the key
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
index 84984a9..a3ef68a 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
@@ -21,9 +21,9 @@
 
 import static androidx.window.common.CommonFoldingFeature.COMMON_STATE_FLAT;
 import static androidx.window.common.CommonFoldingFeature.COMMON_STATE_HALF_OPENED;
-import static androidx.window.util.ExtensionHelper.isZero;
-import static androidx.window.util.ExtensionHelper.rotateRectToDisplayRotation;
-import static androidx.window.util.ExtensionHelper.transformToWindowSpaceRect;
+import static androidx.window.common.ExtensionHelper.isZero;
+import static androidx.window.common.ExtensionHelper.rotateRectToDisplayRotation;
+import static androidx.window.common.ExtensionHelper.transformToWindowSpaceRect;
 
 import android.app.Activity;
 import android.app.ActivityThread;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SampleSidecarImpl.java b/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SampleSidecarImpl.java
index 339908a..b63fd08 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SampleSidecarImpl.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SampleSidecarImpl.java
@@ -25,11 +25,11 @@
 import android.os.IBinder;
 
 import androidx.annotation.NonNull;
+import androidx.window.common.BaseDataProducer;
 import androidx.window.common.CommonFoldingFeature;
 import androidx.window.common.DeviceStateManagerFoldingFeatureProducer;
 import androidx.window.common.EmptyLifecycleCallbacksAdapter;
 import androidx.window.common.RawFoldingFeatureProducer;
-import androidx.window.util.BaseDataProducer;
 
 import java.util.ArrayList;
 import java.util.List;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SidecarHelper.java b/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SidecarHelper.java
index bb6ab47..4fd03e4 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SidecarHelper.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/sidecar/SidecarHelper.java
@@ -17,8 +17,8 @@
 
 import static android.view.Display.DEFAULT_DISPLAY;
 
-import static androidx.window.util.ExtensionHelper.rotateRectToDisplayRotation;
-import static androidx.window.util.ExtensionHelper.transformToWindowSpaceRect;
+import static androidx.window.common.ExtensionHelper.rotateRectToDisplayRotation;
+import static androidx.window.common.ExtensionHelper.transformToWindowSpaceRect;
 
 import android.annotation.NonNull;
 import android.app.Activity;
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/util/ExtensionHelperTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/common/ExtensionHelperTest.java
similarity index 99%
rename from libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/util/ExtensionHelperTest.java
rename to libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/common/ExtensionHelperTest.java
index 3278cdf..b6e9519 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/util/ExtensionHelperTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/common/ExtensionHelperTest.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package androidx.window.util;
+package androidx.window.common;
 
 import static org.junit.Assert.assertEquals;
 
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
index b1c9a77..2b01eac 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
@@ -35,6 +35,7 @@
 ) {
   // All desktop mode related flags will be added here
   DESKTOP_WINDOWING_MODE(Flags::enableDesktopWindowingMode, true),
+  CASCADING_WINDOWS(Flags::enableCascadingWindows, true),
   WALLPAPER_ACTIVITY(Flags::enableDesktopWindowingWallpaperActivity, true),
   MODALS_POLICY(Flags::enableDesktopWindowingModalsPolicy, true),
   THEMED_APP_HEADERS(Flags::enableThemedAppHeaders, true),
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationAdapter.java b/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationAdapter.java
index 8d30db6..86e0f14 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationAdapter.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationAdapter.java
@@ -146,6 +146,11 @@
     /** To be overridden by subclasses to adjust the animation surface change. */
     void onAnimationUpdateInner(@NonNull SurfaceControl.Transaction t) {
         // Update the surface position and alpha.
+        if (com.android.graphics.libgui.flags.Flags.edgeExtensionShader()
+                && mAnimation.getExtensionEdges() != 0) {
+            t.setEdgeExtensionEffect(mLeash, mAnimation.getExtensionEdges());
+        }
+
         mTransformation.getMatrix().postTranslate(mContentRelOffset.x, mContentRelOffset.y);
         t.setMatrix(mLeash, mTransformation.getMatrix(), mMatrix);
         t.setAlpha(mLeash, mTransformation.getAlpha());
@@ -165,7 +170,7 @@
         if (!cropRect.intersect(mWholeAnimationBounds)) {
             // Hide the surface when it is outside of the animation area.
             t.setAlpha(mLeash, 0);
-        } else if (mAnimation.hasExtension()) {
+        } else if (mAnimation.getExtensionEdges() != 0) {
             // Allow the surface to be shown in its original bounds in case we want to use edge
             // extensions.
             cropRect.union(mContentBounds);
@@ -180,6 +185,7 @@
     @CallSuper
     void onAnimationEnd(@NonNull SurfaceControl.Transaction t) {
         onAnimationUpdate(t, mAnimation.getDuration());
+        t.setEdgeExtensionEffect(mLeash, /* edge */ 0);
     }
 
     final long getDurationHint() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationRunner.java b/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationRunner.java
index 5696a54..d2cef4b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationRunner.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingAnimationRunner.java
@@ -144,8 +144,10 @@
             // ending states.
             prepareForJumpCut(info, startTransaction);
         } else {
-            addEdgeExtensionIfNeeded(startTransaction, finishTransaction,
-                    postStartTransactionCallbacks, adapters);
+            if (!com.android.graphics.libgui.flags.Flags.edgeExtensionShader()) {
+                addEdgeExtensionIfNeeded(startTransaction, finishTransaction,
+                        postStartTransactionCallbacks, adapters);
+            }
             addBackgroundColorIfNeeded(info, startTransaction, finishTransaction, adapters);
             for (ActivityEmbeddingAnimationAdapter adapter : adapters) {
                 duration = Math.max(duration, adapter.getDurationHint());
@@ -341,7 +343,7 @@
             @NonNull List<ActivityEmbeddingAnimationAdapter> adapters) {
         for (ActivityEmbeddingAnimationAdapter adapter : adapters) {
             final Animation animation = adapter.mAnimation;
-            if (!animation.hasExtension()) {
+            if (animation.getExtensionEdges() == 0) {
                 continue;
             }
             if (adapter.mChange.hasFlags(FLAG_TRANSLUCENT)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
index 972dce5..24c568c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
@@ -169,6 +169,8 @@
                     R.layout.bubble_overflow_container, null /* root */);
             mOverflowView.initialize(expandedViewManager, positioner);
             addView(mOverflowView);
+            // Don't show handle for overflow
+            mHandleView.setVisibility(View.GONE);
         } else {
             mTaskView = bubbleTaskView.getTaskView();
             mBubbleTaskViewHelper = new BubbleTaskViewHelper(mContext, expandedViewManager,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
index 4fbb574..d7d19f7 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
@@ -315,7 +315,7 @@
                         }
                     }
                     if (!mImeShowing) {
-                        removeImeSurface();
+                        removeImeSurface(mDisplayId);
                     }
                 }
             } else if (!android.view.inputmethod.Flags.refactorInsetsController()
@@ -617,7 +617,7 @@
                                 || hasLeash) {
                             t.hide(mImeSourceControl.getLeash());
                         }
-                        removeImeSurface();
+                        removeImeSurface(mDisplayId);
                         ImeTracker.forLogging().onHidden(mStatsToken);
                     } else if (mAnimationDirection == DIRECTION_SHOW && !mCancelled) {
                         ImeTracker.forLogging().onShown(mStatsToken);
@@ -671,10 +671,10 @@
         }
     }
 
-    void removeImeSurface() {
+    void removeImeSurface(int displayId) {
         // Remove the IME surface to make the insets invisible for
         // non-client controlled insets.
-        InputMethodManagerGlobal.removeImeSurface(
+        InputMethodManagerGlobal.removeImeSurface(displayId,
                 e -> Slog.e(TAG, "Failed to remove IME surface.", e));
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/ImeListener.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ImeListener.kt
new file mode 100644
index 0000000..a34d7be
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ImeListener.kt
@@ -0,0 +1,70 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.common
+
+import android.graphics.Rect
+import android.view.InsetsSource
+import android.view.InsetsState
+import com.android.wm.shell.common.DisplayInsetsController.OnInsetsChangedListener
+
+abstract class ImeListener(
+    private val mDisplayController: DisplayController,
+    private val mDisplayId: Int
+) : OnInsetsChangedListener {
+    // The last insets state
+    private val mInsetsState = InsetsState()
+    private val mTmpBounds = Rect()
+
+    override fun insetsChanged(insetsState: InsetsState) {
+        if (mInsetsState == insetsState) {
+            return
+        }
+
+        // Get the stable bounds that account for display cutout and system bars to calculate the
+        // relative IME height
+        val layout = mDisplayController.getDisplayLayout(mDisplayId)
+        if (layout == null) {
+            return
+        }
+        layout.getStableBounds(mTmpBounds)
+
+        val wasVisible = getImeVisibilityAndHeight(mInsetsState).first
+        val oldHeight = getImeVisibilityAndHeight(mInsetsState).second
+
+        val isVisible = getImeVisibilityAndHeight(insetsState).first
+        val newHeight = getImeVisibilityAndHeight(insetsState).second
+
+        mInsetsState.set(insetsState, true)
+        if (wasVisible != isVisible || oldHeight != newHeight) {
+            onImeVisibilityChanged(isVisible, newHeight)
+        }
+    }
+
+    private fun getImeVisibilityAndHeight(
+            insetsState: InsetsState): Pair<Boolean, Int> {
+        val source = insetsState.peekSource(InsetsSource.ID_IME)
+        val frame = if (source != null && source.isVisible) source.frame else null
+        val height = if (frame != null) mTmpBounds.bottom - frame.top else 0
+        val visible = source?.isVisible ?: false
+        return Pair(visible, height)
+    }
+
+    /**
+     * To be overridden by implementations to handle IME changes.
+     */
+    protected abstract fun onImeVisibilityChanged(imeVisible: Boolean, imeHeight: Int)
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
index 5097ed8..19a109e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
@@ -31,6 +31,7 @@
 import android.app.ActivityManager;
 import android.content.Context;
 import android.content.res.Configuration;
+import android.graphics.Color;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
@@ -201,7 +202,7 @@
     /** Showing resizing hint. */
     public void onResizing(ActivityManager.RunningTaskInfo resizingTask, Rect newBounds,
             Rect sideBounds, SurfaceControl.Transaction t, int offsetX, int offsetY,
-            boolean immediately, float[] veilColor) {
+            boolean immediately) {
         if (mResizingIconView == null) {
             return;
         }
@@ -234,7 +235,7 @@
         if (mBackgroundLeash == null) {
             mBackgroundLeash = SurfaceUtils.makeColorLayer(mHostLeash,
                     RESIZING_BACKGROUND_SURFACE_NAME, mSurfaceSession);
-            t.setColor(mBackgroundLeash, veilColor)
+            t.setColor(mBackgroundLeash, getResizingBackgroundColor(resizingTask))
                     .setLayer(mBackgroundLeash, Integer.MAX_VALUE - 1);
         }
 
@@ -245,7 +246,7 @@
             mGapBackgroundLeash = SurfaceUtils.makeColorLayer(mHostLeash,
                     GAP_BACKGROUND_SURFACE_NAME, mSurfaceSession);
             // Fill up another side bounds area.
-            t.setColor(mGapBackgroundLeash, veilColor)
+            t.setColor(mGapBackgroundLeash, getResizingBackgroundColor(resizingTask))
                     .setLayer(mGapBackgroundLeash, Integer.MAX_VALUE - 2)
                     .setPosition(mGapBackgroundLeash, left, top)
                     .setWindowCrop(mGapBackgroundLeash, sideBounds.width(), sideBounds.height());
@@ -486,4 +487,9 @@
             mIcon = null;
         }
     }
+
+    private static float[] getResizingBackgroundColor(ActivityManager.RunningTaskInfo taskInfo) {
+        final int taskBgColor = taskInfo.taskDescription.getBackgroundColor();
+        return Color.valueOf(taskBgColor == -1 ? Color.WHITE : taskBgColor).getComponents();
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitScreenUtils.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitScreenUtils.java
index e8226051..f9259e7 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitScreenUtils.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitScreenUtils.java
@@ -16,6 +16,8 @@
 
 package com.android.wm.shell.common.split;
 
+import static android.content.pm.LauncherApps.ShortcutQuery.FLAG_MATCH_ALL_KINDS_WITH_ALL_PINNED;
+
 import static com.android.wm.shell.common.split.SplitScreenConstants.CONTROLLED_ACTIVITY_TYPES;
 import static com.android.wm.shell.common.split.SplitScreenConstants.CONTROLLED_WINDOWING_MODES;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
@@ -24,18 +26,25 @@
 
 import android.app.ActivityManager;
 import android.app.PendingIntent;
+import android.content.ComponentName;
 import android.content.Intent;
+import android.content.pm.LauncherApps;
+import android.content.pm.ShortcutInfo;
 import android.content.res.Configuration;
 import android.content.res.Resources;
-import android.graphics.Color;
 import android.graphics.Rect;
+import android.os.UserHandle;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.internal.util.ArrayUtils;
 import com.android.wm.shell.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
 
+import java.util.Arrays;
+import java.util.List;
+
 /** Helper utility class for split screen components to use. */
 public class SplitScreenUtils {
     /** Reverse the split position. */
@@ -128,10 +137,4 @@
             return isLandscape;
         }
     }
-
-    /** Returns the specified background color that matches a RunningTaskInfo. */
-    public static Color getResizingBackgroundColor(ActivityManager.RunningTaskInfo taskInfo) {
-        final int taskBgColor = taskInfo.taskDescription.getBackgroundColor();
-        return Color.valueOf(taskBgColor == -1 ? Color.WHITE : taskBgColor);
-    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
index f32683d..d289ef2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
@@ -95,7 +95,7 @@
         mCallback = callback;
         mHasSizeCompat = taskInfo.appCompatTaskInfo.topActivityInSizeCompat;
         if (DESKTOP_WINDOWING_MODE.isEnabled(mContext)
-                && DesktopModeFlags.THEMED_APP_HEADERS.isEnabled(context)) {
+                && DesktopModeFlags.DYNAMIC_INITIAL_BOUNDS.isEnabled(context)) {
             // Don't show the SCM button for freeform tasks
             mHasSizeCompat &= !taskInfo.isFreeform();
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index afe46f5..35d3876 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -569,7 +569,7 @@
             ShellTaskOrganizer shellTaskOrganizer) {
         int maxTaskLimit = DesktopModeStatus.getMaxTaskLimit(context);
         if (!DesktopModeStatus.canEnterDesktopMode(context)
-                || DESKTOP_WINDOWING_MODE.isEnabled(context)
+                || !DESKTOP_WINDOWING_MODE.isEnabled(context)
                 || maxTaskLimit <= 0) {
             return Optional.empty();
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
index 247cc42..4299841 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -32,77 +32,77 @@
 import java.util.concurrent.Executor
 import java.util.function.Consumer
 
-/** Keeps track of task data related to desktop mode. */
+/** Tracks task data for Desktop Mode. */
 class DesktopModeTaskRepository {
 
-    /** Task data that is tracked per display */
-    private data class DisplayData(
-        /**
-         * Set of task ids that are marked as active in desktop mode. Active tasks in desktop mode
-         * are freeform tasks that are visible or have been visible after desktop mode was
-         * activated. Task gets removed from this list when it vanishes. Or when desktop mode is
-         * turned off.
-         */
+    /**
+     * Task data tracked per desktop.
+     *
+     * @property activeTasks task ids of active tasks currently or previously visible in Desktop
+     * mode session. Tasks become inactive when task closes or when desktop mode session ends.
+     * @property visibleTasks task ids for active freeform tasks that are currently visible. There
+     * might be other active tasks in desktop mode that are not visible.
+     * @property minimizedTasks task ids for active freeform tasks that are currently minimized.
+     * @property closingTasks task ids for tasks that are going to close, but are currently visible.
+     * @property freeformTasksInZOrder list of current freeform task ids ordered from top to bottom
+     * (top is at index 0).
+     */
+    private data class DesktopTaskData(
         val activeTasks: ArraySet<Int> = ArraySet(),
         val visibleTasks: ArraySet<Int> = ArraySet(),
         val minimizedTasks: ArraySet<Int> = ArraySet(),
-        // Tasks that are closing, but are still visible
         // TODO(b/332682201): Remove when the repository state is updated via TransitionObserver
         val closingTasks: ArraySet<Int> = ArraySet(),
-        // Tasks currently in freeform mode, ordered from top to bottom (top is at index 0).
         val freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
     )
 
-    // Token of the current wallpaper activity, used to remove it when the last task is removed
+    /* Current wallpaper activity token to remove wallpaper activity when last task is removed. */
     var wallpaperActivityToken: WindowContainerToken? = null
+
     private val activeTasksListeners = ArraySet<ActiveTasksListener>()
-    // Track visible tasks separately because a task may be part of the desktop but not visible.
     private val visibleTasksListeners = ArrayMap<VisibleTasksListener, Executor>()
-    // Track corner/caption regions of desktop tasks, used to determine gesture exclusion
+
+    /* Tracks corner/caption regions of desktop tasks, used to determine gesture exclusion. */
     private val desktopExclusionRegions = SparseArray<Region>()
-    // Track last bounds of task before toggled to stable bounds
+
+    /* Tracks last bounds of task before toggled to stable bounds. */
     private val boundsBeforeMaximizeByTaskId = SparseArray<Rect>()
+
     private var desktopGestureExclusionListener: Consumer<Region>? = null
     private var desktopGestureExclusionExecutor: Executor? = null
 
-    private val displayData =
-        object : SparseArray<DisplayData>() {
-            /**
-             * Get the [DisplayData] associated with this [displayId]
-             *
-             * Creates a new instance if one does not exist
-             */
-            fun getOrCreate(displayId: Int): DisplayData {
-                if (!contains(displayId)) {
-                    put(displayId, DisplayData())
-                }
-                return get(displayId)
-            }
-        }
+    private val desktopTaskDataByDisplayId = object : SparseArray<DesktopTaskData>() {
+        /** Gets [DesktopTaskData] for existing [displayId] or creates a new one. */
+        fun getOrCreate(displayId: Int): DesktopTaskData =
+            this[displayId] ?: DesktopTaskData().also { this[displayId] = it }
+    }
 
-    /** Add a [ActiveTasksListener] to be notified of updates to active tasks in the repository. */
+    /** Adds [activeTasksListener] to be notified of updates to active tasks. */
     fun addActiveTaskListener(activeTasksListener: ActiveTasksListener) {
         activeTasksListeners.add(activeTasksListener)
     }
 
-    /** Add a [VisibleTasksListener] to be notified when freeform tasks are visible or not. */
+    /** Adds [visibleTasksListener] to be notified of updates to visible tasks. */
     fun addVisibleTasksListener(visibleTasksListener: VisibleTasksListener, executor: Executor) {
         visibleTasksListeners[visibleTasksListener] = executor
-        displayData.keyIterator().forEach {
+        desktopTaskDataByDisplayId.keyIterator().forEach {
+            val visibleTaskCount = getVisibleTaskCount(it)
             executor.execute {
-                visibleTasksListener.onTasksVisibilityChanged(it, visibleTaskCount(it))
+                visibleTasksListener.onTasksVisibilityChanged(it, visibleTaskCount)
             }
         }
     }
 
-    /** Returns a list of all [DisplayData]. */
-    private fun displayDataList(): Sequence<DisplayData> =
-        displayData.valueIterator().asSequence()
+    /** Updates tasks changes on all the active task listeners for given display id. */
+    private fun updateActiveTasksListeners(displayId: Int) {
+        activeTasksListeners.onEach { it.onActiveTasksChanged(displayId) }
+    }
 
-    /**
-     * Add a Consumer which will inform other classes of changes to exclusion regions for all
-     * Desktop tasks.
-     */
+    /** Returns a list of all [DesktopTaskData] in the repository. */
+    private fun desktopTaskDataSequence(): Sequence<DesktopTaskData> =
+        desktopTaskDataByDisplayId.valueIterator().asSequence()
+
+    /** Adds [regionListener] to inform about changes to exclusion regions for all Desktop tasks. */
     fun setExclusionRegionListener(regionListener: Consumer<Region>, executor: Executor) {
         desktopGestureExclusionListener = regionListener
         desktopGestureExclusionExecutor = executor
@@ -111,7 +111,7 @@
         }
     }
 
-    /** Create a new merged region representative of all exclusion regions in all desktop tasks. */
+    /** Creates a new merged region representative of all exclusion regions in all desktop tasks. */
     private fun calculateDesktopExclusionRegion(): Region {
         val desktopExclusionRegion = Region()
         desktopExclusionRegions.valueIterator().forEach { taskExclusionRegion ->
@@ -120,192 +120,120 @@
         return desktopExclusionRegion
     }
 
-    /** Remove a previously registered [ActiveTasksListener] */
+    /** Remove the previously registered [activeTasksListener] */
     fun removeActiveTasksListener(activeTasksListener: ActiveTasksListener) {
         activeTasksListeners.remove(activeTasksListener)
     }
 
-    /** Remove a previously registered [VisibleTasksListener] */
+    /** Removes the previously registered [visibleTasksListener]. */
     fun removeVisibleTasksListener(visibleTasksListener: VisibleTasksListener) {
         visibleTasksListeners.remove(visibleTasksListener)
     }
 
-    /**
-     * Mark a task with given [taskId] as active on given [displayId]
-     *
-     * @return `true` if the task was not active on given [displayId]
-     */
-    fun addActiveTask(displayId: Int, taskId: Int): Boolean {
-        // Check if task is active on another display, if so, remove it
-        displayData.forEach { id, data ->
-            if (id != displayId && data.activeTasks.remove(taskId)) {
-                activeTasksListeners.onEach { it.onActiveTasksChanged(id) }
+    /** Adds task with [taskId] to the list of active tasks on [displayId]. */
+    fun addActiveTask(displayId: Int, taskId: Int) {
+        // Removes task if it is active on another display excluding [displayId].
+        removeActiveTask(taskId, excludedDisplayId = displayId)
+
+        if (desktopTaskDataByDisplayId.getOrCreate(displayId).activeTasks.add(taskId)) {
+            logD("Adds active task=%d displayId=%d", taskId, displayId)
+            updateActiveTasksListeners(displayId)
+        }
+    }
+
+    /** Removes task from active task list of displays excluding the [excludedDisplayId]. */
+    fun removeActiveTask(taskId: Int, excludedDisplayId: Int? = null) {
+        desktopTaskDataByDisplayId.forEach { displayId, desktopTaskData ->
+            if ((displayId != excludedDisplayId)
+                    && desktopTaskData.activeTasks.remove(taskId)) {
+                logD("Removed active task=%d displayId=%d", taskId, displayId)
+                updateActiveTasksListeners(displayId)
             }
         }
-
-        val added = displayData.getOrCreate(displayId).activeTasks.add(taskId)
-        if (added) {
-            ProtoLog.d(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTaskRepo: add active task=%d displayId=%d",
-                taskId,
-                displayId
-            )
-            activeTasksListeners.onEach { it.onActiveTasksChanged(displayId) }
-        }
-        return added
     }
 
-    /**
-     * Remove task with given [taskId] from active tasks.
-     *
-     * @return `true` if the task was active
-     */
-    fun removeActiveTask(taskId: Int): Boolean {
-        var result = false
-        displayData.forEach { displayId, data ->
-            if (data.activeTasks.remove(taskId)) {
-                activeTasksListeners.onEach { it.onActiveTasksChanged(displayId) }
-                result = true
+    /** Adds given task to the closing task list for [displayId]. */
+    fun addClosingTask(displayId: Int, taskId: Int) {
+        if (desktopTaskDataByDisplayId.getOrCreate(displayId).closingTasks.add(taskId)) {
+            logD("Added closing task=%d displayId=%d", taskId, displayId)
+        } else {
+            // If the task hasn't been removed from closing list after it disappeared.
+            logW("Task with taskId=%d displayId=%d is already closing", taskId, displayId)
+        }
+    }
+
+    /** Removes task from the list of closing tasks for [displayId]. */
+    fun removeClosingTask(taskId: Int) {
+        desktopTaskDataByDisplayId.forEach { displayId, taskInfo ->
+            if (taskInfo.closingTasks.remove(taskId)) {
+                logD("Removed closing task=%d displayId=%d", taskId, displayId)
             }
         }
-        if (result) {
-            ProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTaskRepo: remove active task=%d", taskId)
-        }
-        return result
     }
 
-    /**
-     * Mark a task with given [taskId] as closing on given [displayId]
-     *
-     * @return `true` if the task was not closing on given [displayId]
-     */
-    fun addClosingTask(displayId: Int, taskId: Int): Boolean {
-        val added = displayData.getOrCreate(displayId).closingTasks.add(taskId)
-        if (added) {
-            ProtoLog.d(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTaskRepo: added closing task=%d displayId=%d",
-                taskId,
-                displayId
-            )
-        }
-        return added
-    }
+    fun isActiveTask(taskId: Int) = desktopTaskDataSequence().any { taskId in it.activeTasks }
+    fun isClosingTask(taskId: Int) = desktopTaskDataSequence().any { taskId in it.closingTasks }
+    fun isVisibleTask(taskId: Int) = desktopTaskDataSequence().any { taskId in it.visibleTasks }
+    fun isMinimizedTask(taskId: Int) = desktopTaskDataSequence().any { taskId in it.minimizedTasks }
 
-    /**
-     * Remove task with given [taskId] from closing tasks.
-     *
-     * @return `true` if the task was closing
-     */
-    fun removeClosingTask(taskId: Int): Boolean {
-        var removed = false
-        displayData.forEach { _, data ->
-            if (data.closingTasks.remove(taskId)) {
-                removed = true
-            }
-        }
-        if (removed) {
-            ProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTaskRepo: remove closing task=%d", taskId)
-        }
-        return removed
-    }
-
-    fun isActiveTask(taskId: Int) = displayDataList().any { taskId in it.activeTasks }
-    fun isClosingTask(taskId: Int) = displayDataList().any { taskId in it.closingTasks }
-    fun isVisibleTask(taskId: Int) = displayDataList().any { taskId in it.visibleTasks }
-    fun isMinimizedTask(taskId: Int) = displayDataList().any { taskId in it.minimizedTasks }
-
-    /**
-     * Check if a task with the given [taskId] is the only visible, non-closing, not-minimized task
-     * on its display
-     */
+    /** Checks if a task is the only visible, non-closing, non-minimized task on its display. */
     fun isOnlyVisibleNonClosingTask(taskId: Int): Boolean =
-        displayDataList().any { data ->
-            data.visibleTasks
-                .subtract(data.closingTasks)
-                .subtract(data.minimizedTasks)
-                .singleOrNull() == taskId
+        desktopTaskDataSequence().any { it.visibleTasks
+            .subtract(it.closingTasks)
+            .subtract(it.minimizedTasks)
+            .singleOrNull() == taskId
         }
 
-    /** Get a set of the active tasks for given [displayId] */
-    fun getActiveTasks(displayId: Int): ArraySet<Int> {
-        return ArraySet(displayData[displayId]?.activeTasks)
-    }
+    fun getActiveTasks(displayId: Int): ArraySet<Int> =
+        ArraySet(desktopTaskDataByDisplayId[displayId]?.activeTasks)
 
-    /** Returns the minimized tasks for the given [displayId]. */
     fun getMinimizedTasks(displayId: Int): ArraySet<Int> =
-        ArraySet(displayData[displayId]?.minimizedTasks)
+        ArraySet(desktopTaskDataByDisplayId[displayId]?.minimizedTasks)
 
-    /**
-     * Returns a list of Tasks IDs representing all active non-minimized Tasks on the given display,
-     * ordered from front to back.
-     */
-    fun getActiveNonMinimizedTasksOrderedFrontToBack(displayId: Int): List<Int> {
-        val activeTasks = getActiveTasks(displayId)
-        val allTasksInZOrder = getFreeformTasksInZOrder(displayId)
-        return activeTasks
-            // Don't show already minimized Tasks
-            .filter { taskId -> !isMinimizedTask(taskId) }
-            .sortedBy { taskId -> allTasksInZOrder.indexOf(taskId) }
+    /** Returns all active non-minimized tasks for [displayId] ordered from top to bottom. */
+    fun getActiveNonMinimizedOrderedTasks(displayId: Int): List<Int> =
+        getFreeformTasksInZOrder(displayId).filter { !isMinimizedTask(it) }
+
+    /** Returns a list of freeform tasks, ordered from top-bottom (top at index 0). */
+    fun getFreeformTasksInZOrder(displayId: Int): ArrayList<Int> =
+        ArrayList(desktopTaskDataByDisplayId[displayId]?.freeformTasksInZOrder ?: emptyList())
+
+    /** Removes task from visible tasks of all displays except [excludedDisplayId]. */
+    private fun removeVisibleTask(taskId: Int, excludedDisplayId: Int? = null) {
+        desktopTaskDataByDisplayId.forEach { displayId, data ->
+            if ((displayId != excludedDisplayId) && data.visibleTasks.remove(taskId)) {
+                notifyVisibleTaskListeners(displayId, data.visibleTasks.size)
+            }
+        }
     }
 
-    /** Get a list of freeform tasks, ordered from top-bottom (top at index 0). */
-    fun getFreeformTasksInZOrder(displayId: Int): ArrayList<Int> =
-        ArrayList(displayData[displayId]?.freeformTasksInZOrder ?: emptyList())
-
     /**
-     * Updates whether a freeform task with this id is visible or not and notifies listeners.
+     * Updates visibility of a freeform task with [taskId] on [displayId] and notifies listeners.
      *
-     * If the task was visible on a different display with a different displayId, it is removed from
-     * the set of visible tasks on that display. Listeners will be notified.
+     * If task was visible on a different display with a different [displayId], removes from
+     * the set of visible tasks on that display and notifies listeners.
      */
     fun updateVisibleFreeformTasks(displayId: Int, taskId: Int, visible: Boolean) {
         if (visible) {
-            // Task is visible. Check if we need to remove it from any other display.
-            val otherDisplays = displayData.keyIterator().asSequence().filter { it != displayId }
-            for (otherDisplayId in otherDisplays) {
-                if (displayData[otherDisplayId].visibleTasks.remove(taskId)) {
-                    notifyVisibleTaskListeners(
-                        otherDisplayId,
-                        displayData[otherDisplayId].visibleTasks.size
-                    )
-                }
-            }
+            // If task is visible, remove it from any other display besides [displayId].
+            removeVisibleTask(taskId, excludedDisplayId = displayId)
         } else if (displayId == INVALID_DISPLAY) {
             // Task has vanished. Check which display to remove the task from.
-            displayData.forEach { displayId, data ->
-                if (data.visibleTasks.remove(taskId)) {
-                    notifyVisibleTaskListeners(displayId, data.visibleTasks.size)
-                }
-            }
+            removeVisibleTask(taskId)
             return
         }
-
-        val prevCount = visibleTaskCount(displayId)
+        val prevCount = getVisibleTaskCount(displayId)
         if (visible) {
-            displayData.getOrCreate(displayId).visibleTasks.add(taskId)
+            desktopTaskDataByDisplayId.getOrCreate(displayId).visibleTasks.add(taskId)
             unminimizeTask(displayId, taskId)
         } else {
-            displayData[displayId]?.visibleTasks?.remove(taskId)
+            desktopTaskDataByDisplayId[displayId]?.visibleTasks?.remove(taskId)
         }
-        val newCount = visibleTaskCount(displayId)
-
-        // Check if count changed
+        val newCount = getVisibleTaskCount(displayId)
         if (prevCount != newCount) {
-            ProtoLog.d(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTaskRepo: update task visibility taskId=%d visible=%b displayId=%d",
-                taskId,
-                visible,
-                displayId
-            )
-            ProtoLog.d(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTaskRepo: visibleTaskCount has changed from %d to %d",
-                prevCount,
-                newCount
-            )
+            logD("Update task visibility taskId=%d visible=%b displayId=%d",
+                taskId, visible, displayId)
+            logD("VisibleTaskCount has changed from %d to %d", prevCount, newCount)
             notifyVisibleTaskListeners(displayId, newCount)
         }
     }
@@ -316,72 +244,46 @@
         }
     }
 
-    /** Get number of tasks that are marked as visible on given [displayId] */
-    fun visibleTaskCount(displayId: Int): Int {
-        ProtoLog.d(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTaskRepo: visibleTaskCount= %d",
-            displayData[displayId]?.visibleTasks?.size ?: 0
-        )
-        return displayData[displayId]?.visibleTasks?.size ?: 0
-    }
+    /** Gets number of visible tasks on given [displayId] */
+    fun getVisibleTaskCount(displayId: Int): Int =
+        desktopTaskDataByDisplayId[displayId]?.visibleTasks?.size ?: 0.also {
+            logD("getVisibleTaskCount=$it")
+        }
 
-    /** Add (or move if it already exists) the task to the top of the ordered list. */
-    // TODO(b/342417921): Identify if there is additional checks needed to move tasks for
-    // multi-display scenarios.
+    /** Adds task (or moves if it already exists) to the top of the ordered list. */
     fun addOrMoveFreeformTaskToTop(displayId: Int, taskId: Int) {
-        ProtoLog.d(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTaskRepo: add or move task to top: display=%d, taskId=%d",
-            displayId,
-            taskId
-        )
-        displayData[displayId]?.freeformTasksInZOrder?.remove(taskId)
-        displayData.getOrCreate(displayId).freeformTasksInZOrder.add(0, taskId)
+        logD("Add or move task to top: display=%d taskId=%d", taskId, displayId)
+        desktopTaskDataByDisplayId[displayId]?.freeformTasksInZOrder?.remove(taskId)
+        desktopTaskDataByDisplayId.getOrCreate(displayId).freeformTasksInZOrder.add(0, taskId)
     }
 
-    /** Mark a Task as minimized. */
+    /** Minimizes the task for [taskId] and [displayId] */
     fun minimizeTask(displayId: Int, taskId: Int) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopModeTaskRepository: minimize Task: display=%d, task=%d",
-            displayId,
-            taskId
-        )
-        displayData.getOrCreate(displayId).minimizedTasks.add(taskId)
+        logD("Minimize Task: display=%d, task=%d", displayId, taskId)
+        desktopTaskDataByDisplayId.getOrCreate(displayId).minimizedTasks.add(taskId)
     }
 
-    /** Mark a Task as non-minimized. */
+    /** Unminimizes the task for [taskId] and [displayId] */
     fun unminimizeTask(displayId: Int, taskId: Int) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopModeTaskRepository: unminimize Task: display=%d, task=%d",
-            displayId,
-            taskId
-        )
-        displayData[displayId]?.minimizedTasks?.remove(taskId)
+        logD("Unminimize Task: display=%d, task=%d", displayId, taskId)
+        desktopTaskDataByDisplayId[displayId]?.minimizedTasks?.remove(taskId) ?:
+            logW("Unminimize Task: display=%d, task=%d, no task data", displayId, taskId)
     }
 
-    /** Remove the task from the ordered list. */
+    /** Removes task from the ordered list. */
     fun removeFreeformTask(displayId: Int, taskId: Int) {
-        ProtoLog.d(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTaskRepo: remove freeform task from ordered list: display=%d, taskId=%d",
-            displayId,
-            taskId
-        )
-        displayData[displayId]?.freeformTasksInZOrder?.remove(taskId)
+        logD("Removes freeform task: taskId=%d", taskId)
+        desktopTaskDataByDisplayId[displayId]?.freeformTasksInZOrder?.remove(taskId)
         boundsBeforeMaximizeByTaskId.remove(taskId)
-        ProtoLog.d(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTaskRepo: remaining freeform tasks: %s",
-            displayData[displayId]?.freeformTasksInZOrder?.toDumpString() ?: ""
-        )
+        logD("Remaining freeform tasks: %d",
+            desktopTaskDataByDisplayId[displayId]?.freeformTasksInZOrder?.toDumpString() ?: "")
     }
 
     /**
-     * Updates the active desktop gesture exclusion regions; if desktopExclusionRegions has been
-     * accepted by desktopGestureExclusionListener, it will be updated in the appropriate classes.
+     * Updates active desktop gesture exclusion regions.
+     *
+     * If [desktopExclusionRegions] is accepted by [desktopGestureExclusionListener], updates it in
+     * appropriate classes.
      */
     fun updateTaskExclusionRegions(taskId: Int, taskExclusionRegions: Region) {
         desktopExclusionRegions.put(taskId, taskExclusionRegions)
@@ -391,9 +293,10 @@
     }
 
     /**
-     * Removes the desktop gesture exclusion region for the specified task; if exclusionRegion has
-     * been accepted by desktopGestureExclusionListener, it will be updated in the appropriate
-     * classes.
+     * Removes desktop gesture exclusion region for the specified task.
+     *
+     * If [desktopExclusionRegions] is accepted by [desktopGestureExclusionListener], updates it in
+     * appropriate classes.
      */
     fun removeExclusionRegion(taskId: Int) {
         desktopExclusionRegions.delete(taskId)
@@ -403,26 +306,24 @@
     }
 
     /** Removes and returns the bounds saved before maximizing the given task. */
-    fun removeBoundsBeforeMaximize(taskId: Int): Rect? {
-        return boundsBeforeMaximizeByTaskId.removeReturnOld(taskId)
-    }
+    fun removeBoundsBeforeMaximize(taskId: Int): Rect? =
+        boundsBeforeMaximizeByTaskId.removeReturnOld(taskId)
 
     /** Saves the bounds of the given task before maximizing. */
-    fun saveBoundsBeforeMaximize(taskId: Int, bounds: Rect) {
+    fun saveBoundsBeforeMaximize(taskId: Int, bounds: Rect) =
         boundsBeforeMaximizeByTaskId.set(taskId, Rect(bounds))
-    }
 
     internal fun dump(pw: PrintWriter, prefix: String) {
         val innerPrefix = "$prefix  "
         pw.println("${prefix}DesktopModeTaskRepository")
-        dumpDisplayData(pw, innerPrefix)
+        dumpDesktopTaskData(pw, innerPrefix)
         pw.println("${innerPrefix}activeTasksListeners=${activeTasksListeners.size}")
         pw.println("${innerPrefix}visibleTasksListeners=${visibleTasksListeners.size}")
     }
 
-    private fun dumpDisplayData(pw: PrintWriter, prefix: String) {
+    private fun dumpDesktopTaskData(pw: PrintWriter, prefix: String) {
         val innerPrefix = "$prefix  "
-        displayData.forEach { displayId, data ->
+        desktopTaskDataByDisplayId.forEach { displayId, data ->
             pw.println("${prefix}Display $displayId:")
             pw.println("${innerPrefix}activeTasks=${data.activeTasks.toDumpString()}")
             pw.println("${innerPrefix}visibleTasks=${data.visibleTasks.toDumpString()}")
@@ -432,23 +333,28 @@
         }
     }
 
-    /**
-     * Defines interface for classes that can listen to changes for active tasks in desktop mode.
-     */
+    /** Listens to changes for active tasks in desktop mode. */
     interface ActiveTasksListener {
-        /** Called when the active tasks change in desktop mode. */
         fun onActiveTasksChanged(displayId: Int) {}
     }
 
-    /**
-     * Defines interface for classes that can listen to changes for visible tasks in desktop mode.
-     */
+    /** Listens to changes for visible tasks in desktop mode. */
     interface VisibleTasksListener {
-        /** Called when the desktop changes the number of visible freeform tasks. */
         fun onTasksVisibilityChanged(displayId: Int, visibleTasksCount: Int) {}
     }
+
+    private fun logD(msg: String, vararg arguments: Any?) {
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
+    private fun logW(msg: String, vararg arguments: Any?) {
+        ProtoLog.w(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
+    companion object {
+        private const val TAG = "DesktopModeTaskRepository"
+    }
 }
 
-private fun <T> Iterable<T>.toDumpString(): String {
-    return joinToString(separator = ", ", prefix = "[", postfix = "]")
-}
+private fun <T> Iterable<T>.toDumpString(): String =
+    joinToString(separator = ", ", prefix = "[", postfix = "]")
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt
index da212e7..9fcf73d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt
@@ -52,8 +52,10 @@
     val idealSize = calculateIdealSize(screenBounds, scale)
     // If no top activity exists, apps fullscreen bounds and aspect ratio cannot be calculated.
     // Instead default to the desired initial bounds.
+    val stableBounds = Rect()
+    displayLayout.getStableBoundsForDesktopMode(stableBounds)
     val topActivityInfo =
-        taskInfo.topActivityInfo ?: return positionInScreen(idealSize, screenBounds)
+        taskInfo.topActivityInfo ?: return positionInScreen(idealSize, stableBounds)
 
     val initialSize: Size =
         when (taskInfo.configuration.orientation) {
@@ -100,7 +102,7 @@
             }
         }
 
-    return positionInScreen(initialSize, screenBounds)
+    return positionInScreen(initialSize, stableBounds)
 }
 
 /**
@@ -163,17 +165,11 @@
 }
 
 /** Adjusts bounds to be positioned in the middle of the screen. */
-private fun positionInScreen(desiredSize: Size, screenBounds: Rect): Rect {
-    // TODO(b/325240051): Position apps with bottom heavy offset
-    val heightOffset = (screenBounds.height() - desiredSize.height) / 2
-    val widthOffset = (screenBounds.width() - desiredSize.width) / 2
-    return Rect(
-        widthOffset,
-        heightOffset,
-        desiredSize.width + widthOffset,
-        desiredSize.height + heightOffset
-    )
-}
+private fun positionInScreen(desiredSize: Size, stableBounds: Rect): Rect =
+    Rect(0, 0, desiredSize.width, desiredSize.height).apply {
+        val offset = DesktopTaskPosition.Center.getTopLeftCoordinates(stableBounds, this)
+        offsetTo(offset.x, offset.y)
+    }
 
 /**
  * Adjusts bounds to be positioned in the middle of the area provided, not necessarily the
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTaskPosition.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTaskPosition.kt
new file mode 100644
index 0000000..97abda8
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTaskPosition.kt
@@ -0,0 +1,151 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.desktopmode
+
+import android.app.TaskInfo
+import android.content.res.Resources
+import android.graphics.Point
+import android.graphics.Rect
+import android.view.Gravity
+import com.android.internal.annotations.VisibleForTesting
+import com.android.wm.shell.desktopmode.DesktopTaskPosition.BottomLeft
+import com.android.wm.shell.desktopmode.DesktopTaskPosition.BottomRight
+import com.android.wm.shell.desktopmode.DesktopTaskPosition.Center
+import com.android.wm.shell.desktopmode.DesktopTaskPosition.TopLeft
+import com.android.wm.shell.desktopmode.DesktopTaskPosition.TopRight
+import com.android.wm.shell.R
+
+/**
+ * The position of a task window in desktop mode.
+ */
+sealed class DesktopTaskPosition {
+    data object Center : DesktopTaskPosition() {
+        private const val WINDOW_HEIGHT_PROPORTION = 0.375
+
+        override fun getTopLeftCoordinates(frame: Rect, window: Rect): Point {
+            val x = (frame.width() - window.width()) / 2
+            // Position with more margin at the bottom.
+            val y = (frame.height() - window.height()) * WINDOW_HEIGHT_PROPORTION + frame.top
+            return Point(x, y.toInt())
+        }
+
+        override fun next(): DesktopTaskPosition {
+            return BottomRight
+        }
+    }
+
+    data object BottomRight : DesktopTaskPosition() {
+        override fun getTopLeftCoordinates(frame: Rect, window: Rect): Point {
+            return Point(frame.right - window.width(), frame.bottom - window.height())
+        }
+
+        override fun next(): DesktopTaskPosition {
+            return TopLeft
+        }
+    }
+
+    data object TopLeft : DesktopTaskPosition() {
+        override fun getTopLeftCoordinates(frame: Rect, window: Rect): Point {
+            return Point(frame.left, frame.top)
+        }
+
+        override fun next(): DesktopTaskPosition {
+            return BottomLeft
+        }
+    }
+
+    data object BottomLeft : DesktopTaskPosition() {
+        override fun getTopLeftCoordinates(frame: Rect, window: Rect): Point {
+            return Point(frame.left, frame.bottom - window.height())
+        }
+
+        override fun next(): DesktopTaskPosition {
+            return TopRight
+        }
+    }
+
+    data object TopRight : DesktopTaskPosition() {
+        override fun getTopLeftCoordinates(frame: Rect, window: Rect): Point {
+            return Point(frame.right - window.width(), frame.top)
+        }
+
+        override fun next(): DesktopTaskPosition {
+            return Center
+        }
+    }
+
+    /**
+     * Returns the top left coordinates for the window to be placed in the given
+     * DesktopTaskPosition in the frame.
+     */
+    abstract fun getTopLeftCoordinates(frame: Rect, window: Rect): Point
+
+    abstract fun next(): DesktopTaskPosition
+}
+
+/**
+ * If the app has specified horizontal or vertical gravity layout, don't change the
+ * task position for cascading effect.
+ */
+fun canChangeTaskPosition(taskInfo: TaskInfo): Boolean {
+    taskInfo.topActivityInfo?.windowLayout?.let {
+        val horizontalGravityApplied = it.gravity.and(Gravity.HORIZONTAL_GRAVITY_MASK)
+        val verticalGravityApplied = it.gravity.and(Gravity.VERTICAL_GRAVITY_MASK)
+        return horizontalGravityApplied == 0 && verticalGravityApplied == 0
+    }
+    return true
+}
+
+/**
+ * Returns the current DesktopTaskPosition for a given window in the frame.
+ */
+@VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
+fun Rect.getDesktopTaskPosition(bounds: Rect): DesktopTaskPosition {
+    return when {
+        top == bounds.top && left == bounds.left -> TopLeft
+        top == bounds.top && right == bounds.right -> TopRight
+        bottom == bounds.bottom && left == bounds.left -> BottomLeft
+        bottom == bounds.bottom && right == bounds.right -> BottomRight
+        else -> Center
+    }
+}
+
+internal fun cascadeWindow(res: Resources, frame: Rect, prev: Rect, dest: Rect) {
+    val candidateBounds = Rect(dest)
+    val lastPos = frame.getDesktopTaskPosition(prev)
+    var destCoord = Center.getTopLeftCoordinates(frame, candidateBounds)
+    candidateBounds.offsetTo(destCoord.x, destCoord.y)
+    // If the default center position is not free or if last focused window is not at the
+    // center, get the next cascading window position.
+    if (!prevBoundsMovedAboveThreshold(res, prev, candidateBounds) || Center != lastPos) {
+        val nextCascadingPos = lastPos.next()
+        destCoord = nextCascadingPos.getTopLeftCoordinates(frame, dest)
+    }
+    dest.offsetTo(destCoord.x, destCoord.y)
+}
+
+internal fun prevBoundsMovedAboveThreshold(res: Resources, prev: Rect, newBounds: Rect): Boolean {
+    // This is the required minimum dp for a task to be touchable.
+    val moveThresholdPx = res.getDimensionPixelSize(
+        R.dimen.freeform_required_visible_empty_space_in_header)
+    val leftFar = newBounds.left - prev.left > moveThresholdPx
+    val topFar = newBounds.top - prev.top > moveThresholdPx
+    val rightFar = prev.right - newBounds.right > moveThresholdPx
+    val bottomFar = prev.bottom - newBounds.bottom > moveThresholdPx
+
+    return leftFar || topFar || rightFar || bottomFar
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 2ef045d..5f838d3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -121,7 +121,7 @@
     private val exitDesktopTaskTransitionHandler: ExitDesktopTaskTransitionHandler,
     private val toggleResizeDesktopTaskTransitionHandler: ToggleResizeDesktopTaskTransitionHandler,
     private val dragToDesktopTransitionHandler: DragToDesktopTransitionHandler,
-    private val desktopModeTaskRepository: DesktopModeTaskRepository,
+    private val taskRepository: DesktopModeTaskRepository,
     private val desktopModeLoggerTransitionObserver: DesktopModeLoggerTransitionObserver,
     private val launchAdjacentController: LaunchAdjacentController,
     private val recentsTransitionHandler: RecentsTransitionHandler,
@@ -181,7 +181,7 @@
     }
 
     private fun onInit() {
-        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "Initialize DesktopTasksController")
+        logD("onInit")
         shellCommandHandler.addDumpCallback(this::dump, this)
         shellCommandHandler.addCommandCallback("desktopmode", desktopModeShellCommandHandler, this)
         shellController.addExternalInterface(
@@ -190,16 +190,12 @@
             this
         )
         transitions.addHandler(this)
-        desktopModeTaskRepository.addVisibleTasksListener(taskVisibilityListener, mainExecutor)
+        taskRepository.addVisibleTasksListener(taskVisibilityListener, mainExecutor)
         dragToDesktopTransitionHandler.setDragToDesktopStateListener(dragToDesktopStateListener)
         recentsTransitionHandler.addTransitionStateListener(
             object : RecentsTransitionStateListener {
                 override fun onAnimationStateChanged(running: Boolean) {
-                    ProtoLog.v(
-                        WM_SHELL_DESKTOP_MODE,
-                        "DesktopTasksController: recents animation state changed running=%b",
-                        running
-                    )
+                    logV("Recents animation state changed running=%b", running)
                     recentsAnimationRunning = running
                 }
             }
@@ -227,7 +223,7 @@
     /** Returns the transition type for the given remote transition. */
     private fun transitionType(remoteTransition: RemoteTransition?): Int {
         if (remoteTransition == null) {
-            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: remoteTransition is null")
+            logV("RemoteTransition is null")
             return TRANSIT_NONE
         }
         return TRANSIT_TO_FRONT
@@ -235,7 +231,7 @@
 
     /** Show all tasks, that are part of the desktop, on top of launcher */
     fun showDesktopApps(displayId: Int, remoteTransition: RemoteTransition? = null) {
-        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: showDesktopApps")
+        logV("showDesktopApps")
         val wct = WindowContainerTransaction()
         bringDesktopAppsToFront(displayId, wct)
 
@@ -255,7 +251,7 @@
 
     /** Gets number of visible tasks in [displayId]. */
     fun visibleTaskCount(displayId: Int): Int =
-        desktopModeTaskRepository.visibleTaskCount(displayId)
+        taskRepository.getVisibleTaskCount(displayId)
 
     /** Returns true if any tasks are visible in Desktop Mode. */
     fun isDesktopModeShowing(displayId: Int): Boolean = visibleTaskCount(displayId) > 0
@@ -286,14 +282,8 @@
                     // Fullscreen case where we move the current focused task.
                     moveToDesktop(allFocusedTasks[0].taskId, transitionSource = transitionSource)
                 }
-                else -> {
-                    ProtoLog.w(
-                        WM_SHELL_DESKTOP_MODE,
-                        "DesktopTasksController: Cannot enter desktop, expected less " +
-                            "than 3 focused tasks but found %d",
-                        allFocusedTasks.size
-                    )
-                }
+                else -> logW("Cannot enter desktop, expected < 3 focused tasks, found %d",
+                    allFocusedTasks.size)
             }
         }
     }
@@ -317,11 +307,7 @@
         transitionSource: DesktopModeTransitionSource,
     ): Boolean {
         recentTasksController?.findTaskInBackground(taskId)?.let {
-            ProtoLog.v(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: moveToDesktopFromNonRunningTask taskId=%d",
-                taskId
-            )
+            logV("moveToDesktopFromNonRunningTask with taskId=%d, displayId=%d", taskId)
             // TODO(342378842): Instead of using default display, support multiple displays
             val taskToMinimize =
                 bringDesktopAppsToFrontBeforeShowingNewTask(DEFAULT_DISPLAY, wct, taskId)
@@ -351,18 +337,10 @@
     ) {
         if (DesktopModeFlags.MODALS_POLICY.isEnabled(context)
             && isTopActivityExemptFromDesktopWindowing(context, task)) {
-            ProtoLog.w(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: Cannot enter desktop, " +
-                        "ineligible top activity found."
-            )
+            logW("Cannot enter desktop for taskId %d, ineligible top activity found", task.taskId)
             return
         }
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: moveToDesktop taskId=%d",
-            task.taskId
-        )
+        logV("moveToDesktop taskId=%d", task.taskId)
         exitSplitIfApplicable(wct, task)
         // Bring other apps to front first
         val taskToMinimize =
@@ -386,11 +364,7 @@
         dragToDesktopValueAnimator: MoveToDesktopAnimator,
         taskSurface: SurfaceControl,
     ) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: startDragToDesktop taskId=%d",
-            taskInfo.taskId
-        )
+        logV("startDragToDesktop taskId=%d", taskInfo.taskId)
         interactionJankMonitor.begin(taskSurface, context,
             CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_HOLD)
         dragToDesktopTransitionHandler.startDragToDesktopTransition(
@@ -403,7 +377,7 @@
      * The second part of the animated drag to desktop transition, called after
      * [startDragToDesktop].
      */
-    private fun finalizeDragToDesktop(taskInfo: RunningTaskInfo, freeformBounds: Rect) {
+    private fun finalizeDragToDesktop(taskInfo: RunningTaskInfo) {
         ProtoLog.v(
             WM_SHELL_DESKTOP_MODE,
             "DesktopTasksController: finalizeDragToDesktop taskId=%d",
@@ -415,7 +389,6 @@
         val taskToMinimize =
             bringDesktopAppsToFrontBeforeShowingNewTask(taskInfo.displayId, wct, taskInfo.taskId)
         addMoveToDesktopChanges(wct, taskInfo)
-        wct.setBounds(taskInfo.token, freeformBounds)
         val transition = dragToDesktopTransitionHandler.finishDragToDesktopTransition(wct)
         transition?.let { addPendingMinimizeTransition(it, taskToMinimize) }
     }
@@ -443,17 +416,10 @@
      * @param taskId task id of the window that's being closed
      */
     fun onDesktopWindowClose(wct: WindowContainerTransaction, displayId: Int, taskId: Int) {
-        if (desktopModeTaskRepository.isOnlyVisibleNonClosingTask(taskId)) {
+        if (taskRepository.isOnlyVisibleNonClosingTask(taskId)) {
             removeWallpaperActivity(wct)
         }
-        if (!desktopModeTaskRepository.addClosingTask(displayId, taskId)) {
-            // Could happen if the task hasn't been removed from closing list after it disappeared
-            ProtoLog.w(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: the task with taskId=%d is already closing!",
-                taskId
-            )
-        }
+        taskRepository.addClosingTask(displayId, taskId)
     }
 
     /** Move a task with given `taskId` to fullscreen */
@@ -472,11 +438,7 @@
 
     /** Move a desktop app to split screen. */
     fun moveToSplit(task: RunningTaskInfo) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: moveToSplit taskId=%d",
-            task.taskId
-        )
+        logV( "moveToSplit taskId=%s", task.taskId)
         val wct = WindowContainerTransaction()
         wct.setBounds(task.token, Rect())
         // Rather than set windowing mode to multi-window at task level, set it to
@@ -505,11 +467,7 @@
      * [startDragToDesktop].
      */
     fun cancelDragToDesktop(task: RunningTaskInfo) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: cancelDragToDesktop taskId=%d",
-            task.taskId
-        )
+        logV("cancelDragToDesktop taskId=%d", task.taskId)
         dragToDesktopTransitionHandler.cancelDragToDesktopTransition(
             DragToDesktopTransitionHandler.CancelState.STANDARD_CANCEL
         )
@@ -520,11 +478,7 @@
         position: Point,
         transitionSource: DesktopModeTransitionSource
     ) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: moveToFullscreen with animation taskId=%d",
-            task.taskId
-        )
+        logV("moveToFullscreenWithAnimation taskId=%d", task.taskId)
         val wct = WindowContainerTransaction()
         addMoveToFullscreenChanges(wct, task)
 
@@ -548,12 +502,7 @@
 
     /** Move a task to the front */
     fun moveTaskToFront(taskInfo: RunningTaskInfo) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: moveTaskToFront taskId=%d",
-            taskInfo.taskId
-        )
-
+        logV("moveTaskToFront taskId=%s", taskInfo.taskId)
         val wct = WindowContainerTransaction()
         wct.reorder(taskInfo.token, true)
         val taskToMinimize = addAndGetMinimizeChangesIfNeeded(taskInfo.displayId, wct, taskInfo)
@@ -579,15 +528,10 @@
     fun moveToNextDisplay(taskId: Int) {
         val task = shellTaskOrganizer.getRunningTaskInfo(taskId)
         if (task == null) {
-            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToNextDisplay: taskId=%d not found", taskId)
+            logW("moveToNextDisplay: taskId=%d not found", taskId)
             return
         }
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "moveToNextDisplay: taskId=%d taskDisplayId=%d",
-            taskId,
-            task.displayId
-        )
+        logV("moveToNextDisplay: taskId=%d displayId=%d", taskId, task.displayId)
 
         val displayIds = rootTaskDisplayAreaOrganizer.displayIds.sorted()
         // Get the first display id that is higher than current task display id
@@ -597,7 +541,7 @@
             newDisplayId = displayIds.firstOrNull { displayId -> displayId < task.displayId }
         }
         if (newDisplayId == null) {
-            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToNextDisplay: next display not found")
+            logW("moveToNextDisplay: next display not found")
             return
         }
         moveToDisplay(task, newDisplayId)
@@ -609,21 +553,15 @@
      * No-op if task is already on that display per [RunningTaskInfo.displayId].
      */
     private fun moveToDisplay(task: RunningTaskInfo, displayId: Int) {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "moveToDisplay: taskId=%d displayId=%d",
-            task.taskId,
-            displayId
-        )
-
+        logV("moveToDisplay: taskId=%d displayId=%d", task.taskId, displayId)
         if (task.displayId == displayId) {
-            ProtoLog.d(WM_SHELL_DESKTOP_MODE, "moveToDisplay: task already on display")
+            logD("moveToDisplay: task already on display %d", displayId)
             return
         }
 
         val displayAreaInfo = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(displayId)
         if (displayAreaInfo == null) {
-            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "moveToDisplay: display not found")
+            logW("moveToDisplay: display not found")
             return
         }
 
@@ -660,7 +598,7 @@
             // If the task's pre-maximize stable bounds were saved, toggle the task to those bounds.
             // Otherwise, toggle to the default bounds.
             val taskBoundsBeforeMaximize =
-                desktopModeTaskRepository.removeBoundsBeforeMaximize(taskInfo.taskId)
+                taskRepository.removeBoundsBeforeMaximize(taskInfo.taskId)
             if (taskBoundsBeforeMaximize != null) {
                 destinationBounds.set(taskBoundsBeforeMaximize)
             } else {
@@ -673,7 +611,7 @@
         } else {
             // Save current bounds so that task can be restored back to original bounds if necessary
             // and toggle to the stable bounds.
-            desktopModeTaskRepository.saveBoundsBeforeMaximize(taskInfo.taskId, currentTaskBounds)
+            taskRepository.saveBoundsBeforeMaximize(taskInfo.taskId, currentTaskBounds)
 
             if (taskInfo.isResizeable) {
                 // if resizable then expand to entire stable bounds (full display minus insets)
@@ -778,12 +716,7 @@
         wct: WindowContainerTransaction,
         newTaskIdInFront: Int? = null
     ): RunningTaskInfo? {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: bringDesktopAppsToFront, newTaskIdInFront=%s",
-            newTaskIdInFront ?: "null"
-        )
-
+        logV("bringDesktopAppsToFront, newTaskId=%d", newTaskIdInFront)
         // Move home to front, ensures that we go back home when all desktop windows are closed
         moveHomeTask(wct, toTop = true)
 
@@ -794,7 +727,7 @@
         }
 
         val nonMinimizedTasksOrderedFrontToBack =
-            desktopModeTaskRepository.getActiveNonMinimizedTasksOrderedFrontToBack(displayId)
+            taskRepository.getActiveNonMinimizedOrderedTasks(displayId)
         // If we're adding a new Task we might need to minimize an old one
         val taskToMinimize: RunningTaskInfo? =
             if (newTaskIdInFront != null && desktopTasksLimiter.isPresent) {
@@ -812,7 +745,7 @@
             .filter { taskId -> taskId != taskToMinimize?.taskId }
             .mapNotNull { taskId -> shellTaskOrganizer.getRunningTaskInfo(taskId) }
             .reversed() // Start from the back so the front task is brought forward last
-            .forEach { task -> wct.reorder(task.token, true /* onTop */) }
+            .forEach { task -> wct.reorder(task.token, /* onTop= */ true) }
         return taskToMinimize
     }
 
@@ -820,11 +753,11 @@
         shellTaskOrganizer
             .getRunningTasks(context.displayId)
             .firstOrNull { task -> task.activityType == ACTIVITY_TYPE_HOME }
-            ?.let { homeTask -> wct.reorder(homeTask.getToken(), toTop /* onTop */) }
+            ?.let { homeTask -> wct.reorder(homeTask.getToken(), /* onTop= */ toTop) }
     }
 
     private fun addWallpaperActivity(wct: WindowContainerTransaction) {
-        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: addWallpaper")
+        logV("addWallpaperActivity")
         val intent = Intent(context, DesktopWallpaperActivity::class.java)
         val options =
             ActivityOptions.makeBasic().apply {
@@ -842,8 +775,8 @@
     }
 
     private fun removeWallpaperActivity(wct: WindowContainerTransaction) {
-        desktopModeTaskRepository.wallpaperActivityToken?.let { token ->
-            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: removeWallpaper")
+        taskRepository.wallpaperActivityToken?.let { token ->
+            logV("removeWallpaperActivity")
             wct.removeTask(token)
         }
     }
@@ -881,11 +814,7 @@
         transition: IBinder,
         request: TransitionRequestInfo
     ): WindowContainerTransaction? {
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: handleRequest request=%s",
-            request
-        )
+        logV("handleRequest request=%s", request)
         // Check if we should skip handling this transition
         var reason = ""
         val triggerTask = request.triggerTask
@@ -923,11 +852,7 @@
             }
 
         if (!shouldHandleRequest) {
-            ProtoLog.v(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: skipping handleRequest reason=%s",
-                reason
-            )
+            logV("skipping handleRequest reason=%s", reason)
             return null
         }
 
@@ -947,11 +872,7 @@
                     }
                 }
             }
-        ProtoLog.v(
-            WM_SHELL_DESKTOP_MODE,
-            "DesktopTasksController: handleRequest result=%s",
-            result ?: "null"
-        )
+        logV("handleRequest result=%s", result)
         return result
     }
 
@@ -985,25 +906,20 @@
         task: RunningTaskInfo,
         transition: IBinder
     ): WindowContainerTransaction? {
-        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleFreeformTaskLaunch")
+        logV("handleFreeformTaskLaunch")
         if (keyguardManager.isKeyguardLocked) {
             // Do NOT handle freeform task launch when locked.
             // It will be launched in fullscreen windowing mode (Details: b/160925539)
-            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: skip keyguard is locked")
+            logV("skip keyguard is locked")
             return null
         }
         val wct = WindowContainerTransaction()
         if (!isDesktopModeShowing(task.displayId)) {
-            ProtoLog.d(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: bring desktop tasks to front on transition" +
-                    " taskId=%d",
-                task.taskId
-            )
+            logD("Bring desktop tasks to front on transition=taskId=%d", task.taskId)
             // We are outside of desktop mode and already existing desktop task is being launched.
             // We should make this task go to fullscreen instead of freeform. Note that this means
             // any re-launch of a freeform window outside of desktop will be in fullscreen.
-            if (desktopModeTaskRepository.isActiveTask(task.taskId)) {
+            if (taskRepository.isActiveTask(task.taskId)) {
                 addMoveToFullscreenChanges(wct, task)
                 return wct
             }
@@ -1028,14 +944,9 @@
         task: RunningTaskInfo,
         transition: IBinder
     ): WindowContainerTransaction? {
-        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleFullscreenTaskLaunch")
+        logV("handleFullscreenTaskLaunch")
         if (isDesktopModeShowing(task.displayId)) {
-            ProtoLog.d(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: switch fullscreen task to freeform on transition" +
-                    " taskId=%d",
-                task.taskId
-            )
+            logD("Switch fullscreen task to freeform on transition: taskId=%d", task.taskId)
             return WindowContainerTransaction().also { wct ->
                 addMoveToDesktopChanges(wct, task)
                 // In some launches home task is moved behind new task being launched. Make sure
@@ -1062,33 +973,28 @@
 
     /** Handle task closing by removing wallpaper activity if it's the last active task */
     private fun handleTaskClosing(task: RunningTaskInfo): WindowContainerTransaction? {
-        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleTaskClosing")
+        logV("handleTaskClosing")
         val wct = WindowContainerTransaction()
-        if (desktopModeTaskRepository.isOnlyVisibleNonClosingTask(task.taskId)
-            && desktopModeTaskRepository.wallpaperActivityToken != null) {
+        if (taskRepository.isOnlyVisibleNonClosingTask(task.taskId)
+            && taskRepository.wallpaperActivityToken != null) {
             // Remove wallpaper activity when the last active task is removed
             removeWallpaperActivity(wct)
         }
-        if (!desktopModeTaskRepository.addClosingTask(task.displayId, task.taskId)) {
-            // Could happen if the task hasn't been removed from closing list after it disappeared
-            ProtoLog.w(
-                WM_SHELL_DESKTOP_MODE,
-                "DesktopTasksController: the task with taskId=%d is already closing!",
-                task.taskId
-            )
-        }
+        taskRepository.addClosingTask(task.displayId, task.taskId)
         // If a CLOSE or TO_BACK is triggered on a desktop task, remove the task.
         if (Flags.enableDesktopWindowingBackNavigation() &&
-            desktopModeTaskRepository.isVisibleTask(task.taskId)) {
+            taskRepository.isVisibleTask(task.taskId)) {
             wct.removeTask(task.token)
         }
         return if (wct.isEmpty) null else wct
     }
 
-    private fun addMoveToDesktopChanges(
+    @VisibleForTesting
+    fun addMoveToDesktopChanges(
         wct: WindowContainerTransaction,
         taskInfo: RunningTaskInfo
     ) {
+        val displayLayout = displayController.getDisplayLayout(taskInfo.displayId) ?: return
         val tdaInfo = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(taskInfo.displayId)!!
         val tdaWindowingMode = tdaInfo.configuration.windowConfiguration.windowingMode
         val targetWindowingMode =
@@ -1098,6 +1004,28 @@
             } else {
                 WINDOWING_MODE_FREEFORM
             }
+        val initialBounds = if (DesktopModeFlags.DYNAMIC_INITIAL_BOUNDS.isEnabled(context)) {
+            calculateInitialBounds(displayLayout, taskInfo)
+        } else {
+            getDefaultDesktopTaskBounds(displayLayout)
+        }
+
+        if (DesktopModeFlags.CASCADING_WINDOWS.isEnabled(context)) {
+            val stableBounds = Rect()
+            displayLayout.getStableBoundsForDesktopMode(stableBounds)
+
+            val activeTasks = taskRepository
+                .getActiveNonMinimizedOrderedTasks(taskInfo.displayId)
+            activeTasks.firstOrNull()?.let { activeTask ->
+                shellTaskOrganizer.getRunningTaskInfo(activeTask)?.let {
+                    cascadeWindow(context.resources, stableBounds,
+                        it.configuration.windowConfiguration.bounds, initialBounds)
+                }
+            }
+        }
+        if (canChangeTaskPosition(taskInfo)) {
+            wct.setBounds(taskInfo.token, initialBounds)
+        }
         wct.setWindowingMode(taskInfo.token, targetWindowingMode)
         wct.reorder(taskInfo.token, true /* onTop */)
         if (useDesktopOverrideDensity()) {
@@ -1123,7 +1051,7 @@
         if (useDesktopOverrideDensity()) {
             wct.setDensityDpi(taskInfo.token, getDefaultDensityDpi())
         }
-        if (desktopModeTaskRepository.isOnlyVisibleNonClosingTask(taskInfo.taskId)) {
+        if (taskRepository.isOnlyVisibleNonClosingTask(taskInfo.taskId)) {
             // Remove wallpaper activity when leaving desktop mode
             removeWallpaperActivity(wct)
         }
@@ -1142,7 +1070,7 @@
         // The task's density may have been overridden in freeform; revert it here as we don't
         // want it overridden in multi-window.
         wct.setDensityDpi(taskInfo.token, getDefaultDensityDpi())
-        if (desktopModeTaskRepository.isOnlyVisibleNonClosingTask(taskInfo.taskId)) {
+        if (taskRepository.isOnlyVisibleNonClosingTask(taskInfo.taskId)) {
             // Remove wallpaper activity when leaving desktop mode
             removeWallpaperActivity(wct)
         }
@@ -1357,16 +1285,10 @@
         val indicatorType = indicator.updateIndicatorType(inputCoordinates, taskInfo.windowingMode)
         when (indicatorType) {
             IndicatorType.TO_DESKTOP_INDICATOR -> {
-                val displayLayout = displayController.getDisplayLayout(taskInfo.displayId)
-                    ?: return IndicatorType.NO_INDICATOR
                 // Start a new jank interaction for the drag release to desktop window animation.
                 interactionJankMonitor.begin(taskSurface, context,
                     CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_RELEASE, "to_desktop")
-                if (DesktopModeFlags.DYNAMIC_INITIAL_BOUNDS.isEnabled(context)) {
-                    finalizeDragToDesktop(taskInfo, calculateInitialBounds(displayLayout, taskInfo))
-                } else {
-                    finalizeDragToDesktop(taskInfo, getDefaultDesktopTaskBounds(displayLayout))
-                }
+                finalizeDragToDesktop(taskInfo)
             }
             IndicatorType.NO_INDICATOR,
             IndicatorType.TO_FULLSCREEN_INDICATOR -> {
@@ -1384,12 +1306,12 @@
 
     /** Update the exclusion region for a specified task */
     fun onExclusionRegionChanged(taskId: Int, exclusionRegion: Region) {
-        desktopModeTaskRepository.updateTaskExclusionRegions(taskId, exclusionRegion)
+        taskRepository.updateTaskExclusionRegions(taskId, exclusionRegion)
     }
 
     /** Remove a previously tracked exclusion region for a specified task. */
     fun removeExclusionRegionForTask(taskId: Int) {
-        desktopModeTaskRepository.removeExclusionRegion(taskId)
+        taskRepository.removeExclusionRegion(taskId)
     }
 
     /**
@@ -1399,7 +1321,7 @@
      * @param callbackExecutor the executor to call the listener on.
      */
     fun addVisibleTasksListener(listener: VisibleTasksListener, callbackExecutor: Executor) {
-        desktopModeTaskRepository.addVisibleTasksListener(listener, callbackExecutor)
+        taskRepository.addVisibleTasksListener(listener, callbackExecutor)
     }
 
     /**
@@ -1409,7 +1331,7 @@
      * @param callbackExecutor the executor to call the listener on.
      */
     fun setTaskRegionListener(listener: Consumer<Region>, callbackExecutor: Executor) {
-        desktopModeTaskRepository.setExclusionRegionListener(listener, callbackExecutor)
+        taskRepository.setExclusionRegionListener(listener, callbackExecutor)
     }
 
     override fun onUnhandledDrag(
@@ -1427,7 +1349,7 @@
         if (!multiInstanceHelper.supportsMultiInstanceSplit(launchComponent)) {
             // TODO(b/320797628): Should only return early if there is an existing running task, and
             //                    notify the user as well. But for now, just ignore the drop.
-            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "Dropped intent does not support multi-instance")
+            logV("Dropped intent does not support multi-instance")
             return false
         }
 
@@ -1459,7 +1381,7 @@
     private fun dump(pw: PrintWriter, prefix: String) {
         val innerPrefix = "$prefix  "
         pw.println("${prefix}DesktopTasksController")
-        desktopModeTaskRepository.dump(pw, innerPrefix)
+        taskRepository.dump(pw, innerPrefix)
     }
 
     /** The interface for calls from outside the shell, within the host process. */
@@ -1534,12 +1456,12 @@
                 SingleInstanceRemoteListener<DesktopTasksController, IDesktopTaskListener>(
                     controller,
                     { c ->
-                        c.desktopModeTaskRepository.addVisibleTasksListener(
+                        c.taskRepository.addVisibleTasksListener(
                             listener,
                             c.mainExecutor
                         )
                     },
-                    { c -> c.desktopModeTaskRepository.removeVisibleTasksListener(listener) }
+                    { c -> c.taskRepository.removeVisibleTasksListener(listener) }
                 )
         }
 
@@ -1566,10 +1488,7 @@
         }
 
         override fun hideStashedDesktopApps(displayId: Int) {
-            ProtoLog.w(
-                WM_SHELL_DESKTOP_MODE,
-                "IDesktopModeImpl: hideStashedDesktopApps is deprecated"
-            )
+            ProtoLog.w(WM_SHELL_DESKTOP_MODE, "IDesktopModeImpl: hideStashedDesktopApps is deprecated")
         }
 
         override fun getVisibleTaskCount(displayId: Int): Int {
@@ -1593,11 +1512,7 @@
         }
 
         override fun setTaskListener(listener: IDesktopTaskListener?) {
-            ProtoLog.v(
-                WM_SHELL_DESKTOP_MODE,
-                "IDesktopModeImpl: set task listener=%s",
-                listener ?: "null"
-            )
+            ProtoLog.v(WM_SHELL_DESKTOP_MODE, "IDesktopModeImpl: set task listener=%s", listener)
             executeRemoteCallWithTaskPermission(controller, "setTaskListener") { _ ->
                 listener?.let { remoteListener.register(it) } ?: remoteListener.unregister()
             }
@@ -1610,10 +1525,22 @@
         }
     }
 
+    private fun logV(msg: String, vararg arguments: Any?) {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+    private fun logD(msg: String, vararg arguments: Any?) {
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+    private fun logW(msg: String, vararg arguments: Any?) {
+        ProtoLog.w(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
     companion object {
         @JvmField
         val DESKTOP_MODE_INITIAL_BOUNDS_SCALE =
             SystemProperties.getInt("persist.wm.debug.desktop_mode_initial_bounds_scale", 75) / 100f
+
+        private const val TAG = "DesktopTasksController"
     }
 
     /** The positions on a screen that a task can snap to. */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
index c5ed1be..a011ff5 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
@@ -65,10 +65,10 @@
         }
 
         override fun onTransitionReady(
-                transition: IBinder,
-                info: TransitionInfo,
-                startTransaction: SurfaceControl.Transaction,
-                finishTransaction: SurfaceControl.Transaction
+            transition: IBinder,
+            info: TransitionInfo,
+            startTransaction: SurfaceControl.Transaction,
+            finishTransaction: SurfaceControl.Transaction
         ) {
             val taskToMinimize = mPendingTransitionTokensAndTasks.remove(transition) ?: return
 
@@ -129,8 +129,7 @@
         }
 
         fun removeLeftoverMinimizedTasks(displayId: Int, wct: WindowContainerTransaction) {
-            if (taskRepository
-                .getActiveNonMinimizedTasksOrderedFrontToBack(displayId).isNotEmpty()) {
+            if (taskRepository.getActiveNonMinimizedOrderedTasks(displayId).isNotEmpty()) {
                 return
             }
             val remainingMinimizedTasks = taskRepository.getMinimizedTasks(displayId)
@@ -178,7 +177,7 @@
                 "DesktopTasksLimiter: addMinimizeBackTaskChangesIfNeeded, newFrontTask=%d",
                 newFrontTaskInfo.taskId)
         val newTaskListOrderedFrontToBack = createOrderedTaskListWithGivenTaskInFront(
-                taskRepository.getActiveNonMinimizedTasksOrderedFrontToBack(displayId),
+                taskRepository.getActiveNonMinimizedOrderedTasks(displayId),
                 newFrontTaskInfo.taskId)
         val taskToMinimize = getTaskToMinimizeIfNeeded(newTaskListOrderedFrontToBack)
         if (taskToMinimize != null) {
@@ -242,7 +241,5 @@
     }
 
     @VisibleForTesting
-    fun getTransitionObserver(): TransitionObserver {
-        return minimizeTransitionObserver
-    }
+    fun getTransitionObserver(): TransitionObserver = minimizeTransitionObserver
 }
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
index e4aa115..d03a561 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
@@ -21,11 +21,11 @@
 import static android.content.pm.ActivityInfo.CONFIG_ASSETS_PATHS;
 import static android.content.pm.ActivityInfo.CONFIG_UI_MODE;
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
+import static android.view.ViewTreeObserver.InternalInsetsInfo.TOUCHABLE_INSETS_FRAME;
 import static android.view.ViewTreeObserver.InternalInsetsInfo.TOUCHABLE_INSETS_REGION;
 
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
-import static com.android.wm.shell.common.split.SplitScreenUtils.getResizingBackgroundColor;
 import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_BOTTOM;
 import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_LEFT;
 import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_RIGHT;
@@ -40,6 +40,7 @@
 import android.content.Context;
 import android.content.res.Configuration;
 import android.content.res.Resources;
+import android.graphics.Color;
 import android.graphics.Insets;
 import android.graphics.Point;
 import android.graphics.Rect;
@@ -290,7 +291,7 @@
                 final int activityType = taskInfo1.getActivityType();
                 if (activityType == ACTIVITY_TYPE_STANDARD) {
                     Drawable icon1 = mIconProvider.getIcon(taskInfo1.topActivityInfo);
-                    int bgColor1 = getResizingBackgroundColor(taskInfo1).toArgb();
+                    int bgColor1 = getResizingBackgroundColor(taskInfo1);
                     mDropZoneView1.setAppInfo(bgColor1, icon1);
                     mDropZoneView2.setAppInfo(bgColor1, icon1);
                     mDropZoneView1.setForceIgnoreBottomMargin(false);
@@ -312,10 +313,10 @@
                     mSplitScreenController.getTaskInfo(SPLIT_POSITION_BOTTOM_OR_RIGHT);
             if (topOrLeftTask != null && bottomOrRightTask != null) {
                 Drawable topOrLeftIcon = mIconProvider.getIcon(topOrLeftTask.topActivityInfo);
-                int topOrLeftColor = getResizingBackgroundColor(topOrLeftTask).toArgb();
+                int topOrLeftColor = getResizingBackgroundColor(topOrLeftTask);
                 Drawable bottomOrRightIcon = mIconProvider.getIcon(
                         bottomOrRightTask.topActivityInfo);
-                int bottomOrRightColor = getResizingBackgroundColor(bottomOrRightTask).toArgb();
+                int bottomOrRightColor = getResizingBackgroundColor(bottomOrRightTask);
                 mDropZoneView1.setAppInfo(topOrLeftColor, topOrLeftIcon);
                 mDropZoneView2.setAppInfo(bottomOrRightColor, bottomOrRightIcon);
             }
@@ -586,6 +587,11 @@
         }
     }
 
+    private static int getResizingBackgroundColor(ActivityManager.RunningTaskInfo taskInfo) {
+        final int taskBgColor = taskInfo.taskDescription.getBackgroundColor();
+        return Color.valueOf(taskBgColor == -1 ? Color.WHITE : taskBgColor).toArgb();
+    }
+
     /**
      * Dumps information about this drag layout.
      */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index 4531967..229d972 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -101,12 +101,9 @@
                 repository.addOrMoveFreeformTaskToTop(taskInfo.displayId, taskInfo.taskId);
                 repository.unminimizeTask(taskInfo.displayId, taskInfo.taskId);
                 if (taskInfo.isVisible) {
-                    if (repository.addActiveTask(taskInfo.displayId, taskInfo.taskId)) {
-                        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
-                                "Adding active freeform task: #%d", taskInfo.taskId);
-                    }
+                    repository.addActiveTask(taskInfo.displayId, taskInfo.taskId);
                     repository.updateVisibleFreeformTasks(taskInfo.displayId, taskInfo.taskId,
-                            true);
+                        true);
                 }
             });
         }
@@ -122,10 +119,7 @@
             mDesktopModeTaskRepository.ifPresent(repository -> {
                 repository.removeFreeformTask(taskInfo.displayId, taskInfo.taskId);
                 repository.unminimizeTask(taskInfo.displayId, taskInfo.taskId);
-                if (repository.removeActiveTask(taskInfo.taskId)) {
-                    ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
-                            "Removing active freeform task: #%d", taskInfo.taskId);
-                }
+                repository.removeActiveTask(taskInfo.taskId, /* excludedDisplayId= */ null);
                 repository.updateVisibleFreeformTasks(taskInfo.displayId, taskInfo.taskId, false);
             });
         }
@@ -146,14 +140,9 @@
         if (DesktopModeStatus.canEnterDesktopMode(mContext)) {
             mDesktopModeTaskRepository.ifPresent(repository -> {
                 if (taskInfo.isVisible) {
-                    if (repository.addActiveTask(taskInfo.displayId, taskInfo.taskId)) {
-                        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
-                                "Adding active freeform task: #%d", taskInfo.taskId);
-                    }
-                } else if (repository.isClosingTask(taskInfo.taskId)
-                        && repository.removeClosingTask(taskInfo.taskId)) {
-                    ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
-                            "Removing closing freeform task: #%d", taskInfo.taskId);
+                    repository.addActiveTask(taskInfo.displayId, taskInfo.taskId);
+                } else if (repository.isClosingTask(taskInfo.taskId)) {
+                    repository.removeClosingTask(taskInfo.taskId);
                 }
                 repository.updateVisibleFreeformTasks(taskInfo.displayId, taskInfo.taskId,
                         taskInfo.isVisible);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipMotionHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipMotionHelper.java
index df3803d..999ab95 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipMotionHelper.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipMotionHelper.java
@@ -416,6 +416,17 @@
         // location now.
         mSpringingToTouch = false;
 
+        // Boost the velocityX if it's zero to forcefully push it towards the nearest edge.
+        // We don't simply change the xEndValue below since the PhysicsAnimator would rely on the
+        // same velocityX to find out which edge to snap to.
+        if (velocityX == 0) {
+            final int motionCenterX = mPipBoundsState
+                    .getMotionBoundsState().getBoundsInMotion().centerX();
+            final int displayCenterX = mPipBoundsState
+                    .getDisplayBounds().centerX();
+            velocityX = (motionCenterX < displayCenterX) ? -0.001f : 0.001f;
+        }
+
         mTemporaryBoundsPhysicsAnimator
                 .spring(FloatProperties.RECT_WIDTH, getBounds().width(), mSpringConfig)
                 .spring(FloatProperties.RECT_HEIGHT, getBounds().height(), mSpringConfig)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipController.java
index b939b16..8aa0933 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipController.java
@@ -44,6 +44,7 @@
 import com.android.wm.shell.common.DisplayInsetsController;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.ExternalInterfaceBinder;
+import com.android.wm.shell.common.ImeListener;
 import com.android.wm.shell.common.RemoteCallable;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SingleInstanceRemoteListener;
@@ -56,7 +57,6 @@
 import com.android.wm.shell.common.pip.PipDisplayLayoutState;
 import com.android.wm.shell.common.pip.PipUtils;
 import com.android.wm.shell.pip.Pip;
-import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.sysui.ConfigurationChangeListener;
 import com.android.wm.shell.sysui.ShellCommandHandler;
@@ -201,6 +201,11 @@
                                         .getDisplayLayout(mPipDisplayLayoutState.getDisplayId()));
                     }
                 });
+        mDisplayInsetsController.addInsetsChangedListener(mPipDisplayLayoutState.getDisplayId(),
+                new ImeListener(mDisplayController, mPipDisplayLayoutState.getDisplayId()) {
+                    @Override
+                    public void onImeVisibilityChanged(boolean imeVisible, int imeHeight) {}
+                });
 
         // Allow other outside processes to bind to PiP controller using the key below.
         mShellController.addExternalInterface(KEY_EXTRA_SHELL_PIP,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java
index ea02de9d..e1e072a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipMotionHelper.java
@@ -416,6 +416,17 @@
         // location now.
         mSpringingToTouch = false;
 
+        // Boost the velocityX if it's zero to forcefully push it towards the nearest edge.
+        // We don't simply change the xEndValue below since the PhysicsAnimator would rely on the
+        // same velocityX to find out which edge to snap to.
+        if (velocityX == 0) {
+            final int motionCenterX = mPipBoundsState
+                    .getMotionBoundsState().getBoundsInMotion().centerX();
+            final int displayCenterX = mPipBoundsState
+                    .getDisplayBounds().centerX();
+            velocityX = (motionCenterX < displayCenterX) ? -0.001f : 0.001f;
+        }
+
         mTemporaryBoundsPhysicsAnimator
                 .spring(FloatProperties.RECT_WIDTH, getBounds().width(), mSpringConfig)
                 .spring(FloatProperties.RECT_HEIGHT, getBounds().height(), mSpringConfig)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index d7ee563..2531ff1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -41,7 +41,6 @@
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
 import static com.android.wm.shell.common.split.SplitScreenConstants.splitPositionToString;
-import static com.android.wm.shell.common.split.SplitScreenUtils.getResizingBackgroundColor;
 import static com.android.wm.shell.common.split.SplitScreenUtils.reverseSplitPosition;
 import static com.android.wm.shell.common.split.SplitScreenUtils.splitFailureMessage;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN;
@@ -2458,13 +2457,8 @@
         updateSurfaceBounds(layout, t, shouldUseParallaxEffect);
         getMainStageBounds(mTempRect1);
         getSideStageBounds(mTempRect2);
-        // TODO (b/307490004): "commonColor" below is a temporary fix to ensure the colors on both
-        //  sides match. When b/307490004 is fixed, this code can be reverted.
-        float[] commonColor = getResizingBackgroundColor(mSideStage.mRootTaskInfo).getComponents();
-        mMainStage.onResizing(
-                mTempRect1, mTempRect2, t, offsetX, offsetY, mShowDecorImmediately, commonColor);
-        mSideStage.onResizing(
-                mTempRect2, mTempRect1, t, offsetX, offsetY, mShowDecorImmediately, commonColor);
+        mMainStage.onResizing(mTempRect1, mTempRect2, t, offsetX, offsetY, mShowDecorImmediately);
+        mSideStage.onResizing(mTempRect2, mTempRect1, t, offsetX, offsetY, mShowDecorImmediately);
         t.apply();
         mTransactionPool.release(t);
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index 1076eca..d1ab3e9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -314,10 +314,10 @@
     }
 
     void onResizing(Rect newBounds, Rect sideBounds, SurfaceControl.Transaction t, int offsetX,
-            int offsetY, boolean immediately, float[] veilColor) {
+            int offsetY, boolean immediately) {
         if (mSplitDecorManager != null && mRootTaskInfo != null) {
             mSplitDecorManager.onResizing(mRootTaskInfo, newBounds, sideBounds, t, offsetX,
-                    offsetY, immediately, veilColor);
+                    offsetY, immediately);
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
index 7784784..d8c8c60 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
@@ -502,7 +502,8 @@
                 backgroundColorForTransition = getTransitionBackgroundColorIfSet(info, change, a,
                         backgroundColorForTransition);
 
-                if (!isTask && a.hasExtension()) {
+                if (!com.android.graphics.libgui.flags.Flags.edgeExtensionShader() && !isTask
+                        && a.getExtensionEdges() != 0) {
                     if (!TransitionUtil.isOpeningType(mode)) {
                         // Can screenshot now (before startTransaction is applied)
                         edgeExtendWindow(change, a, startTransaction, finishTransaction);
@@ -512,6 +513,8 @@
                         postStartTransactionCallbacks
                                 .add(t -> edgeExtendWindow(change, a, t, finishTransaction));
                     }
+                } else if (com.android.graphics.libgui.flags.Flags.edgeExtensionShader()) {
+                    finishTransaction.setEdgeExtensionEffect(change.getLeash(), /* edge */ 0);
                 }
 
                 final Rect clipRect = TransitionUtil.isClosingType(mode)
@@ -1008,6 +1011,10 @@
             Point position, float cornerRadius, @Nullable Rect immutableClipRect) {
         tmpTransformation.clear();
         anim.getTransformation(time, tmpTransformation);
+        if (anim.getExtensionEdges() != 0
+                && com.android.graphics.libgui.flags.Flags.edgeExtensionShader()) {
+            t.setEdgeExtensionEffect(leash, anim.getExtensionEdges());
+        }
         if (position != null) {
             tmpTransformation.getMatrix().postTranslate(position.x, position.y);
         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
index 2c0aa12..764d5a9 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
@@ -83,7 +83,7 @@
             }
         }, mExecutor) {
             @Override
-            void removeImeSurface() { }
+            void removeImeSurface(int displayId) { }
         }.new PerDisplay(DEFAULT_DISPLAY, ROTATION_0);
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/ImeListenerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/ImeListenerTest.kt
new file mode 100644
index 0000000..3b0a072
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/ImeListenerTest.kt
@@ -0,0 +1,152 @@
+/*
+ * Copyright (C) 2024 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.wm.shell.common
+
+import android.content.res.Configuration
+import android.content.res.Resources
+import android.graphics.Insets
+import android.graphics.Rect
+import android.testing.AndroidTestingRunner
+import android.view.DisplayCutout
+import android.view.DisplayInfo
+import android.view.InsetsSource.ID_IME
+import android.view.InsetsState
+import android.view.Surface
+import android.view.WindowInsets.Type
+import androidx.test.filters.SmallTest
+import com.android.internal.R
+import com.android.wm.shell.ShellTestCase
+import junit.framework.Assert.assertEquals
+import junit.framework.Assert.assertFalse
+import junit.framework.Assert.assertTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito
+import org.mockito.kotlin.whenever
+
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class ImeListenerTest : ShellTestCase() {
+    private lateinit var imeListener: CachingImeListener
+    private lateinit var displayLayout: DisplayLayout
+
+    @Mock private lateinit var displayController: DisplayController
+    @Before
+    fun setUp() {
+        val resources = createResources(40, 50, false)
+        val displayInfo = createDisplayInfo(1000, 1500, 0, Surface.ROTATION_0)
+        displayLayout = DisplayLayout(displayInfo, resources, false, false)
+        whenever(displayController.getDisplayLayout(DEFAULT_DISPLAY_ID)).thenReturn(displayLayout)
+        imeListener = CachingImeListener(displayController, DEFAULT_DISPLAY_ID)
+    }
+
+    @Test
+    fun testImeAppears() {
+        val insetsState = createInsetsStateWithIme(true, DEFAULT_IME_HEIGHT)
+        imeListener.insetsChanged(insetsState)
+        assertTrue("Ime insets source should become visible", imeListener.cachedImeVisible)
+        assertEquals(DEFAULT_IME_HEIGHT, imeListener.cachedImeHeight)
+    }
+
+    @Test
+    fun testImeAppears_thenDisappears() {
+        // Send insetsState with an IME as a visible source.
+        val insetsStateWithIme = createInsetsStateWithIme(true, DEFAULT_IME_HEIGHT)
+        imeListener.insetsChanged(insetsStateWithIme)
+
+        // Send insetsState without IME.
+        val insetsStateWithoutIme = createInsetsStateWithIme(false, 0)
+        imeListener.insetsChanged(insetsStateWithoutIme)
+
+        assertFalse("Ime insets source should become invisible",
+                imeListener.cachedImeVisible)
+        assertEquals(0, imeListener.cachedImeHeight)
+    }
+
+    private fun createInsetsStateWithIme(isVisible: Boolean, imeHeight: Int): InsetsState {
+        val stableBounds = Rect()
+        displayLayout.getStableBounds(stableBounds)
+        val insetsState = InsetsState()
+
+        val insetsSource = insetsState.getOrCreateSource(ID_IME, Type.ime())
+        insetsSource.setVisible(isVisible)
+        insetsSource.setFrame(stableBounds.left, stableBounds.bottom - imeHeight,
+                stableBounds.right, stableBounds.bottom)
+        return insetsState
+    }
+
+    private fun createDisplayInfo(width: Int, height: Int, cutoutHeight: Int,
+                                  rotation: Int): DisplayInfo {
+        val info = DisplayInfo()
+        info.logicalWidth = width
+        info.logicalHeight = height
+        info.rotation = rotation
+        if (cutoutHeight > 0) {
+            info.displayCutout = DisplayCutout(
+                    Insets.of(0, cutoutHeight, 0, 0) /* safeInsets */,
+                    null /* boundLeft */,
+                    Rect(width / 2 - cutoutHeight, 0, width / 2 + cutoutHeight,
+                            cutoutHeight) /* boundTop */, null /* boundRight */,
+                    null /* boundBottom */)
+        } else {
+            info.displayCutout = DisplayCutout.NO_CUTOUT
+        }
+        info.logicalDensityDpi = 300
+        return info
+    }
+
+    private fun createResources(navLand: Int, navPort: Int, navMoves: Boolean): Resources {
+        val cfg = Configuration()
+        cfg.uiMode = Configuration.UI_MODE_TYPE_NORMAL
+        val res = Mockito.mock(Resources::class.java)
+        Mockito.doReturn(navLand).whenever(res).getDimensionPixelSize(
+                R.dimen.navigation_bar_height_landscape_car_mode)
+        Mockito.doReturn(navPort).whenever(res).getDimensionPixelSize(
+                R.dimen.navigation_bar_height_car_mode)
+        Mockito.doReturn(navLand).whenever(res).getDimensionPixelSize(
+                R.dimen.navigation_bar_width_car_mode)
+        Mockito.doReturn(navLand).whenever(res).getDimensionPixelSize(
+                R.dimen.navigation_bar_height_landscape)
+        Mockito.doReturn(navPort).whenever(res).getDimensionPixelSize(
+                R.dimen.navigation_bar_height)
+        Mockito.doReturn(navLand).whenever(res).getDimensionPixelSize(
+                R.dimen.navigation_bar_width)
+        Mockito.doReturn(navMoves).whenever(res).getBoolean(R.bool.config_navBarCanMove)
+        Mockito.doReturn(cfg).whenever(res).configuration
+        return res
+    }
+
+    private class CachingImeListener(
+            displayController: DisplayController,
+            displayId: Int
+    ) : ImeListener(displayController, displayId) {
+        var cachedImeVisible = false
+        var cachedImeHeight = 0
+        public override fun onImeVisibilityChanged(imeVisible: Boolean, imeHeight: Int) {
+            cachedImeVisible = imeVisible
+            cachedImeHeight = imeHeight
+        }
+    }
+
+    companion object {
+        private const val DEFAULT_DISPLAY_ID = 0
+        private const val DEFAULT_IME_HEIGHT = 500
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
index 18b08bf..0a5672d 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
@@ -41,32 +41,44 @@
     }
 
     @Test
-    fun addActiveTask_listenerNotifiedAndTaskIsActive() {
+    fun addActiveTask_notifiesListener() {
         val listener = TestListener()
         repo.addActiveTaskListener(listener)
 
         repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
+
         assertThat(listener.activeChangesOnDefaultDisplay).isEqualTo(1)
+    }
+
+    @Test
+    fun addActiveTask_taskIsActive() {
+        val listener = TestListener()
+        repo.addActiveTaskListener(listener)
+
+        repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
+
         assertThat(repo.isActiveTask(1)).isTrue()
     }
 
     @Test
-    fun addActiveTask_sameTaskDoesNotNotify() {
+    fun addSameActiveTaskTwice_notifiesOnce() {
         val listener = TestListener()
         repo.addActiveTaskListener(listener)
 
         repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
         repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
+
         assertThat(listener.activeChangesOnDefaultDisplay).isEqualTo(1)
     }
 
     @Test
-    fun addActiveTask_multipleTasksAddedNotifiesForEach() {
+    fun addActiveTask_multipleTasksAdded_notifiesForAllTasks() {
         val listener = TestListener()
         repo.addActiveTaskListener(listener)
 
         repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
         repo.addActiveTask(DEFAULT_DISPLAY, taskId = 2)
+
         assertThat(listener.activeChangesOnDefaultDisplay).isEqualTo(2)
     }
 
@@ -84,22 +96,35 @@
     }
 
     @Test
-    fun removeActiveTask_listenerNotifiedAndTaskNotActive() {
+    fun removeActiveTask_notifiesListener() {
         val listener = TestListener()
         repo.addActiveTaskListener(listener)
-
         repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
+
         repo.removeActiveTask(1)
+
         // Notify once for add and once for remove
         assertThat(listener.activeChangesOnDefaultDisplay).isEqualTo(2)
+    }
+
+    @Test
+    fun removeActiveTask_taskNotActive() {
+        val listener = TestListener()
+        repo.addActiveTaskListener(listener)
+        repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
+
+        repo.removeActiveTask(1)
+
         assertThat(repo.isActiveTask(1)).isFalse()
     }
 
     @Test
-    fun removeActiveTask_removeNotExistingTaskDoesNotNotify() {
+    fun removeActiveTask_nonExistingTask_doesNotNotify() {
         val listener = TestListener()
         repo.addActiveTaskListener(listener)
+
         repo.removeActiveTask(99)
+
         assertThat(listener.activeChangesOnDefaultDisplay).isEqualTo(0)
     }
 
@@ -108,32 +133,38 @@
         val listener = TestListener()
         repo.addActiveTaskListener(listener)
         repo.addActiveTask(DEFAULT_DISPLAY, taskId = 1)
+
         repo.removeActiveTask(1)
+
         assertThat(listener.activeChangesOnSecondaryDisplay).isEqualTo(0)
         assertThat(repo.isActiveTask(1)).isFalse()
     }
 
     @Test
-    fun isActiveTask_notExistingTaskReturnsFalse() {
+    fun isActiveTask_nonExistingTask_returnsFalse() {
         assertThat(repo.isActiveTask(99)).isFalse()
     }
 
     @Test
-    fun isOnlyVisibleNonClosingTask_noTasks() {
+    fun isOnlyVisibleNonClosingTask_noTasks_returnsFalse() {
         // No visible tasks
         assertThat(repo.isOnlyVisibleNonClosingTask(1)).isFalse()
+    }
+
+    @Test
+    fun isClosingTask_noTasks_returnsFalse() {
+        // No visible tasks
         assertThat(repo.isClosingTask(1)).isFalse()
     }
 
     @Test
-    fun isOnlyVisibleNonClosingTask_singleVisibleNonClosingTask() {
+    fun updateVisibleFreeformTasks_singleVisibleNonClosingTask_updatesTasksCorrectly() {
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
 
-        // The only visible task
         assertThat(repo.isVisibleTask(1)).isTrue()
         assertThat(repo.isClosingTask(1)).isFalse()
         assertThat(repo.isOnlyVisibleNonClosingTask(1)).isTrue()
-        // Not a visible task
+
         assertThat(repo.isVisibleTask(99)).isFalse()
         assertThat(repo.isClosingTask(99)).isFalse()
         assertThat(repo.isOnlyVisibleNonClosingTask(99)).isFalse()
@@ -207,10 +238,11 @@
     }
 
     @Test
-    fun addListener_notifiesVisibleFreeformTask() {
+    fun addVisibleTasksListener_notifiesVisibleFreeformTask() {
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
         val listener = TestVisibilityListener()
         val executor = TestShellExecutor()
+
         repo.addVisibleTasksListener(listener, executor)
         executor.flushAll()
 
@@ -236,6 +268,7 @@
         val listener = TestVisibilityListener()
         val executor = TestShellExecutor()
         repo.addVisibleTasksListener(listener, executor)
+
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 2, visible = true)
         executor.flushAll()
@@ -303,6 +336,7 @@
         executor.flushAll()
 
         assertThat(listener.visibleTasksCountOnDefaultDisplay).isEqualTo(2)
+
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = false)
         executor.flushAll()
 
@@ -329,6 +363,7 @@
         executor.flushAll()
 
         assertThat(listener.visibleTasksCountOnDefaultDisplay).isEqualTo(2)
+
         repo.updateVisibleFreeformTasks(INVALID_DISPLAY, taskId = 1, visible = false)
         executor.flushAll()
 
@@ -337,65 +372,73 @@
     }
 
     @Test
-    fun visibleTaskCount_defaultDisplay_returnsCorrectCount() {
+    fun getVisibleTaskCount_defaultDisplay_returnsCorrectCount() {
         // No tasks, count is 0
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
 
         // New task increments count to 1
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
 
         // Visibility update to same task does not increase count
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
 
         // Second task visible increments count
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 2, visible = true)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(2)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(2)
 
         // Hiding a task decrements count
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = false)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
 
         // Hiding all tasks leaves count at 0
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 2, visible = false)
-        assertThat(repo.visibleTaskCount(displayId = 9)).isEqualTo(0)
+        assertThat(repo.getVisibleTaskCount(displayId = 9)).isEqualTo(0)
 
         // Hiding a not existing task, count remains at 0
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 999, visible = false)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
     }
 
     @Test
-    fun visibleTaskCount_multipleDisplays_returnsCorrectCount() {
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
-        assertThat(repo.visibleTaskCount(SECOND_DISPLAY)).isEqualTo(0)
+    fun getVisibleTaskCount_multipleDisplays_returnsCorrectCount() {
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
+        assertThat(repo.getVisibleTaskCount(SECOND_DISPLAY)).isEqualTo(0)
 
         // New task on default display increments count for that display only
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
-        assertThat(repo.visibleTaskCount(SECOND_DISPLAY)).isEqualTo(0)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
+        assertThat(repo.getVisibleTaskCount(SECOND_DISPLAY)).isEqualTo(0)
 
         // New task on secondary display, increments count for that display only
         repo.updateVisibleFreeformTasks(SECOND_DISPLAY, taskId = 2, visible = true)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
-        assertThat(repo.visibleTaskCount(SECOND_DISPLAY)).isEqualTo(1)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
+        assertThat(repo.getVisibleTaskCount(SECOND_DISPLAY)).isEqualTo(1)
 
         // Marking task visible on another display, updates counts for both displays
         repo.updateVisibleFreeformTasks(SECOND_DISPLAY, taskId = 1, visible = true)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
-        assertThat(repo.visibleTaskCount(SECOND_DISPLAY)).isEqualTo(2)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
+        assertThat(repo.getVisibleTaskCount(SECOND_DISPLAY)).isEqualTo(2)
 
         // Marking task that is on secondary display, hidden on default display, does not affect
         // secondary display
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = false)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
-        assertThat(repo.visibleTaskCount(SECOND_DISPLAY)).isEqualTo(2)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
+        assertThat(repo.getVisibleTaskCount(SECOND_DISPLAY)).isEqualTo(2)
 
         // Hiding a task on that display, decrements count
         repo.updateVisibleFreeformTasks(SECOND_DISPLAY, taskId = 1, visible = false)
-        assertThat(repo.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
-        assertThat(repo.visibleTaskCount(SECOND_DISPLAY)).isEqualTo(1)
+
+        assertThat(repo.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
+        assertThat(repo.getVisibleTaskCount(SECOND_DISPLAY)).isEqualTo(1)
     }
 
     @Test
@@ -428,7 +471,9 @@
     fun removeFreeformTask_removesTaskBoundsBeforeMaximize() {
         val taskId = 1
         repo.saveBoundsBeforeMaximize(taskId, Rect(0, 0, 200, 200))
+
         repo.removeFreeformTask(THIRD_DISPLAY, taskId)
+
         assertThat(repo.removeBoundsBeforeMaximize(taskId)).isNull()
     }
 
@@ -436,7 +481,9 @@
     fun saveBoundsBeforeMaximize_boundsSavedByTaskId() {
         val taskId = 1
         val bounds = Rect(0, 0, 200, 200)
+
         repo.saveBoundsBeforeMaximize(taskId, bounds)
+
         assertThat(repo.removeBoundsBeforeMaximize(taskId)).isEqualTo(bounds)
     }
 
@@ -446,17 +493,20 @@
         val bounds = Rect(0, 0, 200, 200)
         repo.saveBoundsBeforeMaximize(taskId, bounds)
         repo.removeBoundsBeforeMaximize(taskId)
-        assertThat(repo.removeBoundsBeforeMaximize(taskId)).isNull()
+
+        val boundsBeforeMaximize = repo.removeBoundsBeforeMaximize(taskId)
+
+        assertThat(boundsBeforeMaximize).isNull()
     }
 
     @Test
-    fun minimizeTaskNotCalled_noTasksMinimized() {
+    fun isMinimizedTask_minimizeTaskNotCalled_noTasksMinimized() {
         assertThat(repo.isMinimizedTask(taskId = 0)).isFalse()
         assertThat(repo.isMinimizedTask(taskId = 1)).isFalse()
     }
 
     @Test
-    fun minimizeTask_onlyThatTaskIsMinimized() {
+    fun minimizeTask_minimizesCorrectTask() {
         repo.minimizeTask(displayId = 0, taskId = 0)
 
         assertThat(repo.isMinimizedTask(taskId = 0)).isTrue()
@@ -465,8 +515,9 @@
     }
 
     @Test
-    fun unminimizeTask_taskNoLongerMinimized() {
+    fun unminimizeTask_unminimizesTask() {
         repo.minimizeTask(displayId = 0, taskId = 0)
+
         repo.unminimizeTask(displayId = 0, taskId = 0)
 
         assertThat(repo.isMinimizedTask(taskId = 0)).isFalse()
@@ -478,6 +529,7 @@
     fun unminimizeTask_nonExistentTask_doesntCrash() {
         repo.unminimizeTask(displayId = 0, taskId = 0)
 
+        // No change
         assertThat(repo.isMinimizedTask(taskId = 0)).isFalse()
         assertThat(repo.isMinimizedTask(taskId = 1)).isFalse()
         assertThat(repo.isMinimizedTask(taskId = 2)).isFalse()
@@ -485,41 +537,44 @@
 
 
     @Test
-    fun updateVisibleFreeformTasks_toVisible_taskIsUnminimized() {
+    fun updateVisibleFreeformTasks_minimizedTaskBecomesVisible_unminimizesTask() {
         repo.minimizeTask(displayId = 10, taskId = 2)
-
         repo.updateVisibleFreeformTasks(displayId = 10, taskId = 2, visible = true)
 
-        assertThat(repo.isMinimizedTask(taskId = 2)).isFalse()
+        val isMinimizedTask = repo.isMinimizedTask(taskId = 2)
+
+        assertThat(isMinimizedTask).isFalse()
     }
 
     @Test
-    fun getActiveNonMinimizedTasksOrderedFrontToBack_returnsFreeformTasksInCorrectOrder() {
+    fun getActiveNonMinimizedOrderedTasks_returnsFreeformTasksInCorrectOrder() {
         repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 1)
         repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 2)
         repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 3)
-        // The front-most task will be the one added last through addOrMoveFreeformTaskToTop
+        // The front-most task will be the one added last through `addOrMoveFreeformTaskToTop`
         repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 3)
         repo.addOrMoveFreeformTaskToTop(displayId = 0, taskId = 2)
         repo.addOrMoveFreeformTaskToTop(displayId = 0, taskId = 1)
 
-        assertThat(repo.getActiveNonMinimizedTasksOrderedFrontToBack(displayId = 0))
-            .containsExactly(1, 2, 3).inOrder()
+        val tasks = repo.getActiveNonMinimizedOrderedTasks(displayId = 0)
+
+        assertThat(tasks).containsExactly(1, 2, 3).inOrder()
     }
 
     @Test
-    fun getActiveNonMinimizedTasksOrderedFrontToBack_minimizedTaskNotIncluded() {
+    fun getActiveNonMinimizedOrderedTasks_excludesMinimizedTasks() {
         repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 1)
         repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 2)
         repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 3)
-        // The front-most task will be the one added last through addOrMoveFreeformTaskToTop
+        // The front-most task will be the one added last through `addOrMoveFreeformTaskToTop`
         repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 3)
         repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 2)
         repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 1)
         repo.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = 2)
 
-        assertThat(repo.getActiveNonMinimizedTasksOrderedFrontToBack(
-            displayId = DEFAULT_DISPLAY)).containsExactly(1, 3).inOrder()
+        val tasks = repo.getActiveNonMinimizedOrderedTasks(displayId = DEFAULT_DISPLAY)
+
+        assertThat(tasks).containsExactly(1, 3).inOrder()
     }
 
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index e66018f..2368cef 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -43,6 +43,7 @@
 import android.platform.test.flag.junit.SetFlagsRule
 import android.testing.AndroidTestingRunner
 import android.view.Display.DEFAULT_DISPLAY
+import android.view.Gravity
 import android.view.SurfaceControl
 import android.view.WindowManager
 import android.view.WindowManager.TRANSIT_CHANGE
@@ -70,6 +71,7 @@
 import com.android.window.flags.Flags
 import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE
 import com.android.wm.shell.MockToken
+import com.android.wm.shell.R
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.ShellTestCase
@@ -185,12 +187,12 @@
 
   private val DISPLAY_DIMENSION_SHORT = 1600
   private val DISPLAY_DIMENSION_LONG = 2560
-  private val DEFAULT_LANDSCAPE_BOUNDS = Rect(320, 200, 2240, 1400)
-  private val DEFAULT_PORTRAIT_BOUNDS = Rect(200, 320, 1400, 2240)
-  private val RESIZABLE_LANDSCAPE_BOUNDS = Rect(25, 680, 1575, 1880)
-  private val RESIZABLE_PORTRAIT_BOUNDS = Rect(680, 200, 1880, 1400)
-  private val UNRESIZABLE_LANDSCAPE_BOUNDS = Rect(25, 699, 1575, 1861)
-  private val UNRESIZABLE_PORTRAIT_BOUNDS = Rect(830, 200, 1730, 1400)
+  private val DEFAULT_LANDSCAPE_BOUNDS = Rect(320, 75, 2240, 1275)
+  private val DEFAULT_PORTRAIT_BOUNDS = Rect(200, 165, 1400, 2085)
+  private val RESIZABLE_LANDSCAPE_BOUNDS = Rect(25, 435, 1575, 1635)
+  private val RESIZABLE_PORTRAIT_BOUNDS = Rect(680, 75, 1880, 1275)
+  private val UNRESIZABLE_LANDSCAPE_BOUNDS = Rect(25, 449, 1575, 1611)
+  private val UNRESIZABLE_PORTRAIT_BOUNDS = Rect(830, 75, 1730, 1275)
 
   @Before
   fun setUp() {
@@ -590,6 +592,161 @@
   }
 
   @Test
+  fun addMoveToDesktopChanges_gravityLeft_noBoundsApplied() {
+    setUpLandscapeDisplay()
+    val task = setUpFullscreenTask(gravity = Gravity.LEFT)
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(finalBounds).isEqualTo(Rect())
+  }
+
+  @Test
+  fun addMoveToDesktopChanges_gravityRight_noBoundsApplied() {
+    setUpLandscapeDisplay()
+    val task = setUpFullscreenTask(gravity = Gravity.RIGHT)
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(finalBounds).isEqualTo(Rect())
+  }
+
+  @Test
+  fun addMoveToDesktopChanges_gravityTop_noBoundsApplied() {
+    setUpLandscapeDisplay()
+    val task = setUpFullscreenTask(gravity = Gravity.TOP)
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(finalBounds).isEqualTo(Rect())
+  }
+
+  @Test
+  fun addMoveToDesktopChanges_gravityBottom_noBoundsApplied() {
+    setUpLandscapeDisplay()
+    val task = setUpFullscreenTask(gravity = Gravity.BOTTOM)
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(finalBounds).isEqualTo(Rect())
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_CASCADING_WINDOWS)
+  fun addMoveToDesktopChanges_positionBottomRight() {
+    setUpLandscapeDisplay()
+    val stableBounds = Rect()
+    displayLayout.getStableBoundsForDesktopMode(stableBounds)
+
+    setUpFreeformTask(bounds = DEFAULT_LANDSCAPE_BOUNDS)
+
+    val task = setUpFullscreenTask()
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(stableBounds.getDesktopTaskPosition(finalBounds!!))
+      .isEqualTo(DesktopTaskPosition.BottomRight)
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_CASCADING_WINDOWS)
+  fun addMoveToDesktopChanges_positionTopLeft() {
+    setUpLandscapeDisplay()
+    val stableBounds = Rect()
+    displayLayout.getStableBoundsForDesktopMode(stableBounds)
+
+    addFreeformTaskAtPosition(DesktopTaskPosition.BottomRight, stableBounds)
+
+    val task = setUpFullscreenTask()
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(stableBounds.getDesktopTaskPosition(finalBounds!!))
+      .isEqualTo(DesktopTaskPosition.TopLeft)
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_CASCADING_WINDOWS)
+  fun addMoveToDesktopChanges_positionBottomLeft() {
+    setUpLandscapeDisplay()
+    val stableBounds = Rect()
+    displayLayout.getStableBoundsForDesktopMode(stableBounds)
+
+    addFreeformTaskAtPosition(DesktopTaskPosition.TopLeft, stableBounds)
+
+    val task = setUpFullscreenTask()
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(stableBounds.getDesktopTaskPosition(finalBounds!!))
+      .isEqualTo(DesktopTaskPosition.BottomLeft)
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_CASCADING_WINDOWS)
+  fun addMoveToDesktopChanges_positionTopRight() {
+    setUpLandscapeDisplay()
+    val stableBounds = Rect()
+    displayLayout.getStableBoundsForDesktopMode(stableBounds)
+
+    addFreeformTaskAtPosition(DesktopTaskPosition.BottomLeft, stableBounds)
+
+    val task = setUpFullscreenTask()
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(stableBounds.getDesktopTaskPosition(finalBounds!!))
+      .isEqualTo(DesktopTaskPosition.TopRight)
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_CASCADING_WINDOWS)
+  fun addMoveToDesktopChanges_positionResetsToCenter() {
+    setUpLandscapeDisplay()
+    val stableBounds = Rect()
+    displayLayout.getStableBoundsForDesktopMode(stableBounds)
+
+    addFreeformTaskAtPosition(DesktopTaskPosition.TopRight, stableBounds)
+
+    val task = setUpFullscreenTask()
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(stableBounds.getDesktopTaskPosition(finalBounds!!))
+      .isEqualTo(DesktopTaskPosition.Center)
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_CASCADING_WINDOWS)
+  fun addMoveToDesktopChanges_defaultToCenterIfFree() {
+    setUpLandscapeDisplay()
+    val stableBounds = Rect()
+    displayLayout.getStableBoundsForDesktopMode(stableBounds)
+
+    val minTouchTarget = context.resources.getDimensionPixelSize(
+      R.dimen.freeform_required_visible_empty_space_in_header)
+    addFreeformTaskAtPosition(DesktopTaskPosition.Center, stableBounds,
+      Rect(0, 0, 1600, 1200), Point(0, minTouchTarget + 1))
+
+    val task = setUpFullscreenTask()
+    val wct = WindowContainerTransaction()
+    controller.addMoveToDesktopChanges(wct, task)
+
+    val finalBounds = findBoundsChange(wct, task)
+    assertThat(stableBounds.getDesktopTaskPosition(finalBounds!!))
+      .isEqualTo(DesktopTaskPosition.Center)
+  }
+
+  @Test
   fun moveToDesktop_tdaFullscreen_windowingModeSetToFreeform() {
     val task = setUpFullscreenTask()
     val tda = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY)!!
@@ -2406,6 +2563,18 @@
   private val desktopWallpaperIntent: Intent
     get() = Intent(context, DesktopWallpaperActivity::class.java)
 
+  private fun addFreeformTaskAtPosition(
+    pos: DesktopTaskPosition,
+    stableBounds: Rect,
+    bounds: Rect = DEFAULT_LANDSCAPE_BOUNDS,
+    offsetPos: Point = Point(0, 0)
+  ): RunningTaskInfo {
+    val offset = pos.getTopLeftCoordinates(stableBounds, bounds)
+    val prevTaskBounds = Rect(bounds)
+    prevTaskBounds.offsetTo(offset.x + offsetPos.x, offset.y + offsetPos.y)
+    return setUpFreeformTask(bounds = prevTaskBounds)
+  }
+
   private fun setUpFreeformTask(
       displayId: Int = DEFAULT_DISPLAY,
       bounds: Rect? = null
@@ -2434,11 +2603,13 @@
       windowingMode: Int = WINDOWING_MODE_FULLSCREEN,
       deviceOrientation: Int = ORIENTATION_LANDSCAPE,
       screenOrientation: Int = SCREEN_ORIENTATION_UNSPECIFIED,
-      shouldLetterbox: Boolean = false
+      shouldLetterbox: Boolean = false,
+      gravity: Int = Gravity.NO_GRAVITY
   ): RunningTaskInfo {
     val task = createFullscreenTask(displayId)
     val activityInfo = ActivityInfo()
     activityInfo.screenOrientation = screenOrientation
+    activityInfo.windowLayout = ActivityInfo.WindowLayout(0, 0F, 0, 0F, gravity, 0, 0)
     with(task) {
       topActivityInfo = activityInfo
       isResizeable = isResizable
@@ -2479,11 +2650,23 @@
   private fun setUpLandscapeDisplay() {
     whenever(displayLayout.width()).thenReturn(DISPLAY_DIMENSION_LONG)
     whenever(displayLayout.height()).thenReturn(DISPLAY_DIMENSION_SHORT)
+    val stableBounds = Rect(0, 0, DISPLAY_DIMENSION_LONG,
+      DISPLAY_DIMENSION_SHORT - Companion.TASKBAR_FRAME_HEIGHT
+    )
+    whenever(displayLayout.getStableBoundsForDesktopMode(any())).thenAnswer { i ->
+      (i.arguments.first() as Rect).set(stableBounds)
+    }
   }
 
   private fun setUpPortraitDisplay() {
     whenever(displayLayout.width()).thenReturn(DISPLAY_DIMENSION_SHORT)
     whenever(displayLayout.height()).thenReturn(DISPLAY_DIMENSION_LONG)
+    val stableBounds = Rect(0, 0, DISPLAY_DIMENSION_SHORT,
+      DISPLAY_DIMENSION_LONG - Companion.TASKBAR_FRAME_HEIGHT
+    )
+    whenever(displayLayout.getStableBoundsForDesktopMode(any())).thenAnswer { i ->
+      (i.arguments.first() as Rect).set(stableBounds)
+    }
   }
 
   private fun setUpSplitScreenTask(displayId: Int = DEFAULT_DISPLAY): RunningTaskInfo {
@@ -2601,6 +2784,7 @@
     const val SECOND_DISPLAY = 2
     val STABLE_BOUNDS = Rect(0, 0, 1000, 1000)
     const val MAX_TASK_LIMIT = 6
+    private const val TASKBAR_FRAME_HEIGHT = 200
   }
 }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
index 409b877..81e6d07 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
@@ -442,6 +442,27 @@
     }
 
     @Test
+    public void testTransitionFilterAnimOverride() {
+        TransitionFilter filter = new TransitionFilter();
+        filter.mRequirements =
+                new TransitionFilter.Requirement[]{new TransitionFilter.Requirement()};
+        filter.mRequirements[0].mCustomAnimation = true;
+        filter.mRequirements[0].mModes = new int[]{TRANSIT_OPEN, TRANSIT_TO_FRONT};
+
+        final RunningTaskInfo taskInf = createTaskInfo(1);
+        final TransitionInfo openTask = new TransitionInfoBuilder(TRANSIT_OPEN)
+                .addChange(TRANSIT_OPEN, taskInf).build();
+        assertFalse(filter.matches(openTask));
+
+        final TransitionInfo.AnimationOptions overOpts =
+                TransitionInfo.AnimationOptions.makeCustomAnimOptions("pakname", 0, 0, 0, true);
+        final TransitionInfo openTaskOpts = new TransitionInfoBuilder(TRANSIT_OPEN)
+                .addChange(TRANSIT_OPEN, taskInf).build();
+        openTaskOpts.getChanges().get(0).setAnimationOptions(overOpts);
+        assertTrue(filter.matches(openTaskOpts));
+    }
+
+    @Test
     public void testRegisteredRemoteTransition() {
         Transitions transitions = createTestTransitions();
         transitions.replaceDefaultHandlerForTest(mDefaultHandler);
diff --git a/media/java/android/media/session/MediaSession.java b/media/java/android/media/session/MediaSession.java
index 1d6e38d..8a877b8 100644
--- a/media/java/android/media/session/MediaSession.java
+++ b/media/java/android/media/session/MediaSession.java
@@ -165,7 +165,7 @@
 
     /**
      * Creates a new session. The session will automatically be registered with
-     * the system but will not be published until {@link #setActive(boolean)
+     * the system, but will not be published until {@link #setActive(boolean)
      * setActive(true)} is called. You must call {@link #release()} when
      * finished with the session.
      * <p>
diff --git a/mms/java/android/telephony/MmsManager.java b/mms/java/android/telephony/MmsManager.java
index b893b45..ac29277 100644
--- a/mms/java/android/telephony/MmsManager.java
+++ b/mms/java/android/telephony/MmsManager.java
@@ -26,6 +26,7 @@
 import android.os.Bundle;
 import android.os.RemoteException;
 import android.os.ServiceManager;
+import android.os.UserHandle;
 
 import com.android.internal.telephony.IMms;
 
@@ -69,9 +70,9 @@
                 return;
             }
 
-            iMms.sendMessage(subId, ActivityThread.currentPackageName(), contentUri,
-                    locationUrl, configOverrides, sentIntent, messageId,
-                    mContext.getAttributionTag());
+            iMms.sendMessage(subId, /* placeholder callingUser= */ UserHandle.USER_NULL,
+                    ActivityThread.currentPackageName(), contentUri, locationUrl,
+                    configOverrides, sentIntent, messageId, mContext.getAttributionTag());
         } catch (RemoteException e) {
             // Ignore it
         }
@@ -101,9 +102,9 @@
             if (iMms == null) {
                 return;
             }
-            iMms.downloadMessage(subId, ActivityThread.currentPackageName(),
-                    locationUrl, contentUri, configOverrides, downloadedIntent,
-                    messageId, mContext.getAttributionTag());
+            iMms.downloadMessage(subId, /* placeholder callingUser= */ UserHandle.USER_NULL,
+                    ActivityThread.currentPackageName(), locationUrl, contentUri,
+                    configOverrides, downloadedIntent, messageId, mContext.getAttributionTag());
         } catch (RemoteException e) {
             // Ignore it
         }
diff --git a/mms/java/com/android/internal/telephony/IMms.aidl b/mms/java/com/android/internal/telephony/IMms.aidl
index 3cdde10..1c75951 100644
--- a/mms/java/com/android/internal/telephony/IMms.aidl
+++ b/mms/java/com/android/internal/telephony/IMms.aidl
@@ -29,6 +29,7 @@
      * Send an MMS message with attribution tag.
      *
      * @param subId the SIM id
+     * @param callingUser user id of the calling app
      * @param callingPkg the package name of the calling app
      * @param contentUri the content uri from which to read MMS message encoded in standard MMS
      *  PDU format
@@ -40,7 +41,7 @@
      * @param messageId An id that uniquely identifies the message requested to be sent.
      * @param attributionTag a tag that attributes the call to a client App.
      */
-    void sendMessage(int subId, String callingPkg, in Uri contentUri,
+    void sendMessage(int subId, in int callingUser, String callingPkg, in Uri contentUri,
             String locationUrl, in Bundle configOverrides, in PendingIntent sentIntent,
             in long messageId, String attributionTag);
 
@@ -48,6 +49,7 @@
      * Download an MMS message using known location and transaction id
      *
      * @param subId the SIM id
+     * @param callingUser user id of the calling app
      * @param callingPkg the package name of the calling app
      * @param locationUrl the location URL of the MMS message to be downloaded, usually obtained
      *  from the MMS WAP push notification
@@ -60,7 +62,7 @@
      * @param messageId An id that uniquely identifies the message requested to be downloaded.
      * @param attributionTag a tag that attributes the call to a client App.
     */
-    void downloadMessage(int subId, String callingPkg, String locationUrl,
+    void downloadMessage(int subId, in int callingUser, String callingPkg, String locationUrl,
             in Uri contentUri, in Bundle configOverrides,
             in PendingIntent downloadedIntent, in long messageId, String attributionTag);
 
@@ -82,6 +84,7 @@
     /**
       * Import a multimedia message into system's MMS store
       *
+     * @param callingUser user id of the calling app
       * @param callingPkg the package name of the calling app
       * @param contentUri the content uri from which to read PDU of the message to import
       * @param messageId the optional message id
@@ -90,7 +93,7 @@
       * @param read if the message is read
       * @return the message URI, null if failed
       */
-    Uri importMultimediaMessage(String callingPkg, in Uri contentUri, String messageId,
+    Uri importMultimediaMessage(in int callingUser, String callingPkg, in Uri contentUri, String messageId,
             long timestampSecs, boolean seen, boolean read);
 
     /**
@@ -146,11 +149,12 @@
     /**
      * Add a multimedia message draft to system MMS store
      *
+     * @param callingUser user id of the calling app
      * @param callingPkg the package name of the calling app
      * @param contentUri the content Uri from which to read PDU data of the draft MMS
      * @return the URI of the stored draft message
      */
-    Uri addMultimediaMessageDraft(String callingPkg, in Uri contentUri);
+    Uri addMultimediaMessageDraft(in int callingUser, String callingPkg, in Uri contentUri);
 
     /**
      * Send a system stored MMS message
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
index b4a9172..ce02404 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
@@ -129,7 +129,7 @@
             }
         }
 
-        override fun onChanged(reason: Int) = onKeyChanged(null, reason)
+        override fun onChanged(observable: Observable, reason: Int) = onKeyChanged(null, reason)
 
         override fun onKeyChanged(key: Any?, reason: Int) {
             notifyBackupManager(key, reason)
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt
index ede7c63..4ce1d37 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/KeyedObserver.kt
@@ -103,7 +103,7 @@
 }
 
 /** A thread safe implementation of [KeyedObservable]. */
-class KeyedDataObservable<K> : KeyedObservable<K> {
+open class KeyedDataObservable<K> : KeyedObservable<K> {
     // Instead of @GuardedBy("this"), guarded by itself because KeyedDataObservable object could be
     // synchronized outside by the holder
     @GuardedBy("itself") private val observers = WeakHashMap<KeyedObserver<K?>, Executor>()
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/ObservableBackupRestoreStorage.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/ObservableBackupRestoreStorage.kt
index 0e399c0..300d240 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/ObservableBackupRestoreStorage.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/ObservableBackupRestoreStorage.kt
@@ -21,5 +21,7 @@
  *
  * This class provides the [Observable] implementations on top of [DataObservable] by delegation.
  */
-abstract class ObservableBackupRestoreStorage :
-    BackupRestoreStorage(), Observable by DataObservable()
+abstract class ObservableBackupRestoreStorage : BackupRestoreStorage(), ObservableDelegation {
+
+    final override val observableDelegate: Observable = DataObservable(this)
+}
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/Observer.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/Observer.kt
index 98d0f6e..6af3d1c 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/Observer.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/Observer.kt
@@ -32,10 +32,11 @@
      *
      * This callback will run in the given [Executor] when observer is added.
      *
+     * @param observable observable of the change
      * @param reason the reason of change
      * @see [Observable.addObserver] for the notices.
      */
-    fun onChanged(reason: Int)
+    fun onChanged(observable: Observable, reason: Int)
 }
 
 /** An observable object allows to observe change with [Observer]. */
@@ -68,8 +69,21 @@
     fun notifyChange(reason: Int)
 }
 
+/** Delegation of [Observable]. */
+interface ObservableDelegation : Observable {
+    /** [Observable] to delegate. */
+    val observableDelegate: Observable
+
+    override fun addObserver(observer: Observer, executor: Executor) =
+        observableDelegate.addObserver(observer, executor)
+
+    override fun removeObserver(observer: Observer) = observableDelegate.removeObserver(observer)
+
+    override fun notifyChange(reason: Int) = observableDelegate.notifyChange(reason)
+}
+
 /** A thread safe implementation of [Observable]. */
-class DataObservable : Observable {
+class DataObservable(private val observable: Observable) : Observable {
     // Instead of @GuardedBy("this"), guarded by itself because DataObservable object could be
     // synchronized outside by the holder
     @GuardedBy("itself") private val observers = WeakHashMap<Observer, Executor>()
@@ -90,7 +104,7 @@
         val entries = synchronized(observers) { observers.entries.toTypedArray() }
         for (entry in entries) {
             val observer = entry.key // avoid reference "entry"
-            entry.value.execute { observer.onChanged(reason) }
+            entry.value.execute { observer.onChanged(observable, reason) }
         }
     }
 }
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesObservable.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesObservable.kt
new file mode 100644
index 0000000..e70ec5b
--- /dev/null
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesObservable.kt
@@ -0,0 +1,45 @@
+/*
+ * Copyright (C) 2024 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.settingslib.datastore
+
+import android.content.SharedPreferences
+
+/** [SharedPreferences] based [KeyedDataObservable]. */
+class SharedPreferencesObservable(private val sharedPreferences: SharedPreferences) :
+    KeyedDataObservable<String>(), AutoCloseable {
+
+    private val listener = createSharedPreferenceListener()
+
+    init {
+        sharedPreferences.registerOnSharedPreferenceChangeListener(listener)
+    }
+
+    override fun close() {
+        sharedPreferences.unregisterOnSharedPreferenceChangeListener(listener)
+    }
+}
+
+/** Creates [SharedPreferences.OnSharedPreferenceChangeListener] for [KeyedObservable]. */
+internal fun KeyedObservable<String>.createSharedPreferenceListener() =
+    SharedPreferences.OnSharedPreferenceChangeListener { _, key ->
+        if (key != null) {
+            notifyChange(key, DataChangeReason.UPDATE)
+        } else {
+            // On Android >= R, SharedPreferences.Editor.clear() will trigger this case
+            notifyChange(DataChangeReason.DELETE)
+        }
+    }
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
index 20a95d7..0ca91cd 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
@@ -80,15 +80,7 @@
             return context.getSharedPreferences(intermediateName, Context.MODE_MULTI_PROCESS)
         }
 
-    private val sharedPreferencesListener =
-        SharedPreferences.OnSharedPreferenceChangeListener { _, key ->
-            if (key != null) {
-                notifyChange(key, DataChangeReason.UPDATE)
-            } else {
-                // On Android >= R, SharedPreferences.Editor.clear() will trigger this case
-                notifyChange(DataChangeReason.DELETE)
-            }
-        }
+    private val sharedPreferencesListener = createSharedPreferenceListener()
 
     init {
         // listener is weakly referenced, so unregister is optional
@@ -191,8 +183,7 @@
                 else -> {
                     Log.e(
                         LOG_TAG,
-                        "[$name] $operation $key=$value, unknown type: ${value?.javaClass}"
-                    )
+                        "[$name] $operation $key=$value, unknown type: ${value?.javaClass}")
                 }
             }
         }
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt
index 19c574a..97b473c 100644
--- a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt
@@ -159,7 +159,7 @@
 
         verify(keyedObserver).onKeyChanged("key", DataChangeReason.RESTORE)
         verify(anyKeyObserver).onKeyChanged(null, DataChangeReason.RESTORE)
-        verify(observer).onChanged(DataChangeReason.RESTORE)
+        verify(observer).onChanged(fileStorage, DataChangeReason.RESTORE)
         if (isRobolectric()) {
             Shadows.shadowOf(BackupManager(application)).apply {
                 assertThat(isDataChanged).isFalse()
@@ -187,7 +187,7 @@
         }
 
         fileStorage.notifyChange(DataChangeReason.UPDATE)
-        verify(observer).onChanged(DataChangeReason.UPDATE)
+        verify(observer).onChanged(fileStorage, DataChangeReason.UPDATE)
         verify(keyedObserver, never()).onKeyChanged(any(), any())
         verify(anyKeyObserver, never()).onKeyChanged(any(), any())
         reset(observer)
@@ -197,7 +197,7 @@
         }
 
         keyedStorage.notifyChange("key", DataChangeReason.DELETE)
-        verify(observer, never()).onChanged(any())
+        verify(observer, never()).onChanged(any(), any())
         verify(keyedObserver).onKeyChanged("key", DataChangeReason.DELETE)
         verify(anyKeyObserver).onKeyChanged("key", DataChangeReason.DELETE)
         backupManager?.apply {
@@ -209,7 +209,7 @@
         // backup manager is not notified for restore event
         fileStorage.notifyChange(DataChangeReason.RESTORE)
         keyedStorage.notifyChange("key", DataChangeReason.RESTORE)
-        verify(observer).onChanged(DataChangeReason.RESTORE)
+        verify(observer).onChanged(fileStorage, DataChangeReason.RESTORE)
         verify(keyedObserver).onKeyChanged("key", DataChangeReason.RESTORE)
         verify(anyKeyObserver).onKeyChanged("key", DataChangeReason.RESTORE)
         backupManager?.apply {
@@ -225,7 +225,10 @@
     }
 
     private class FileStorage(override val name: String) :
-        BackupRestoreFileStorage(getApplicationContext(), "file"), Observable by DataObservable()
+        BackupRestoreFileStorage(getApplicationContext(), "file"), ObservableDelegation {
+
+        override val observableDelegate: Observable = DataObservable(this)
+    }
 
     private class DummyBackupAgentHelper : BackupAgentHelper() {
         val backupHelpers = mutableMapOf<String, BackupHelper>()
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
index 5d0303c..bd114d1 100644
--- a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
@@ -24,6 +24,7 @@
 import org.junit.Assert.assertThrows
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.any
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
 import org.mockito.kotlin.reset
@@ -33,10 +34,11 @@
 class ObserverTest {
     private val observer1 = mock<Observer>()
     private val observer2 = mock<Observer>()
+    private val originalObservable = mock<Observable>()
 
     private val executor1: Executor = MoreExecutors.directExecutor()
     private val executor2: Executor = MoreExecutors.newDirectExecutorService()
-    private val observable = DataObservable()
+    private val observable = DataObservable(originalObservable)
 
     @Test
     fun addObserver_sameExecutor() {
@@ -55,7 +57,7 @@
     @Test
     fun addObserver_weaklyReferenced() {
         val counter = AtomicInteger()
-        var observer: Observer? = Observer { counter.incrementAndGet() }
+        var observer: Observer? = Observer { _, _ -> counter.incrementAndGet() }
         observable.addObserver(observer!!, executor1)
 
         observable.notifyChange(DataChangeReason.UPDATE)
@@ -77,21 +79,21 @@
 
         observable.notifyChange(DataChangeReason.DELETE)
 
-        verify(observer1).onChanged(DataChangeReason.DELETE)
-        verify(observer2).onChanged(DataChangeReason.DELETE)
+        verify(observer1).onChanged(originalObservable, DataChangeReason.DELETE)
+        verify(observer2).onChanged(originalObservable, DataChangeReason.DELETE)
 
         reset(observer1, observer2)
         observable.removeObserver(observer2)
 
         observable.notifyChange(DataChangeReason.UPDATE)
-        verify(observer1).onChanged(DataChangeReason.UPDATE)
-        verify(observer2, never()).onChanged(DataChangeReason.UPDATE)
+        verify(observer1).onChanged(originalObservable, DataChangeReason.UPDATE)
+        verify(observer2, never()).onChanged(any(), any())
     }
 
     @Test
     fun notifyChange_addObserverWithinCallback() {
         // ConcurrentModificationException is raised if it is not implemented correctly
-        val observer = Observer { observable.addObserver(observer1, executor1) }
+        val observer = Observer { _, _ -> observable.addObserver(observer1, executor1) }
         observable.addObserver(observer, executor1)
         observable.notifyChange(DataChangeReason.UPDATE)
         observable.removeObserver(observer)
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt b/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt
index 49b974f..87ab6b3 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt
@@ -73,15 +73,22 @@
     }
 
     fun activateMode(id: String) {
-        val oldMode = mutableModesFlow.value.find { it.id == id } ?: return
-        removeMode(id)
-        mutableModesFlow.value += TestModeBuilder(oldMode).setActive(true).build()
+        updateModeActiveState(id = id, isActive = true)
     }
 
     fun deactivateMode(id: String) {
-        val oldMode = mutableModesFlow.value.find { it.id == id } ?: return
-        removeMode(id)
-        mutableModesFlow.value += TestModeBuilder(oldMode).setActive(false).build()
+        updateModeActiveState(id = id, isActive = false)
+    }
+
+    // Update the active state while maintaining the mode's position in the list
+    private fun updateModeActiveState(id: String, isActive: Boolean) {
+        val modes = mutableModesFlow.value.toMutableList()
+        val index = modes.indexOfFirst { it.id == id }
+        if (index < 0) {
+            throw IllegalArgumentException("mode $id not found")
+        }
+        modes[index] = TestModeBuilder(modes[index]).setActive(isActive).build()
+        mutableModesFlow.value = modes
     }
 }
 
@@ -101,7 +108,8 @@
             suppressedVisualEffects,
             state,
             priorityConversationSenders,
-        ))
+        )
+    )
 
 private fun newMode(id: String, active: Boolean = false): ZenMode {
     return TestModeBuilder().setId(id).setName("Mode $id").setActive(active).build()
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java
index 2f7cdd6..a06f084 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java
@@ -16,6 +16,9 @@
 
 package com.android.settingslib.notification.modes;
 
+import static android.service.notification.ZenModeConfig.UPDATE_ORIGIN_UNKNOWN;
+import static android.service.notification.ZenModeConfig.UPDATE_ORIGIN_USER;
+
 import android.app.AutomaticZenRule;
 import android.app.NotificationManager;
 import android.content.ComponentName;
@@ -144,8 +147,15 @@
     }
 
     public TestModeBuilder setEnabled(boolean enabled) {
+        return setEnabled(enabled, /* byUser= */ false);
+    }
+
+    public TestModeBuilder setEnabled(boolean enabled, boolean byUser) {
         mRule.setEnabled(enabled);
         mConfigZenRule.enabled = enabled;
+        if (!enabled) {
+            mConfigZenRule.disabledOrigin = byUser ? UPDATE_ORIGIN_USER : UPDATE_ORIGIN_UNKNOWN;
+        }
         return this;
     }
 
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
index 960df63..271d5c4 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
@@ -35,7 +35,6 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
-import androidx.appcompat.content.res.AppCompatResources;
 
 import com.google.common.util.concurrent.FluentFuture;
 import com.google.common.util.concurrent.Futures;
@@ -104,7 +103,7 @@
                 return context.getDrawable(iconResId);
             } else {
                 Context appContext = context.createPackageContext(pkg, 0);
-                Drawable appDrawable = AppCompatResources.getDrawable(appContext, iconResId);
+                Drawable appDrawable = appContext.getDrawable(iconResId);
                 return getMonochromeIconIfPresent(appDrawable);
             }
         })).catching(Exception.class, ex -> {
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
index 990a2d4..732b358 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
@@ -16,6 +16,8 @@
 
 package com.android.settingslib.notification.modes;
 
+import static android.app.AutomaticZenRule.TYPE_SCHEDULE_CALENDAR;
+import static android.app.AutomaticZenRule.TYPE_SCHEDULE_TIME;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_ALL;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
 import static android.service.notification.SystemZenRules.getTriggerDescriptionForScheduleEvent;
@@ -298,6 +300,17 @@
         return mIsManualDnd;
     }
 
+    /**
+     * A <em>custom manual</em> mode is a mode created by the user, and not yet assigned an
+     * automatic trigger condition (neither time schedule nor a calendar).
+     */
+    public boolean isCustomManual() {
+        return isSystemOwned()
+                && getType() != TYPE_SCHEDULE_TIME
+                && getType() != TYPE_SCHEDULE_CALENDAR
+                && !isManualDnd();
+    }
+
     public boolean isEnabled() {
         return mRule.isEnabled();
     }
diff --git a/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioSharingRepository.kt b/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioSharingRepository.kt
index e78b8a7..99d5891 100644
--- a/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioSharingRepository.kt
+++ b/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioSharingRepository.kt
@@ -62,6 +62,9 @@
     /** Whether the device is in audio sharing. */
     val inAudioSharing: Flow<Boolean>
 
+    /** The primary headset groupId in audio sharing. */
+    val primaryGroupId: StateFlow<Int>
+
     /** The secondary headset groupId in audio sharing. */
     val secondaryGroupId: StateFlow<Int>
 
@@ -109,6 +112,16 @@
         awaitClose { contentResolver.unregisterContentObserver(callback) }
     }
 
+    override val primaryGroupId: StateFlow<Int> =
+        primaryChange
+            .map { BluetoothUtils.getPrimaryGroupIdForBroadcast(contentResolver) }
+            .onStart { emit(BluetoothUtils.getPrimaryGroupIdForBroadcast(contentResolver)) }
+            .flowOn(backgroundCoroutineContext)
+            .stateIn(
+                coroutineScope,
+                SharingStarted.WhileSubscribed(),
+                BluetoothUtils.getPrimaryGroupIdForBroadcast(contentResolver))
+
     override val secondaryGroupId: StateFlow<Int> =
         merge(
                 btManager.profileManager.leAudioBroadcastAssistantProfile
@@ -121,7 +134,7 @@
                                 BluetoothProfile.LE_AUDIO_BROADCAST_ASSISTANT
                     }
                     .map { getSecondaryGroupId() },
-                primaryChange.map { getSecondaryGroupId() })
+                primaryGroupId.map { getSecondaryGroupId() })
             .onStart { emit(getSecondaryGroupId()) }
             .flowOn(backgroundCoroutineContext)
             .stateIn(coroutineScope, SharingStarted.WhileSubscribed(), getSecondaryGroupId())
@@ -193,6 +206,8 @@
 
 class AudioSharingRepositoryEmptyImpl : AudioSharingRepository {
     override val inAudioSharing: Flow<Boolean> = flowOf(false)
+    override val primaryGroupId: StateFlow<Int> =
+        MutableStateFlow(BluetoothCsipSetCoordinator.GROUP_ID_INVALID)
     override val secondaryGroupId: StateFlow<Int> =
         MutableStateFlow(BluetoothCsipSetCoordinator.GROUP_ID_INVALID)
     override val volumeMap: StateFlow<GroupIdToVolumes> = MutableStateFlow(emptyMap())
diff --git a/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioSharingRepositoryTest.kt b/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioSharingRepositoryTest.kt
index 94595d3..c54a2e4 100644
--- a/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioSharingRepositoryTest.kt
+++ b/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioSharingRepositoryTest.kt
@@ -131,6 +131,10 @@
         `when`(deviceManager.findDevice(device2)).thenReturn(cachedDevice2)
         `when`(receiveState.bisSyncState).thenReturn(arrayListOf(TEST_RECEIVE_STATE_CONTENT))
         `when`(assistant.getAllSources(any())).thenReturn(listOf(receiveState))
+        Settings.Secure.putInt(
+            contentResolver,
+            BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
+            TEST_GROUP_ID_INVALID)
         underTest =
             AudioSharingRepositoryImpl(
                 contentResolver,
@@ -156,6 +160,22 @@
     }
 
     @Test
+    fun primaryGroupIdChange_emitValues() {
+        testScope.runTest {
+            val groupIds = mutableListOf<Int?>()
+            underTest.primaryGroupId.onEach { groupIds.add(it) }.launchIn(backgroundScope)
+            runCurrent()
+            triggerContentObserverChange()
+            runCurrent()
+
+            Truth.assertThat(groupIds)
+                .containsExactly(
+                    TEST_GROUP_ID_INVALID,
+                    TEST_GROUP_ID2)
+        }
+    }
+
+    @Test
     fun secondaryGroupIdChange_emitValues() {
         testScope.runTest {
             val groupIds = mutableListOf<Int?>()
@@ -217,7 +237,7 @@
     fun setSecondaryVolume_setValue() {
         testScope.runTest {
             Settings.Secure.putInt(
-                context.contentResolver,
+                contentResolver,
                 BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
                 TEST_GROUP_ID2)
             `when`(assistant.allConnectedDevices).thenReturn(listOf(device1, device2))
@@ -248,7 +268,7 @@
     private fun triggerSourceAdded() {
         verify(assistant).registerServiceCallBack(any(), assistantCallbackCaptor.capture())
         Settings.Secure.putInt(
-            context.contentResolver,
+            contentResolver,
             BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
             TEST_GROUP_ID1)
         `when`(assistant.allConnectedDevices).thenReturn(listOf(device1, device2))
@@ -259,7 +279,7 @@
         verify(assistant).registerServiceCallBack(any(), assistantCallbackCaptor.capture())
         `when`(assistant.allConnectedDevices).thenReturn(listOf(device1))
         Settings.Secure.putInt(
-            context.contentResolver,
+            contentResolver,
             BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
             TEST_GROUP_ID1)
         assistantCallbackCaptor.value.sourceRemoved(device2)
@@ -269,7 +289,7 @@
         verify(eventManager).registerCallback(btCallbackCaptor.capture())
         `when`(assistant.allConnectedDevices).thenReturn(listOf(device1))
         Settings.Secure.putInt(
-            context.contentResolver,
+            contentResolver,
             BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
             TEST_GROUP_ID1)
         btCallbackCaptor.value.onProfileConnectionStateChanged(cachedDevice2, state, profile)
@@ -283,7 +303,7 @@
                 contentObserverCaptor.capture())
         `when`(assistant.allConnectedDevices).thenReturn(listOf(device1, device2))
         Settings.Secure.putInt(
-            context.contentResolver,
+            contentResolver,
             BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
             TEST_GROUP_ID2)
         contentObserverCaptor.value.primaryChanged()
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
index e705f97..651e57c 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
@@ -27,6 +27,7 @@
 import android.net.Uri;
 import android.os.Parcel;
 import android.service.notification.Condition;
+import android.service.notification.SystemZenRules;
 import android.service.notification.ZenModeConfig;
 import android.service.notification.ZenPolicy;
 
@@ -109,6 +110,61 @@
     }
 
     @Test
+    public void isCustomManual_customManualMode() {
+        AutomaticZenRule rule = new AutomaticZenRule.Builder("Mode", Uri.parse("x"))
+                .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                .setType(AutomaticZenRule.TYPE_OTHER)
+                .build();
+        ZenMode mode = new ZenMode("id", rule, zenConfigRuleFor(rule, false));
+
+        assertThat(mode.isCustomManual()).isTrue();
+    }
+
+    @Test
+    public void isCustomManual_scheduleTime_false() {
+        AutomaticZenRule rule = new AutomaticZenRule.Builder("Mode", Uri.parse("x"))
+                .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                .setType(AutomaticZenRule.TYPE_SCHEDULE_TIME)
+                .build();
+        ZenMode mode = new ZenMode("id", rule, zenConfigRuleFor(rule, false));
+
+        assertThat(mode.isCustomManual()).isFalse();
+    }
+
+    @Test
+    public void isCustomManual_scheduleCalendar_false() {
+        AutomaticZenRule rule = new AutomaticZenRule.Builder("Mode", Uri.parse("x"))
+                .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                .setType(AutomaticZenRule.TYPE_SCHEDULE_CALENDAR)
+                .build();
+        ZenMode mode = new ZenMode("id", rule, zenConfigRuleFor(rule, false));
+
+        assertThat(mode.isCustomManual()).isFalse();
+    }
+
+    @Test
+    public void isCustomManual_appProvidedMode_false() {
+        AutomaticZenRule rule = new AutomaticZenRule.Builder("Mode", Uri.parse("x"))
+                .setPackage("com.some.package")
+                .setType(AutomaticZenRule.TYPE_OTHER)
+                .build();
+        ZenMode mode = new ZenMode("id", rule, zenConfigRuleFor(rule, false));
+
+        assertThat(mode.isCustomManual()).isFalse();
+    }
+
+    @Test
+    public void isCustomManual_manualDnd_false() {
+        AutomaticZenRule dndRule = new AutomaticZenRule.Builder("Mode", Uri.parse("x"))
+                .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                .setType(AutomaticZenRule.TYPE_OTHER)
+                .build();
+        ZenMode mode = ZenMode.manualDndMode(dndRule, false);
+
+        assertThat(mode.isCustomManual()).isFalse();
+    }
+
+    @Test
     public void getPolicy_interruptionFilterPriority_returnsZenPolicy() {
         AutomaticZenRule azr = new AutomaticZenRule.Builder("Rule", Uri.EMPTY)
                 .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
diff --git a/packages/SettingsProvider/src/android/provider/settings/backup/SecureSettings.java b/packages/SettingsProvider/src/android/provider/settings/backup/SecureSettings.java
index 2b8b23e..40a8199 100644
--- a/packages/SettingsProvider/src/android/provider/settings/backup/SecureSettings.java
+++ b/packages/SettingsProvider/src/android/provider/settings/backup/SecureSettings.java
@@ -253,6 +253,7 @@
         Settings.Secure.CUSTOM_BUGREPORT_HANDLER_APP,
         Settings.Secure.CUSTOM_BUGREPORT_HANDLER_USER,
         Settings.Secure.CONTEXTUAL_SCREEN_TIMEOUT_ENABLED,
+        Settings.Secure.HINGE_ANGLE_LIDEVENT_ENABLED,
         Settings.Secure.LOCK_SCREEN_WEATHER_ENABLED,
         Settings.Secure.HEARING_AID_RINGTONE_ROUTING,
         Settings.Secure.HEARING_AID_CALL_ROUTING,
diff --git a/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java b/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java
index cc5302b..3b9c683 100644
--- a/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java
+++ b/packages/SettingsProvider/src/android/provider/settings/validators/SecureSettingsValidators.java
@@ -406,6 +406,7 @@
         VALIDATORS.put(Secure.CUSTOM_BUGREPORT_HANDLER_USER, ANY_INTEGER_VALIDATOR);
         VALIDATORS.put(Secure.LOCK_SCREEN_WEATHER_ENABLED, BOOLEAN_VALIDATOR);
         VALIDATORS.put(Secure.CONTEXTUAL_SCREEN_TIMEOUT_ENABLED, BOOLEAN_VALIDATOR);
+        VALIDATORS.put(Secure.HINGE_ANGLE_LIDEVENT_ENABLED, BOOLEAN_VALIDATOR);
         VALIDATORS.put(Secure.HEARING_AID_RINGTONE_ROUTING,
                 new DiscreteValueValidator(new String[] {"0", "1", "2"}));
         VALIDATORS.put(Secure.HEARING_AID_CALL_ROUTING,
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 2f90ccc..cfd8f635 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -258,6 +258,7 @@
         "tests/src/**/systemui/statusbar/policy/LocationControllerImplTest.java",
         "tests/src/**/systemui/statusbar/policy/RemoteInputViewTest.java",
         "tests/src/**/systemui/statusbar/policy/SmartReplyViewTest.java",
+        "tests/src/**/systemui/statusbar/policy/ui/dialog/ModesDialogDelegateTest.kt",
         "tests/src/**/systemui/statusbar/StatusBarStateControllerImplTest.kt",
         "tests/src/**/systemui/theme/ThemeOverlayApplierTest.java",
         "tests/src/**/systemui/touch/TouchInsetManagerTest.java",
diff --git a/packages/SystemUI/AndroidManifest.xml b/packages/SystemUI/AndroidManifest.xml
index 666d939..92abc4c 100644
--- a/packages/SystemUI/AndroidManifest.xml
+++ b/packages/SystemUI/AndroidManifest.xml
@@ -482,7 +482,7 @@
             android:exported="true"
             android:theme="@style/Theme.AppCompat.NoActionBar">
             <intent-filter>
-                <action android:name="com.android.systemui.action.TOUCHPAD_TUTORIAL"/>
+                <action android:name="com.android.systemui.action.TOUCHPAD_KEYBOARD_TUTORIAL"/>
                 <category android:name="android.intent.category.DEFAULT"/>
             </intent-filter>
         </activity>
diff --git a/packages/SystemUI/accessibility/accessibilitymenu/aconfig/accessibility.aconfig b/packages/SystemUI/accessibility/accessibilitymenu/aconfig/accessibility.aconfig
index ba84287..c1e43c9 100644
--- a/packages/SystemUI/accessibility/accessibilitymenu/aconfig/accessibility.aconfig
+++ b/packages/SystemUI/accessibility/accessibilitymenu/aconfig/accessibility.aconfig
@@ -19,3 +19,13 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "action_bar_wrap_content"
+    namespace: "accessibility"
+    description: "Applies WRAP_CONTENT to the action bar in A11yMenu settings to better fit large fonts"
+    bug: "347911378"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/activity/A11yMenuSettingsActivity.java b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/activity/A11yMenuSettingsActivity.java
index ab8f97a..c71ef83 100644
--- a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/activity/A11yMenuSettingsActivity.java
+++ b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/activity/A11yMenuSettingsActivity.java
@@ -26,6 +26,7 @@
 import android.provider.Browser;
 import android.provider.Settings;
 import android.view.View;
+import android.view.ViewGroup;
 import android.widget.TextView;
 import android.window.OnBackInvokedCallback;
 
@@ -35,6 +36,7 @@
 import androidx.preference.PreferenceFragmentCompat;
 import androidx.preference.PreferenceManager;
 
+import com.android.systemui.accessibility.accessibilitymenu.Flags;
 import com.android.systemui.accessibility.accessibilitymenu.R;
 
 /**
@@ -60,6 +62,18 @@
         ((TextView) findViewById(R.id.action_bar_title)).setText(
                 getResources().getString(R.string.accessibility_menu_settings_name)
         );
+        if (Flags.actionBarWrapContent()) {
+            setHeightWrapContent(findViewById(com.android.internal.R.id.action_bar));
+            setHeightWrapContent(findViewById(com.android.internal.R.id.action_bar_container));
+        }
+    }
+
+    private void setHeightWrapContent(View view) {
+        if (view != null) {
+            ViewGroup.LayoutParams params = view.getLayoutParams();
+            params.height = ViewGroup.LayoutParams.WRAP_CONTENT;
+            view.setLayoutParams(params);
+        }
     }
 
     @Override
diff --git a/packages/SystemUI/aconfig/accessibility.aconfig b/packages/SystemUI/aconfig/accessibility.aconfig
index 0861454..8860452 100644
--- a/packages/SystemUI/aconfig/accessibility.aconfig
+++ b/packages/SystemUI/aconfig/accessibility.aconfig
@@ -90,6 +90,16 @@
 }
 
 flag {
+    name: "update_corner_radius_on_display_changed"
+    namespace: "accessibility"
+    description: "Updates the corner radius to the magnification fullscreen border when the display changes."
+    bug: "335113174"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "hearing_devices_dialog_related_tools"
     namespace: "accessibility"
     description: "Shows the related tools for hearing devices dialog."
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 462db34..1f1495a 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -601,6 +601,13 @@
 }
 
 flag {
+    name: "screenshot_ui_controller_refactor"
+    namespace: "systemui"
+    description: "Simplify and refactor ScreenshotController"
+    bug: "354711957"
+}
+
+flag {
    name: "run_fingerprint_detect_on_dismissible_keyguard"
    namespace: "systemui"
    description: "Run fingerprint detect instead of authenticate if the keyguard is dismissible."
@@ -998,6 +1005,16 @@
 }
 
 flag {
+  name: "communal_timer_flicker_fix"
+  namespace: "systemui"
+  description: "fixes timers on the hub flickering when pausing"
+  bug: "353801573"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
+
+flag {
   name: "app_clips_backlinks"
   namespace: "systemui"
   description: "Enables Backlinks improvement feature in App Clips"
@@ -1210,6 +1227,9 @@
   namespace: "systemui"
   description: "Enables fullscreen vertical swiping in hub mode to bring up and down the bouncer and shade"
   bug: "340177049"
+  metadata {
+      purpose: PURPOSE_BUGFIX
+  }
 }
 
 flag {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
index 35db9e0..3be5231 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
@@ -7,6 +7,7 @@
 import androidx.compose.animation.core.rememberInfiniteTransition
 import androidx.compose.animation.core.tween
 import androidx.compose.foundation.background
+import androidx.compose.foundation.focusable
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.isSystemInDarkTheme
 import androidx.compose.foundation.layout.Box
@@ -25,6 +26,9 @@
 import androidx.compose.ui.graphics.Brush
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.platform.LocalDensity
+import androidx.compose.ui.semantics.clearAndSetSemantics
+import androidx.compose.ui.semantics.contentDescription
+import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.Back
@@ -41,6 +45,7 @@
 import com.android.compose.animation.scene.observableTransitionState
 import com.android.compose.animation.scene.transitions
 import com.android.compose.theme.LocalAndroidColorScheme
+import com.android.internal.R.attr.focusable
 import com.android.systemui.Flags.glanceableHubBackGesture
 import com.android.systemui.communal.shared.model.CommunalBackgroundType
 import com.android.systemui.communal.shared.model.CommunalScenes
@@ -88,12 +93,12 @@
     }
     to(CommunalScenes.Communal) {
         spec = tween(durationMillis = 1000)
-        translate(Communal.Elements.Grid, Edge.Right)
+        translate(Communal.Elements.Grid, Edge.End)
         timestampRange(startMillis = 167, endMillis = 334) { fade(AllElements) }
     }
     to(CommunalScenes.Blank) {
         spec = tween(durationMillis = 1000)
-        translate(Communal.Elements.Grid, Edge.Right)
+        translate(Communal.Elements.Grid, Edge.End)
         timestampRange(endMillis = 167) {
             fade(Communal.Elements.Grid)
             fade(Communal.Elements.IndicationArea)
@@ -186,9 +191,7 @@
         scene(
             CommunalScenes.Blank,
             userActions =
-                mapOf(
-                    Swipe(SwipeDirection.Left, fromSource = Edge.Right) to CommunalScenes.Communal
-                )
+                mapOf(Swipe(SwipeDirection.Start, fromSource = Edge.End) to CommunalScenes.Communal)
         ) {
             // This scene shows nothing only allowing for transitions to the communal scene.
             Box(modifier = Modifier.fillMaxSize())
@@ -197,11 +200,11 @@
         val userActions =
             if (glanceableHubBackGesture()) {
                 mapOf(
-                    Swipe(SwipeDirection.Right) to CommunalScenes.Blank,
+                    Swipe(SwipeDirection.End) to CommunalScenes.Blank,
                     Back to CommunalScenes.Blank,
                 )
             } else {
-                mapOf(Swipe(SwipeDirection.Right) to CommunalScenes.Blank)
+                mapOf(Swipe(SwipeDirection.End) to CommunalScenes.Blank)
             }
 
         scene(CommunalScenes.Communal, userActions = userActions) {
@@ -209,6 +212,8 @@
                 backgroundType = backgroundType,
                 colors = colors,
                 content = content,
+                viewModel = viewModel,
+                modifier = Modifier.horizontalNestedScrollToScene(),
             )
         }
     }
@@ -224,17 +229,41 @@
     backgroundType: CommunalBackgroundType,
     colors: CommunalColors,
     content: CommunalContent,
+    viewModel: CommunalViewModel,
     modifier: Modifier = Modifier,
 ) {
-    Box(modifier = Modifier.element(Communal.Elements.Scrim).fillMaxSize()) {
+    val isFocusable by viewModel.isFocusable.collectAsStateWithLifecycle(initialValue = false)
+
+    Box(
+        modifier =
+            Modifier.element(Communal.Elements.Scrim)
+                .fillMaxSize()
+                .then(
+                    if (isFocusable) {
+                        Modifier.focusable()
+                    } else {
+                        Modifier.semantics { contentDescription = "" }.clearAndSetSemantics {}
+                    }
+                )
+    ) {
         when (backgroundType) {
             CommunalBackgroundType.STATIC -> DefaultBackground(colors = colors)
             CommunalBackgroundType.STATIC_GRADIENT -> StaticLinearGradient()
             CommunalBackgroundType.ANIMATED -> AnimatedLinearGradient()
             CommunalBackgroundType.NONE -> BackgroundTopScrim()
         }
+
+        with(content) {
+            Content(
+                modifier =
+                    modifier.focusable(isFocusable).semantics {
+                        if (!isFocusable) {
+                            contentDescription = ""
+                        }
+                    }
+            )
+        }
     }
-    with(content) { Content(modifier = modifier) }
 }
 
 /** Default background of the hub, a single color */
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
index fc95754..4c29e79 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
@@ -52,6 +52,7 @@
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.RowScope
 import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.calculateStartPadding
 import androidx.compose.foundation.layout.fillMaxHeight
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
@@ -68,8 +69,10 @@
 import androidx.compose.foundation.lazy.grid.LazyGridState
 import androidx.compose.foundation.lazy.grid.LazyHorizontalGrid
 import androidx.compose.foundation.lazy.grid.rememberLazyGridState
+import androidx.compose.foundation.rememberScrollState
 import androidx.compose.foundation.selection.selectable
 import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.foundation.verticalScroll
 import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.filled.Add
 import androidx.compose.material.icons.filled.Check
@@ -91,6 +94,7 @@
 import androidx.compose.material3.Text
 import androidx.compose.material3.rememberModalBottomSheetState
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.State
 import androidx.compose.runtime.derivedStateOf
@@ -127,6 +131,7 @@
 import androidx.compose.ui.platform.LocalConfiguration
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalDensity
+import androidx.compose.ui.platform.LocalLayoutDirection
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.res.dimensionResource
 import androidx.compose.ui.res.stringResource
@@ -143,9 +148,11 @@
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
 import androidx.compose.ui.unit.times
 import androidx.compose.ui.util.fastAll
 import androidx.compose.ui.viewinterop.AndroidView
+import androidx.compose.ui.viewinterop.NoOpUpdate
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import androidx.window.layout.WindowMetricsCalculator
 import com.android.compose.animation.Easings.Emphasized
@@ -153,6 +160,7 @@
 import com.android.compose.theme.LocalAndroidColorScheme
 import com.android.compose.ui.graphics.painter.rememberDrawablePainter
 import com.android.internal.R.dimen.system_app_widget_background_radius
+import com.android.systemui.Flags.communalTimerFlickerFix
 import com.android.systemui.communal.domain.model.CommunalContentModel
 import com.android.systemui.communal.shared.model.CommunalContentSize
 import com.android.systemui.communal.shared.model.CommunalScenes
@@ -209,7 +217,14 @@
 
     ObserveScrollEffect(gridState, viewModel)
 
-    if (!viewModel.isEditMode) {
+    val context = LocalContext.current
+    val windowMetrics = WindowMetricsCalculator.getOrCreate().computeCurrentWindowMetrics(context)
+    val screenWidth = windowMetrics.bounds.width()
+    val layoutDirection = LocalLayoutDirection.current
+
+    if (viewModel.isEditMode) {
+        ObserveNewWidgetAddedEffect(communalContent, gridState, viewModel)
+    } else {
         ScrollOnUpdatedLiveContentEffect(communalContent, gridState)
     }
 
@@ -230,7 +245,7 @@
                 .testTag(COMMUNAL_HUB_TEST_TAG)
                 .fillMaxSize()
                 .nestedScroll(nestedScrollConnection)
-                .pointerInput(gridState, contentOffset, contentListState) {
+                .pointerInput(layoutDirection, gridState, contentOffset, contentListState) {
                     awaitPointerEventScope {
                         while (true) {
                             var event = awaitFirstDown(requireUnconsumed = false)
@@ -261,7 +276,13 @@
                     // If not in edit mode, don't allow selecting items.
                     if (!viewModel.isEditMode) return@pointerInput
                     observeTaps { offset ->
-                        val adjustedOffset = offset - contentOffset
+                        // if RTL, flip offset direction from Left side to Right
+                        val adjustedOffset =
+                            Offset(
+                                if (layoutDirection == LayoutDirection.Rtl) screenWidth - offset.x
+                                else offset.x,
+                                offset.y
+                            ) - contentOffset
                         val index = firstIndexAtOffset(gridState, adjustedOffset)
                         val key = index?.let { keyAtIndexIfEditable(contentListState.list, index) }
                         viewModel.setSelectedKey(key)
@@ -279,7 +300,12 @@
                             // offset.
                             val adjustedOffset =
                                 gridCoordinates?.let {
-                                    offset - it.positionInWindow() - contentOffset
+                                    Offset(
+                                        if (layoutDirection == LayoutDirection.Rtl)
+                                            screenWidth - offset.x
+                                        else offset.x,
+                                        offset.y
+                                    ) - it.positionInWindow() - contentOffset
                                 }
                             val index = adjustedOffset?.let { firstIndexAtOffset(gridState, it) }
                             val key = index?.let { keyAtIndexIfEditable(communalContent, index) }
@@ -330,6 +356,7 @@
                             viewModel = viewModel,
                             contentPadding = contentPadding,
                             contentOffset = contentOffset,
+                            screenWidth = screenWidth,
                             setGridCoordinates = { gridCoordinates = it },
                             updateDragPositionForRemove = { offset ->
                                 isPointerWithinEnabledRemoveButton(
@@ -528,6 +555,56 @@
     }
 }
 
+/**
+ * Observes communal content and determines whether a new widget has been added, upon which case:
+ * - Announce for accessibility
+ * - Scroll if the new widget is not visible
+ */
+@Composable
+private fun ObserveNewWidgetAddedEffect(
+    communalContent: List<CommunalContentModel>,
+    gridState: LazyGridState,
+    viewModel: BaseCommunalViewModel,
+) {
+    val coroutineScope = rememberCoroutineScope()
+    val widgetKeys = remember { mutableListOf<String>() }
+    var communalContentPending by remember { mutableStateOf(true) }
+
+    LaunchedEffect(communalContent) {
+        // Do nothing until any communal content comes in
+        if (communalContentPending && communalContent.isEmpty()) {
+            return@LaunchedEffect
+        }
+
+        val oldWidgetKeys = widgetKeys.toList()
+        val widgets = communalContent.filterIsInstance<CommunalContentModel.WidgetContent.Widget>()
+        widgetKeys.clear()
+        widgetKeys.addAll(widgets.map { it.key })
+
+        // Do nothing on first communal content since we don't have a delta
+        if (communalContentPending) {
+            communalContentPending = false
+            return@LaunchedEffect
+        }
+
+        // Do nothing if there is no new widget
+        val indexOfFirstNewWidget = widgetKeys.indexOfFirst { !oldWidgetKeys.contains(it) }
+        if (indexOfFirstNewWidget < 0) {
+            return@LaunchedEffect
+        }
+
+        viewModel.onNewWidgetAdded(widgets[indexOfFirstNewWidget].providerInfo)
+
+        // Scroll if the new widget is not visible
+        val lastVisibleItemIndex = gridState.layoutInfo.visibleItemsInfo.lastOrNull()?.index
+        if (lastVisibleItemIndex != null && indexOfFirstNewWidget > lastVisibleItemIndex) {
+            // Launching with a scope to prevent the job from being canceled in the case of a
+            // recomposition during scrolling
+            coroutineScope.launch { gridState.animateScrollToItem(indexOfFirstNewWidget) }
+        }
+    }
+}
+
 @OptIn(ExperimentalFoundationApi::class)
 @Composable
 private fun BoxScope.CommunalHubLazyGrid(
@@ -535,6 +612,7 @@
     viewModel: BaseCommunalViewModel,
     contentPadding: PaddingValues,
     selectedKey: State<String?>,
+    screenWidth: Int,
     contentOffset: Offset,
     gridState: LazyGridState,
     contentListState: ContentListState,
@@ -557,7 +635,15 @@
                 updateDragPositionForRemove = updateDragPositionForRemove
             )
         gridModifier =
-            gridModifier.fillMaxSize().dragContainer(dragDropState, contentOffset, viewModel)
+            gridModifier
+                .fillMaxSize()
+                .dragContainer(
+                    dragDropState,
+                    LocalLayoutDirection.current,
+                    screenWidth,
+                    contentOffset,
+                    viewModel
+                )
         // for widgets dropped from other activities
         val dragAndDropTargetState =
             rememberDragAndDropTargetState(
@@ -940,7 +1026,9 @@
         shape = RoundedCornerShape(68.adjustedDp, 34.adjustedDp, 68.adjustedDp, 34.adjustedDp)
     ) {
         Column(
-            modifier = Modifier.fillMaxSize().padding(vertical = 32.dp, horizontal = 50.dp),
+            modifier =
+                Modifier.fillMaxSize()
+                    .padding(vertical = 32.adjustedDp, horizontal = 50.adjustedDp),
             verticalArrangement = Arrangement.Center,
             horizontalAlignment = Alignment.CenterHorizontally,
         ) {
@@ -949,47 +1037,57 @@
                 contentDescription = stringResource(R.string.cta_label_to_open_widget_picker),
                 modifier = Modifier.size(Dimensions.IconSize).clearAndSetSemantics {},
             )
-            Spacer(modifier = Modifier.size(6.dp))
+            Spacer(modifier = Modifier.size(6.adjustedDp))
             Text(
                 text = stringResource(R.string.cta_label_to_edit_widget),
                 style = MaterialTheme.typography.titleLarge,
                 fontSize = nonScalableTextSize(22.dp),
                 lineHeight = nonScalableTextSize(28.dp),
+                modifier = Modifier.verticalScroll(rememberScrollState()).weight(1F)
             )
-            Spacer(modifier = Modifier.size(16.dp))
+            Spacer(modifier = Modifier.size(16.adjustedDp))
             Row(
-                modifier = Modifier.fillMaxWidth().height(56.dp),
-                horizontalArrangement = Arrangement.spacedBy(16.dp, Alignment.CenterHorizontally),
+                modifier = Modifier.fillMaxWidth().height(56.adjustedDp),
+                horizontalArrangement =
+                    Arrangement.spacedBy(16.adjustedDp, Alignment.CenterHorizontally),
             ) {
-                OutlinedButton(
-                    modifier = Modifier.fillMaxHeight(),
-                    colors =
-                        ButtonDefaults.buttonColors(
-                            contentColor = colors.onPrimary,
-                        ),
-                    border = BorderStroke(width = 1.0.dp, color = colors.primaryContainer),
-                    contentPadding = PaddingValues(26.dp, 8.dp),
-                    onClick = viewModel::onDismissCtaTile,
+                CompositionLocalProvider(
+                    LocalDensity provides
+                        Density(
+                            LocalDensity.current.density,
+                            LocalDensity.current.fontScale.coerceIn(0f, 1.25f)
+                        )
                 ) {
-                    Text(
-                        text = stringResource(R.string.cta_tile_button_to_dismiss),
-                        fontSize = nonScalableTextSize(14.dp),
-                    )
-                }
-                Button(
-                    modifier = Modifier.fillMaxHeight(),
-                    colors =
-                        ButtonDefaults.buttonColors(
-                            containerColor = colors.primaryContainer,
-                            contentColor = colors.onPrimaryContainer,
-                        ),
-                    contentPadding = PaddingValues(26.dp, 8.dp),
-                    onClick = viewModel::onOpenWidgetEditor
-                ) {
-                    Text(
-                        text = stringResource(R.string.cta_tile_button_to_open_widget_editor),
-                        fontSize = nonScalableTextSize(14.dp),
-                    )
+                    OutlinedButton(
+                        modifier = Modifier.fillMaxHeight().weight(1F),
+                        colors =
+                            ButtonDefaults.buttonColors(
+                                contentColor = colors.onPrimary,
+                            ),
+                        border = BorderStroke(width = 1.0.dp, color = colors.primaryContainer),
+                        onClick = viewModel::onDismissCtaTile,
+                        contentPadding = PaddingValues(0.dp, 0.dp, 0.dp, 0.dp),
+                    ) {
+                        Text(
+                            text = stringResource(R.string.cta_tile_button_to_dismiss),
+                            fontSize = 14.sp,
+                        )
+                    }
+                    Button(
+                        modifier = Modifier.fillMaxHeight().weight(1F),
+                        colors =
+                            ButtonDefaults.buttonColors(
+                                containerColor = colors.primaryContainer,
+                                contentColor = colors.onPrimaryContainer,
+                            ),
+                        onClick = viewModel::onOpenWidgetEditor,
+                        contentPadding = PaddingValues(0.dp, 0.dp, 0.dp, 0.dp),
+                    ) {
+                        Text(
+                            text = stringResource(R.string.cta_tile_button_to_open_widget_editor),
+                            fontSize = 14.sp,
+                        )
+                    }
                 }
             }
         }
@@ -1260,9 +1358,15 @@
         factory = { context ->
             SmartspaceAppWidgetHostView(context).apply {
                 interactionHandler?.let { setInteractionHandler(it) }
-                updateAppWidget(model.remoteViews)
+                if (!communalTimerFlickerFix()) {
+                    updateAppWidget(model.remoteViews)
+                }
             }
         },
+        update =
+            if (communalTimerFlickerFix()) {
+                { view: SmartspaceAppWidgetHostView -> view.updateAppWidget(model.remoteViews) }
+            } else NoOpUpdate,
         // For reusing composition in lazy lists.
         onReset = {},
     )
@@ -1371,7 +1475,7 @@
 private fun beforeContentPadding(paddingValues: PaddingValues): ContentPaddingInPx {
     return with(LocalDensity.current) {
         ContentPaddingInPx(
-            start = paddingValues.calculateLeftPadding(LayoutDirection.Ltr).toPx(),
+            start = paddingValues.calculateStartPadding(LocalLayoutDirection.current).toPx(),
             top = paddingValues.calculateTopPadding().toPx()
         )
     }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt
index e0fc340..5886d7d 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt
@@ -48,7 +48,7 @@
     override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
         MutableStateFlow<Map<UserAction, UserActionResult>>(
                 mapOf(
-                    Swipe(SwipeDirection.Right) to UserActionResult(Scenes.Lockscreen),
+                    Swipe(SwipeDirection.End) to UserActionResult(Scenes.Lockscreen),
                 )
             )
             .asStateFlow()
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
index 07898b0..20ee131 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/GridDragDropState.kt
@@ -37,7 +37,9 @@
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.graphics.graphicsLayer
 import androidx.compose.ui.input.pointer.pointerInput
+import androidx.compose.ui.platform.LocalLayoutDirection
 import androidx.compose.ui.unit.IntOffset
+import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.toOffset
 import androidx.compose.ui.unit.toSize
 import com.android.systemui.communal.ui.compose.extensions.firstItemAtOffset
@@ -47,6 +49,9 @@
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.launch
 
+private fun Float.directional(origin: LayoutDirection, current: LayoutDirection): Float =
+    if (origin == current) this else -this
+
 @Composable
 fun rememberGridDragDropState(
     gridState: LazyGridState,
@@ -113,14 +118,24 @@
      *
      * @return {@code True} if dragging a grid item, {@code False} otherwise.
      */
-    internal fun onDragStart(offset: Offset, contentOffset: Offset): Boolean {
+    internal fun onDragStart(
+        offset: Offset,
+        screenWidth: Int,
+        layoutDirection: LayoutDirection,
+        contentOffset: Offset
+    ): Boolean {
+        val normalizedOffset =
+            Offset(
+                if (layoutDirection == LayoutDirection.Ltr) offset.x else screenWidth - offset.x,
+                offset.y
+            )
         state.layoutInfo.visibleItemsInfo
             .filter { item -> contentListState.isItemEditable(item.index) }
             // grid item offset is based off grid content container so we need to deduct
             // before content padding from the initial pointer position
-            .firstItemAtOffset(offset - contentOffset)
+            .firstItemAtOffset(normalizedOffset - contentOffset)
             ?.apply {
-                dragStartPointerOffset = offset - this.offset.toOffset()
+                dragStartPointerOffset = normalizedOffset - this.offset.toOffset()
                 draggingItemIndex = index
                 draggingItemInitialOffset = this.offset.toOffset()
                 return true
@@ -145,8 +160,10 @@
         dragStartPointerOffset = Offset.Zero
     }
 
-    internal fun onDrag(offset: Offset) {
-        draggingItemDraggedDelta += offset
+    internal fun onDrag(offset: Offset, layoutDirection: LayoutDirection) {
+        // Adjust offset to match the layout direction
+        draggingItemDraggedDelta +=
+            Offset(offset.x.directional(LayoutDirection.Ltr, layoutDirection), offset.y)
 
         val draggingItem = draggingItemLayoutInfo ?: return
         val startOffset = draggingItem.offset.toOffset() + draggingItemOffset
@@ -213,6 +230,8 @@
 
 fun Modifier.dragContainer(
     dragDropState: GridDragDropState,
+    layoutDirection: LayoutDirection,
+    screenWidth: Int,
     contentOffset: Offset,
     viewModel: BaseCommunalViewModel,
 ): Modifier {
@@ -221,10 +240,17 @@
             detectDragGesturesAfterLongPress(
                 onDrag = { change, offset ->
                     change.consume()
-                    dragDropState.onDrag(offset = offset)
+                    dragDropState.onDrag(offset, layoutDirection)
                 },
                 onDragStart = { offset ->
-                    if (dragDropState.onDragStart(offset, contentOffset)) {
+                    if (
+                        dragDropState.onDragStart(
+                            offset,
+                            screenWidth,
+                            layoutDirection,
+                            contentOffset
+                        )
+                    ) {
                         viewModel.onReorderWidgetStart()
                     }
                 },
@@ -262,10 +288,12 @@
             targetValue = if (dragDropState.isDraggingToRemove) 0.5f else 1f,
             label = "DraggableItemAlpha"
         )
+    val direction = LocalLayoutDirection.current
     val draggingModifier =
         if (dragging) {
             Modifier.graphicsLayer {
-                translationX = dragDropState.draggingItemOffset.x
+                translationX =
+                    dragDropState.draggingItemOffset.x.directional(LayoutDirection.Ltr, direction)
                 translationY = dragDropState.draggingItemOffset.y
                 alpha = itemAlpha
             }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/blueprint/DefaultBlueprint.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/blueprint/DefaultBlueprint.kt
index c01039d..a9e63c6 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/blueprint/DefaultBlueprint.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/blueprint/DefaultBlueprint.kt
@@ -31,6 +31,7 @@
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.modifiers.padding
+import com.android.compose.modifiers.thenIf
 import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.keyguard.ui.composable.LockscreenLongPress
 import com.android.systemui.keyguard.ui.composable.section.AmbientIndicationSection
@@ -94,9 +95,12 @@
                             with(topAreaSection) {
                                 DefaultClockLayout(
                                     modifier =
-                                        Modifier.fillMaxWidth(0.5f).graphicsLayer {
-                                            translationX = unfoldTranslations.start
-                                        }
+                                        Modifier.thenIf(isShadeLayoutWide) {
+                                                Modifier.fillMaxWidth(0.5f)
+                                            }
+                                            .graphicsLayer {
+                                                translationX = unfoldTranslations.start
+                                            }
                                 )
                             }
                             if (isShadeLayoutWide) {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt
index b40bccb..6feaf6d 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/BottomAreaSection.kt
@@ -32,6 +32,7 @@
 import androidx.core.content.res.ResourcesCompat
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.SceneScope
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.animation.view.LaunchableImageView
 import com.android.systemui.keyguard.ui.binder.KeyguardIndicationAreaBinder
 import com.android.systemui.keyguard.ui.binder.KeyguardQuickAffordanceViewBinder
@@ -55,6 +56,7 @@
     private val vibratorHelper: VibratorHelper,
     private val indicationController: KeyguardIndicationController,
     private val indicationAreaViewModel: KeyguardIndicationAreaViewModel,
+    private val shortcutsLogger: KeyguardQuickAffordancesLogger,
 ) {
     /**
      * Renders a single lockscreen shortcut.
@@ -162,6 +164,7 @@
                         transitionAlpha,
                         falsingManager,
                         vibratorHelper,
+                        shortcutsLogger,
                     ) {
                         indicationController.showTransientIndication(it)
                     }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaScenePicker.kt b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaScenePicker.kt
index 0bd0809..7b497e8 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaScenePicker.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaScenePicker.kt
@@ -19,6 +19,7 @@
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.ElementScenePicker
 import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.SceneTransitionLayoutState
 import com.android.compose.animation.scene.TransitionState
 import com.android.systemui.scene.shared.model.Scenes
 
@@ -45,22 +46,28 @@
             shouldElevateMedia(transition) -> {
                 Scenes.Shade
             }
-
-            // TODO: 345467290 - update with the actual scene picking
+            transition.isTransitioningBetween(Scenes.Lockscreen, Scenes.Communal) -> {
+                Scenes.Lockscreen
+            }
             transition.isTransitioningBetween(Scenes.QuickSettings, Scenes.Shade) -> {
                 Scenes.QuickSettings
             }
-
-            // TODO: 340216785 - update with the actual scene picking
-            else -> pickSingleSceneIn(scenes, transition, element)
+            else -> {
+                when {
+                    scenes.contains(transition.toScene) -> transition.toScene
+                    scenes.contains(transition.fromScene) -> transition.fromScene
+                    else -> null
+                }
+            }
         }
     }
 
     /** Returns true when the media should be laid on top of the rest for the given [transition]. */
-    fun shouldElevateMedia(transition: TransitionState.Transition?): Boolean {
-        if (transition == null) {
-            return false
-        }
+    fun shouldElevateMedia(transition: TransitionState.Transition): Boolean {
         return transition.isTransitioningBetween(Scenes.Lockscreen, Scenes.Shade)
     }
 }
+
+fun MediaScenePicker.shouldElevateMedia(layoutState: SceneTransitionLayoutState): Boolean {
+    return layoutState.currentTransition?.let { shouldElevateMedia(it) } ?: false
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
index c4970c5..76a7a10 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
@@ -290,6 +290,7 @@
     val isCurrentGestureOverscroll =
         viewModel.isCurrentGestureOverscroll.collectAsStateWithLifecycle(false)
     val expansionFraction by viewModel.expandFraction.collectAsStateWithLifecycle(0f)
+    val shadeToQsFraction by viewModel.shadeToQsFraction.collectAsStateWithLifecycle(0f)
 
     val topPadding = dimensionResource(id = R.dimen.notification_side_paddings)
     val navBarHeight = WindowInsets.systemBars.asPaddingValues().calculateBottomPadding()
@@ -385,14 +386,26 @@
             modifier
                 .element(Notifications.Elements.NotificationScrim)
                 .offset {
-                    // if scrim is expanded while transitioning to Gone scene, increase the offset
-                    // in step with the transition so that it is 0 when it completes.
+                    // if scrim is expanded while transitioning to Gone or QS scene, increase the
+                    // offset in step with the corresponding transition so that it is 0 when it
+                    // completes.
                     if (
                         scrimOffset.value < 0 &&
                             layoutState.isTransitioning(from = Scenes.Shade, to = Scenes.Gone) ||
                             layoutState.isTransitioning(from = Scenes.Shade, to = Scenes.Lockscreen)
                     ) {
                         IntOffset(x = 0, y = (scrimOffset.value * expansionFraction).roundToInt())
+                    } else if (
+                        scrimOffset.value < 0 &&
+                            layoutState.isTransitioning(
+                                from = Scenes.Shade,
+                                to = Scenes.QuickSettings
+                            )
+                    ) {
+                        IntOffset(
+                            x = 0,
+                            y = (scrimOffset.value * (1 - shadeToQsFraction)).roundToInt()
+                        )
                     } else {
                         IntOffset(x = 0, y = scrimOffset.value.roundToInt())
                     }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToSplitShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToSplitShadeTransition.kt
index f14ff76..2f8c248 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToSplitShadeTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToSplitShadeTransition.kt
@@ -16,10 +16,58 @@
 
 package com.android.systemui.scene.ui.composable.transitions
 
+import androidx.compose.animation.core.Spring
+import androidx.compose.animation.core.spring
+import androidx.compose.animation.core.tween
+import androidx.compose.foundation.gestures.Orientation
+import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.TransitionBuilder
+import com.android.compose.animation.scene.UserActionDistance
+import com.android.compose.animation.scene.UserActionDistanceScope
+import com.android.systemui.media.controls.ui.composable.MediaCarousel
+import com.android.systemui.media.controls.ui.composable.MediaScenePicker
+import com.android.systemui.notifications.ui.composable.Notifications
+import com.android.systemui.qs.ui.composable.QuickSettings
+import com.android.systemui.shade.ui.composable.Shade
+import com.android.systemui.shade.ui.composable.ShadeHeader
+import kotlin.time.Duration.Companion.milliseconds
 
 fun TransitionBuilder.goneToSplitShadeTransition(
     durationScale: Double = 1.0,
 ) {
-    toSplitShadeTransition(durationScale)
+    spec = tween(durationMillis = (DefaultDuration * durationScale).inWholeMilliseconds.toInt())
+    swipeSpec =
+        spring(
+            stiffness = Spring.StiffnessMediumLow,
+            visibilityThreshold = Shade.Dimensions.ScrimVisibilityThreshold,
+        )
+    distance =
+        object : UserActionDistance {
+            override fun UserActionDistanceScope.absoluteDistance(
+                fromSceneSize: IntSize,
+                orientation: Orientation,
+            ): Float {
+                return fromSceneSize.height.toFloat() * 2 / 3f
+            }
+        }
+
+    fractionRange(end = .33f) { fade(Shade.Elements.BackgroundScrim) }
+
+    fractionRange(start = .33f) {
+        val qsTranslation = ShadeHeader.Dimensions.CollapsedHeight * MediaScenePicker.SHADE_FRACTION
+        val qsExpansionDiff =
+            ShadeHeader.Dimensions.ExpandedHeight - ShadeHeader.Dimensions.CollapsedHeight
+        translate(MediaCarousel.Elements.Content, y = -(qsExpansionDiff + qsTranslation))
+        fade(MediaCarousel.Elements.Content)
+
+        fade(ShadeHeader.Elements.Clock)
+        fade(ShadeHeader.Elements.CollapsedContentStart)
+        fade(ShadeHeader.Elements.CollapsedContentEnd)
+        fade(ShadeHeader.Elements.PrivacyChip)
+        fade(QuickSettings.Elements.SplitShadeQuickSettings)
+        fade(QuickSettings.Elements.FooterActions)
+        fade(Notifications.Elements.NotificationScrim)
+    }
 }
+
+private val DefaultDuration = 500.milliseconds
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToCommunalTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToCommunalTransition.kt
index 0021bf5..5401936 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToCommunalTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToCommunalTransition.kt
@@ -25,8 +25,8 @@
     spec = tween(durationMillis = 500)
 
     // Translate lockscreen to the left.
-    translate(Scenes.Lockscreen.rootElementKey, Edge.Left)
+    translate(Scenes.Lockscreen.rootElementKey, Edge.Start)
 
     // Translate communal from the right.
-    translate(Scenes.Communal.rootElementKey, Edge.Right)
+    translate(Scenes.Communal.rootElementKey, Edge.End)
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToSplitShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToSplitShadeTransition.kt
index 70c343c..1486ea7 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToSplitShadeTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToSplitShadeTransition.kt
@@ -16,10 +16,50 @@
 
 package com.android.systemui.scene.ui.composable.transitions
 
+import androidx.compose.animation.core.Spring
+import androidx.compose.animation.core.spring
+import androidx.compose.animation.core.tween
+import androidx.compose.foundation.gestures.Orientation
+import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.TransitionBuilder
+import com.android.compose.animation.scene.UserActionDistance
+import com.android.compose.animation.scene.UserActionDistanceScope
+import com.android.systemui.notifications.ui.composable.Notifications
+import com.android.systemui.qs.ui.composable.QuickSettings
+import com.android.systemui.shade.ui.composable.Shade
+import com.android.systemui.shade.ui.composable.ShadeHeader
+import kotlin.time.Duration.Companion.milliseconds
 
 fun TransitionBuilder.lockscreenToSplitShadeTransition(
     durationScale: Double = 1.0,
 ) {
-    toSplitShadeTransition(durationScale = durationScale)
+    spec = tween(durationMillis = (DefaultDuration * durationScale).inWholeMilliseconds.toInt())
+    swipeSpec =
+        spring(
+            stiffness = Spring.StiffnessMediumLow,
+            visibilityThreshold = Shade.Dimensions.ScrimVisibilityThreshold,
+        )
+    distance =
+        object : UserActionDistance {
+            override fun UserActionDistanceScope.absoluteDistance(
+                fromSceneSize: IntSize,
+                orientation: Orientation,
+            ): Float {
+                return fromSceneSize.height.toFloat() * 2 / 3f
+            }
+        }
+
+    fractionRange(end = .33f) { fade(Shade.Elements.BackgroundScrim) }
+
+    fractionRange(start = .33f) {
+        fade(ShadeHeader.Elements.Clock)
+        fade(ShadeHeader.Elements.CollapsedContentStart)
+        fade(ShadeHeader.Elements.CollapsedContentEnd)
+        fade(ShadeHeader.Elements.PrivacyChip)
+        fade(QuickSettings.Elements.SplitShadeQuickSettings)
+        fade(QuickSettings.Elements.FooterActions)
+        fade(Notifications.Elements.NotificationScrim)
+    }
 }
+
+private val DefaultDuration = 500.milliseconds
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToSplitShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToSplitShadeTransition.kt
deleted file mode 100644
index a8315c0..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToSplitShadeTransition.kt
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * Copyright (C) 2024 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.scene.ui.composable.transitions
-
-import androidx.compose.animation.core.Spring
-import androidx.compose.animation.core.spring
-import androidx.compose.animation.core.tween
-import androidx.compose.foundation.gestures.Orientation
-import androidx.compose.ui.unit.IntSize
-import com.android.compose.animation.scene.TransitionBuilder
-import com.android.compose.animation.scene.UserActionDistance
-import com.android.compose.animation.scene.UserActionDistanceScope
-import com.android.systemui.notifications.ui.composable.Notifications
-import com.android.systemui.qs.ui.composable.QuickSettings
-import com.android.systemui.shade.ui.composable.Shade
-import com.android.systemui.shade.ui.composable.ShadeHeader
-import kotlin.time.Duration.Companion.milliseconds
-
-fun TransitionBuilder.toSplitShadeTransition(
-    durationScale: Double = 1.0,
-) {
-    spec = tween(durationMillis = (DefaultDuration * durationScale).inWholeMilliseconds.toInt())
-    swipeSpec =
-        spring(
-            stiffness = Spring.StiffnessMediumLow,
-            visibilityThreshold = Shade.Dimensions.ScrimVisibilityThreshold,
-        )
-    distance =
-        object : UserActionDistance {
-            override fun UserActionDistanceScope.absoluteDistance(
-                fromSceneSize: IntSize,
-                orientation: Orientation,
-            ): Float {
-                return fromSceneSize.height.toFloat() * 2 / 3f
-            }
-        }
-
-    fractionRange(end = .33f) { fade(Shade.Elements.BackgroundScrim) }
-
-    fractionRange(start = .33f) {
-        fade(ShadeHeader.Elements.Clock)
-        fade(ShadeHeader.Elements.CollapsedContentStart)
-        fade(ShadeHeader.Elements.CollapsedContentEnd)
-        fade(ShadeHeader.Elements.PrivacyChip)
-        fade(QuickSettings.Elements.SplitShadeQuickSettings)
-        fade(QuickSettings.Elements.FooterActions)
-        fade(Notifications.Elements.NotificationScrim)
-    }
-}
-
-private val DefaultDuration = 500.milliseconds
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
index 21e3431..18ca0f7 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
@@ -60,6 +60,7 @@
 import androidx.compose.ui.platform.LocalLifecycleOwner
 import androidx.compose.ui.res.colorResource
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.zIndex
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.LowestZIndexScenePicker
@@ -81,10 +82,12 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.media.controls.ui.composable.MediaCarousel
 import com.android.systemui.media.controls.ui.composable.MediaScenePicker
+import com.android.systemui.media.controls.ui.composable.shouldElevateMedia
 import com.android.systemui.media.controls.ui.controller.MediaCarouselController
 import com.android.systemui.media.controls.ui.controller.MediaHierarchyManager
 import com.android.systemui.media.controls.ui.view.MediaHost
 import com.android.systemui.media.controls.ui.view.MediaHostState
+import com.android.systemui.media.dagger.MediaModule.QS_PANEL
 import com.android.systemui.media.dagger.MediaModule.QUICK_QS_PANEL
 import com.android.systemui.notifications.ui.composable.NotificationScrollingStack
 import com.android.systemui.notifications.ui.composable.NotificationStackCutoffGuideline
@@ -148,7 +151,8 @@
     private val batteryMeterViewControllerFactory: BatteryMeterViewController.Factory,
     private val statusBarIconController: StatusBarIconController,
     private val mediaCarouselController: MediaCarouselController,
-    @Named(QUICK_QS_PANEL) private val mediaHost: MediaHost,
+    @Named(QUICK_QS_PANEL) private val qqsMediaHost: MediaHost,
+    @Named(QS_PANEL) private val qsMediaHost: MediaHost,
 ) : ComposableScene {
 
     override val key = Scenes.Shade
@@ -172,15 +176,20 @@
             createBatteryMeterViewController = batteryMeterViewControllerFactory::create,
             statusBarIconController = statusBarIconController,
             mediaCarouselController = mediaCarouselController,
-            mediaHost = mediaHost,
+            qqsMediaHost = qqsMediaHost,
+            qsMediaHost = qsMediaHost,
             modifier = modifier,
             shadeSession = shadeSession,
         )
 
     init {
-        mediaHost.expansion = MediaHostState.EXPANDED
-        mediaHost.showsOnlyActiveMedia = true
-        mediaHost.init(MediaHierarchyManager.LOCATION_QQS)
+        qqsMediaHost.expansion = MediaHostState.EXPANDED
+        qqsMediaHost.showsOnlyActiveMedia = true
+        qqsMediaHost.init(MediaHierarchyManager.LOCATION_QQS)
+
+        qsMediaHost.expansion = MediaHostState.EXPANDED
+        qsMediaHost.showsOnlyActiveMedia = false
+        qsMediaHost.init(MediaHierarchyManager.LOCATION_QS)
     }
 }
 
@@ -193,7 +202,8 @@
     createBatteryMeterViewController: (ViewGroup, StatusBarLocation) -> BatteryMeterViewController,
     statusBarIconController: StatusBarIconController,
     mediaCarouselController: MediaCarouselController,
-    mediaHost: MediaHost,
+    qqsMediaHost: MediaHost,
+    qsMediaHost: MediaHost,
     modifier: Modifier = Modifier,
     shadeSession: SaveableSession,
 ) {
@@ -208,7 +218,7 @@
                 createBatteryMeterViewController = createBatteryMeterViewController,
                 statusBarIconController = statusBarIconController,
                 mediaCarouselController = mediaCarouselController,
-                mediaHost = mediaHost,
+                mediaHost = qqsMediaHost,
                 modifier = modifier,
                 shadeSession = shadeSession,
             )
@@ -221,7 +231,7 @@
                 createBatteryMeterViewController = createBatteryMeterViewController,
                 statusBarIconController = statusBarIconController,
                 mediaCarouselController = mediaCarouselController,
-                mediaHost = mediaHost,
+                mediaHost = qsMediaHost,
                 modifier = modifier,
                 shadeSession = shadeSession,
             )
@@ -366,7 +376,7 @@
 
             layout(constraints.maxWidth, constraints.maxHeight) {
                 val qsZIndex =
-                    if (MediaScenePicker.shouldElevateMedia(layoutState.currentTransition)) {
+                    if (MediaScenePicker.shouldElevateMedia(layoutState)) {
                         1f
                     } else {
                         0f
@@ -475,17 +485,20 @@
 
     val brightnessMirrorShowingModifier = Modifier.graphicsLayer { alpha = contentAlpha }
 
-    Box(
-        modifier =
-            modifier
-                .fillMaxSize()
-                .element(Shade.Elements.BackgroundScrim)
-                // Cannot set the alpha of the whole element to 0, because the mirror should be
-                // in the QS column.
-                .background(
-                    colorResource(R.color.shade_scrim_background_dark).copy(alpha = contentAlpha)
-                )
-    ) {
+    Box {
+        Box(
+            modifier =
+                modifier
+                    .fillMaxSize()
+                    .element(Shade.Elements.BackgroundScrim)
+                    // Cannot set the alpha of the whole element to 0, because the mirror should be
+                    // in the QS column.
+                    .background(
+                        colorResource(R.color.shade_scrim_background_dark)
+                            .copy(alpha = contentAlpha)
+                    )
+        )
+
         Column(
             modifier = Modifier.fillMaxSize(),
         ) {
@@ -545,11 +558,15 @@
                                     squishiness = { tileSquishiness },
                                 )
                             }
-
                             MediaCarousel(
                                 isVisible = isMediaVisible,
                                 mediaHost = mediaHost,
-                                modifier = Modifier.fillMaxWidth(),
+                                modifier =
+                                    Modifier.fillMaxWidth().thenIf(
+                                        MediaScenePicker.shouldElevateMedia(layoutState)
+                                    ) {
+                                        Modifier.zIndex(1f)
+                                    },
                                 carouselController = mediaCarouselController,
                             )
                         }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
index 324e7bd..b329534 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
@@ -16,12 +16,16 @@
 
 package com.android.compose.animation.scene
 
+import androidx.annotation.VisibleForTesting
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.gestures.awaitHorizontalTouchSlopOrCancellation
 import androidx.compose.foundation.gestures.awaitVerticalTouchSlopOrCancellation
 import androidx.compose.runtime.Stable
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
+import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
+import androidx.compose.ui.input.nestedscroll.NestedScrollDispatcher
+import androidx.compose.ui.input.nestedscroll.NestedScrollSource
 import androidx.compose.ui.input.pointer.AwaitPointerEventScope
 import androidx.compose.ui.input.pointer.PointerEvent
 import androidx.compose.ui.input.pointer.PointerEventPass
@@ -36,13 +40,11 @@
 import androidx.compose.ui.input.pointer.util.VelocityTracker
 import androidx.compose.ui.input.pointer.util.addPointerInputChange
 import androidx.compose.ui.node.CompositionLocalConsumerModifierNode
-import androidx.compose.ui.node.DelegatableNode
 import androidx.compose.ui.node.DelegatingNode
 import androidx.compose.ui.node.ModifierNodeElement
 import androidx.compose.ui.node.ObserverModifierNode
 import androidx.compose.ui.node.PointerInputModifierNode
 import androidx.compose.ui.node.currentValueOf
-import androidx.compose.ui.node.findNearestAncestor
 import androidx.compose.ui.node.observeReads
 import androidx.compose.ui.platform.LocalViewConfiguration
 import androidx.compose.ui.unit.IntSize
@@ -51,6 +53,7 @@
 import androidx.compose.ui.util.fastAny
 import androidx.compose.ui.util.fastFirstOrNull
 import androidx.compose.ui.util.fastSumBy
+import com.android.compose.ui.util.SpaceVectorConverter
 import kotlin.coroutines.cancellation.CancellationException
 import kotlin.math.sign
 import kotlinx.coroutines.coroutineScope
@@ -71,6 +74,7 @@
  * dragged) and a second pointer is down and dragged. This is an implementation detail that might
  * change in the future.
  */
+@VisibleForTesting
 @Stable
 internal fun Modifier.multiPointerDraggable(
     orientation: Orientation,
@@ -78,6 +82,7 @@
     startDragImmediately: (startedPosition: Offset) -> Boolean,
     onDragStarted: (startedPosition: Offset, overSlop: Float, pointersDown: Int) -> DragController,
     swipeDetector: SwipeDetector = DefaultSwipeDetector,
+    dispatcher: NestedScrollDispatcher,
 ): Modifier =
     this.then(
         MultiPointerDraggableElement(
@@ -86,6 +91,7 @@
             startDragImmediately,
             onDragStarted,
             swipeDetector,
+            dispatcher,
         )
     )
 
@@ -96,6 +102,7 @@
     private val onDragStarted:
         (startedPosition: Offset, overSlop: Float, pointersDown: Int) -> DragController,
     private val swipeDetector: SwipeDetector,
+    private val dispatcher: NestedScrollDispatcher,
 ) : ModifierNodeElement<MultiPointerDraggableNode>() {
     override fun create(): MultiPointerDraggableNode =
         MultiPointerDraggableNode(
@@ -104,6 +111,7 @@
             startDragImmediately = startDragImmediately,
             onDragStarted = onDragStarted,
             swipeDetector = swipeDetector,
+            dispatcher = dispatcher,
         )
 
     override fun update(node: MultiPointerDraggableNode) {
@@ -122,11 +130,13 @@
     var onDragStarted:
         (startedPosition: Offset, overSlop: Float, pointersDown: Int) -> DragController,
     var swipeDetector: SwipeDetector = DefaultSwipeDetector,
+    private val dispatcher: NestedScrollDispatcher,
 ) :
     DelegatingNode(),
     PointerInputModifierNode,
     CompositionLocalConsumerModifierNode,
-    ObserverModifierNode {
+    ObserverModifierNode,
+    SpaceVectorConverter {
     private val pointerInputHandler: suspend PointerInputScope.() -> Unit = { pointerInput() }
     private val delegate = delegate(SuspendingPointerInputModifierNode(pointerInputHandler))
     private val velocityTracker = VelocityTracker()
@@ -141,26 +151,22 @@
             }
         }
 
-    private var _toFloat = orientation.toFunctionOffsetToFloat()
+    private var converter = SpaceVectorConverter(orientation)
 
-    private fun Offset.toFloat(): Float = _toFloat(this)
+    override fun Offset.toFloat(): Float = with(converter) { this@toFloat.toFloat() }
 
-    private fun Orientation.toFunctionOffsetToFloat(): (Offset) -> Float =
-        when (this) {
-            Orientation.Vertical -> {
-                { it.y }
-            }
-            Orientation.Horizontal -> {
-                { it.x }
-            }
-        }
+    override fun Velocity.toFloat(): Float = with(converter) { this@toFloat.toFloat() }
+
+    override fun Float.toOffset(): Offset = with(converter) { this@toOffset.toOffset() }
+
+    override fun Float.toVelocity(): Velocity = with(converter) { this@toVelocity.toVelocity() }
 
     var orientation: Orientation = orientation
         set(value) {
             // Reset the pointer input whenever orientation changed.
             if (value != field) {
                 field = value
-                _toFloat = field.toFunctionOffsetToFloat()
+                converter = SpaceVectorConverter(value)
                 delegate.resetPointerInputHandler()
             }
         }
@@ -240,28 +246,32 @@
                                 },
                                 onDrag = { controller, change, amount ->
                                     velocityTracker.addPointerInputChange(change)
-                                    controller.onDrag(amount)
+                                    dispatchScrollEvents(
+                                        availableOnPreScroll = amount,
+                                        onScroll = { controller.onDrag(it) },
+                                        source = NestedScrollSource.UserInput,
+                                    )
                                 },
                                 onDragEnd = { controller ->
-                                    val viewConfiguration = currentValueOf(LocalViewConfiguration)
-                                    val maxVelocity =
-                                        viewConfiguration.maximumFlingVelocity.let {
-                                            Velocity(it, it)
-                                        }
-                                    val velocity = velocityTracker.calculateVelocity(maxVelocity)
-                                    controller.onStop(
-                                        velocity =
-                                            when (orientation) {
-                                                Orientation.Horizontal -> velocity.x
-                                                Orientation.Vertical -> velocity.y
-                                            },
-                                        canChangeScene = true,
+                                    startFlingGesture(
+                                        initialVelocity =
+                                            currentValueOf(LocalViewConfiguration)
+                                                .maximumFlingVelocity
+                                                .let {
+                                                    val maxVelocity = Velocity(it, it)
+                                                    velocityTracker.calculateVelocity(maxVelocity)
+                                                }
+                                                .toFloat(),
+                                        onFling = { controller.onStop(it, canChangeScene = true) }
                                     )
                                 },
                                 onDragCancel = { controller ->
-                                    controller.onStop(velocity = 0f, canChangeScene = true)
+                                    startFlingGesture(
+                                        initialVelocity = 0f,
+                                        onFling = { controller.onStop(it, canChangeScene = true) }
+                                    )
                                 },
-                                swipeDetector = swipeDetector
+                                swipeDetector = swipeDetector,
                             )
                         } catch (exception: CancellationException) {
                             // If the coroutine scope is active, we can just restart the drag cycle.
@@ -276,6 +286,101 @@
     }
 
     /**
+     * Start a fling gesture in another CoroutineScope, this is to ensure that even when the pointer
+     * input scope is reset we will continue any coroutine scope that we started from these methods
+     * while the pointer input scope was active.
+     *
+     * Note: Inspired by [androidx.compose.foundation.gestures.ScrollableNode.onDragStopped]
+     */
+    private fun startFlingGesture(initialVelocity: Float, onFling: (velocity: Float) -> Float) {
+        // Note: [AwaitPointerEventScope] is annotated as @RestrictsSuspension, we need another
+        // CoroutineScope to run the fling gestures.
+        // We do not need to cancel this [Job], the source will take care of emitting an
+        // [onPostFling] before starting a new gesture.
+        dispatcher.coroutineScope.launch {
+            dispatchFlingEvents(availableOnPreFling = initialVelocity, onFling = onFling)
+        }
+    }
+
+    /**
+     * Use the nested scroll system to fire scroll events. This allows us to consume events from our
+     * ancestors during the pre-scroll and post-scroll phases.
+     *
+     * @param availableOnPreScroll amount available before the scroll, this can be partially
+     *   consumed by our ancestors.
+     * @param onScroll function that returns the amount consumed during a scroll given the amount
+     *   available after the [NestedScrollConnection.onPreScroll].
+     * @param source the source of the scroll event
+     * @return Total offset consumed.
+     */
+    private inline fun dispatchScrollEvents(
+        availableOnPreScroll: Float,
+        onScroll: (delta: Float) -> Float,
+        source: NestedScrollSource,
+    ): Float {
+        // PreScroll phase
+        val consumedByPreScroll =
+            dispatcher
+                .dispatchPreScroll(
+                    available = availableOnPreScroll.toOffset(),
+                    source = source,
+                )
+                .toFloat()
+
+        // Scroll phase
+        val availableOnScroll = availableOnPreScroll - consumedByPreScroll
+        val consumedBySelfScroll = onScroll(availableOnScroll)
+
+        // PostScroll phase
+        val availableOnPostScroll = availableOnScroll - consumedBySelfScroll
+        val consumedByPostScroll =
+            dispatcher
+                .dispatchPostScroll(
+                    consumed = consumedBySelfScroll.toOffset(),
+                    available = availableOnPostScroll.toOffset(),
+                    source = source,
+                )
+                .toFloat()
+
+        return consumedByPreScroll + consumedBySelfScroll + consumedByPostScroll
+    }
+
+    /**
+     * Use the nested scroll system to fire fling events. This allows us to consume events from our
+     * ancestors during the pre-fling and post-fling phases.
+     *
+     * @param availableOnPreFling velocity available before the fling, this can be partially
+     *   consumed by our ancestors.
+     * @param onFling function that returns the velocity consumed during the fling given the
+     *   velocity available after the [NestedScrollConnection.onPreFling].
+     * @return Total velocity consumed.
+     */
+    private suspend inline fun dispatchFlingEvents(
+        availableOnPreFling: Float,
+        onFling: (velocity: Float) -> Float,
+    ): Float {
+        // PreFling phase
+        val consumedByPreFling =
+            dispatcher.dispatchPreFling(available = availableOnPreFling.toVelocity()).toFloat()
+
+        // Fling phase
+        val availableOnFling = availableOnPreFling - consumedByPreFling
+        val consumedBySelfFling = onFling(availableOnFling)
+
+        // PostFling phase
+        val availableOnPostFling = availableOnFling - consumedBySelfFling
+        val consumedByPostFling =
+            dispatcher
+                .dispatchPostFling(
+                    consumed = consumedBySelfFling.toVelocity(),
+                    available = availableOnPostFling.toVelocity(),
+                )
+                .toFloat()
+
+        return consumedByPreFling + consumedBySelfFling + consumedByPostFling
+    }
+
+    /**
      * Detect drag gestures in the given [orientation].
      *
      * This function is a mix of [androidx.compose.foundation.gestures.awaitDownAndSlop] and
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeToScene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeToScene.kt
index b8010f2..a2118b2 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeToScene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeToScene.kt
@@ -20,6 +20,7 @@
 import androidx.compose.runtime.Stable
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
+import androidx.compose.ui.input.nestedscroll.NestedScrollDispatcher
 import androidx.compose.ui.input.nestedscroll.nestedScrollModifierNode
 import androidx.compose.ui.input.pointer.PointerEvent
 import androidx.compose.ui.input.pointer.PointerEventPass
@@ -57,6 +58,7 @@
     draggableHandler: DraggableHandlerImpl,
     swipeDetector: SwipeDetector,
 ) : DelegatingNode(), PointerInputModifierNode {
+    private val dispatcher = NestedScrollDispatcher()
     private val multiPointerDraggableNode =
         delegate(
             MultiPointerDraggableNode(
@@ -65,6 +67,7 @@
                 startDragImmediately = ::startDragImmediately,
                 onDragStarted = draggableHandler::onDragStarted,
                 swipeDetector = swipeDetector,
+                dispatcher = dispatcher,
             )
         )
 
@@ -93,7 +96,7 @@
         )
 
     init {
-        delegate(nestedScrollModifierNode(nestedScrollHandlerImpl.connection, dispatcher = null))
+        delegate(nestedScrollModifierNode(nestedScrollHandlerImpl.connection, dispatcher))
         delegate(ScrollBehaviorOwnerNode(draggableHandler.nestedScrollKey, nestedScrollHandlerImpl))
     }
 
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
index b98400a..2d37a0d 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
@@ -28,6 +28,10 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.geometry.Size
+import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
+import androidx.compose.ui.input.nestedscroll.NestedScrollDispatcher
+import androidx.compose.ui.input.nestedscroll.NestedScrollSource
+import androidx.compose.ui.input.nestedscroll.nestedScroll
 import androidx.compose.ui.input.pointer.AwaitPointerEventScope
 import androidx.compose.ui.input.pointer.PointerEventPass
 import androidx.compose.ui.input.pointer.PointerInputChange
@@ -37,6 +41,7 @@
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onRoot
 import androidx.compose.ui.test.performTouchInput
+import androidx.compose.ui.unit.Velocity
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.coroutineScope
@@ -49,17 +54,22 @@
 class MultiPointerDraggableTest {
     @get:Rule val rule = createComposeRule()
 
+    private val emptyConnection = object : NestedScrollConnection {}
+    private val defaultDispatcher = NestedScrollDispatcher()
+
+    private fun Modifier.nestedScrollDispatcher() = nestedScroll(emptyConnection, defaultDispatcher)
+
     private class SimpleDragController(
-        val onDrag: () -> Unit,
-        val onStop: () -> Unit,
+        val onDrag: (delta: Float) -> Unit,
+        val onStop: (velocity: Float) -> Unit,
     ) : DragController {
         override fun onDrag(delta: Float): Float {
-            onDrag()
+            onDrag.invoke(delta)
             return delta
         }
 
         override fun onStop(velocity: Float, canChangeScene: Boolean): Float {
-            onStop()
+            onStop.invoke(velocity)
             return velocity
         }
     }
@@ -79,6 +89,7 @@
             touchSlop = LocalViewConfiguration.current.touchSlop
             Box(
                 Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScrollDispatcher()
                     .multiPointerDraggable(
                         orientation = Orientation.Vertical,
                         enabled = { enabled },
@@ -90,6 +101,7 @@
                                 onStop = { stopped = true },
                             )
                         },
+                        dispatcher = defaultDispatcher,
                     )
             )
         }
@@ -145,6 +157,7 @@
             touchSlop = LocalViewConfiguration.current.touchSlop
             Box(
                 Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScrollDispatcher()
                     .multiPointerDraggable(
                         orientation = Orientation.Vertical,
                         enabled = { true },
@@ -157,6 +170,7 @@
                                 onStop = { stopped = true },
                             )
                         },
+                        dispatcher = defaultDispatcher,
                     )
                     .pointerInput(Unit) {
                         coroutineScope {
@@ -217,6 +231,7 @@
             touchSlop = LocalViewConfiguration.current.touchSlop
             Box(
                 Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScrollDispatcher()
                     .multiPointerDraggable(
                         orientation = Orientation.Vertical,
                         enabled = { true },
@@ -228,6 +243,7 @@
                                 onStop = { stopped = true },
                             )
                         },
+                        dispatcher = defaultDispatcher,
                     )
             ) {
                 if (hasScrollable) {
@@ -335,6 +351,7 @@
             touchSlop = LocalViewConfiguration.current.touchSlop
             Box(
                 Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScrollDispatcher()
                     .multiPointerDraggable(
                         orientation = Orientation.Vertical,
                         enabled = { true },
@@ -346,6 +363,7 @@
                                 onStop = { stopped = true },
                             )
                         },
+                        dispatcher = defaultDispatcher,
                     )
             ) {
                 Box(
@@ -436,6 +454,7 @@
             touchSlop = LocalViewConfiguration.current.touchSlop
             Box(
                 Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScrollDispatcher()
                     .multiPointerDraggable(
                         orientation = Orientation.Vertical,
                         enabled = { true },
@@ -447,6 +466,7 @@
                                 onStop = { verticalStopped = true },
                             )
                         },
+                        dispatcher = defaultDispatcher,
                     )
                     .multiPointerDraggable(
                         orientation = Orientation.Horizontal,
@@ -459,6 +479,7 @@
                                 onStop = { horizontalStopped = true },
                             )
                         },
+                        dispatcher = defaultDispatcher,
                     )
             )
         }
@@ -539,6 +560,7 @@
             touchSlop = LocalViewConfiguration.current.touchSlop
             Box(
                 Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScrollDispatcher()
                     .multiPointerDraggable(
                         orientation = Orientation.Vertical,
                         enabled = { true },
@@ -557,6 +579,7 @@
                                 onStop = { /* do nothing */ },
                             )
                         },
+                        dispatcher = defaultDispatcher,
                     )
             ) {}
         }
@@ -587,4 +610,113 @@
 
         assertThat(started).isTrue()
     }
+
+    @Test
+    fun multiPointerNestedScrollDispatcher() {
+        val size = 200f
+        val middle = Offset(size / 2f, size / 2f)
+        var touchSlop = 0f
+
+        var consumedOnPreScroll = 0f
+
+        var availableOnPreScroll = Float.MIN_VALUE
+        var availableOnPostScroll = Float.MIN_VALUE
+        var availableOnPreFling = Float.MIN_VALUE
+        var availableOnPostFling = Float.MIN_VALUE
+
+        var consumedOnDrag = 0f
+        var consumedOnDragStop = 0f
+
+        val connection =
+            object : NestedScrollConnection {
+                override fun onPreScroll(available: Offset, source: NestedScrollSource): Offset {
+                    availableOnPreScroll = available.y
+                    return Offset(0f, consumedOnPreScroll)
+                }
+
+                override fun onPostScroll(
+                    consumed: Offset,
+                    available: Offset,
+                    source: NestedScrollSource
+                ): Offset {
+                    availableOnPostScroll = available.y
+                    return Offset.Zero
+                }
+
+                override suspend fun onPreFling(available: Velocity): Velocity {
+                    availableOnPreFling = available.y
+                    return Velocity.Zero
+                }
+
+                override suspend fun onPostFling(
+                    consumed: Velocity,
+                    available: Velocity
+                ): Velocity {
+                    availableOnPostFling = available.y
+                    return Velocity.Zero
+                }
+            }
+
+        rule.setContent {
+            touchSlop = LocalViewConfiguration.current.touchSlop
+            Box(
+                Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScroll(connection)
+                    .nestedScrollDispatcher()
+                    .multiPointerDraggable(
+                        orientation = Orientation.Vertical,
+                        enabled = { true },
+                        startDragImmediately = { false },
+                        onDragStarted = { _, _, _ ->
+                            SimpleDragController(
+                                onDrag = { consumedOnDrag = it },
+                                onStop = { consumedOnDragStop = it },
+                            )
+                        },
+                        dispatcher = defaultDispatcher,
+                    )
+            )
+        }
+
+        fun startDrag() {
+            rule.onRoot().performTouchInput {
+                down(middle)
+                moveBy(Offset(0f, touchSlop))
+            }
+        }
+
+        fun continueDrag() {
+            rule.onRoot().performTouchInput { moveBy(Offset(0f, touchSlop)) }
+        }
+
+        fun stopDrag() {
+            rule.onRoot().performTouchInput { up() }
+        }
+
+        startDrag()
+
+        continueDrag()
+        assertThat(availableOnPreScroll).isEqualTo(touchSlop)
+        assertThat(consumedOnDrag).isEqualTo(touchSlop)
+        assertThat(availableOnPostScroll).isEqualTo(0f)
+
+        // Parent node consumes half of the gesture
+        consumedOnPreScroll = touchSlop / 2f
+        continueDrag()
+        assertThat(availableOnPreScroll).isEqualTo(touchSlop)
+        assertThat(consumedOnDrag).isEqualTo(touchSlop / 2f)
+        assertThat(availableOnPostScroll).isEqualTo(0f)
+
+        // Parent node consumes the gesture
+        consumedOnPreScroll = touchSlop
+        continueDrag()
+        assertThat(availableOnPreScroll).isEqualTo(touchSlop)
+        assertThat(consumedOnDrag).isEqualTo(0f)
+        assertThat(availableOnPostScroll).isEqualTo(0f)
+
+        // Parent node can intercept the velocity on stop
+        stopDrag()
+        assertThat(availableOnPreFling).isEqualTo(consumedOnDragStop)
+        assertThat(availableOnPostFling).isEqualTo(0f)
+    }
 }
diff --git a/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt b/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt
index 9e857deb..5eca5b4 100644
--- a/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt
+++ b/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt
@@ -18,9 +18,11 @@
 
 import android.provider.Settings
 import com.android.systemui.shared.settings.data.repository.SecureSettingsRepository
+import com.android.systemui.shared.settings.data.repository.SystemSettingsRepository
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOn
@@ -33,6 +35,7 @@
     private val scope: CoroutineScope,
     private val backgroundDispatcher: CoroutineDispatcher,
     private val secureSettingsRepository: SecureSettingsRepository,
+    private val systemSettingsRepository: SystemSettingsRepository,
 ) {
     val isNotificationHistoryEnabled: Flow<Boolean> =
         secureSettingsRepository
@@ -60,4 +63,15 @@
             )
         }
     }
+
+    val isCooldownEnabled: StateFlow<Boolean> =
+        systemSettingsRepository
+            .intSetting(name = Settings.System.NOTIFICATION_COOLDOWN_ENABLED)
+            .map { it == 1 }
+            .flowOn(backgroundDispatcher)
+            .stateIn(
+                scope = scope,
+                started = SharingStarted.Eagerly,
+                initialValue = false,
+            )
 }
diff --git a/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/domain/interactor/NotificationSettingsInteractor.kt b/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/domain/interactor/NotificationSettingsInteractor.kt
index b4105bd..a274f1d 100644
--- a/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/domain/interactor/NotificationSettingsInteractor.kt
+++ b/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/domain/interactor/NotificationSettingsInteractor.kt
@@ -38,4 +38,5 @@
         val current = repository.isShowNotificationsOnLockScreenEnabled().value
         repository.setShowNotificationsOnLockscreenEnabled(!current)
     }
-}
+
+    val isCooldownEnabled = repository.isCooldownEnabled}
diff --git a/packages/SystemUI/customization/src/com/android/systemui/shared/settings/data/repository/SystemSettingsRepository.kt b/packages/SystemUI/customization/src/com/android/systemui/shared/settings/data/repository/SystemSettingsRepository.kt
new file mode 100644
index 0000000..afe82fb
--- /dev/null
+++ b/packages/SystemUI/customization/src/com/android/systemui/shared/settings/data/repository/SystemSettingsRepository.kt
@@ -0,0 +1,117 @@
+/*
+ * Copyright (C) 2024 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.shared.settings.data.repository
+
+import android.content.ContentResolver
+import android.database.ContentObserver
+import android.provider.Settings
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.withContext
+
+/**
+ * Defines interface for classes that can provide access to data from [Settings.System]. This
+ * repository doesn't guarantee to provide value across different users. For that see:
+ * [UserAwareSecureSettingsRepository] which does that for secure settings.
+ */
+interface SystemSettingsRepository {
+
+    /** Returns a [Flow] tracking the value of a setting as an [Int]. */
+    fun intSetting(
+        name: String,
+        defaultValue: Int = 0,
+    ): Flow<Int>
+
+    /** Updates the value of the setting with the given name. */
+    suspend fun setInt(
+        name: String,
+        value: Int,
+    )
+
+    suspend fun getInt(
+        name: String,
+        defaultValue: Int = 0,
+    ): Int
+
+    suspend fun getString(name: String): String?
+}
+
+class SystemSettingsRepositoryImpl(
+    private val contentResolver: ContentResolver,
+    private val backgroundDispatcher: CoroutineDispatcher,
+) : SystemSettingsRepository {
+
+    override fun intSetting(
+        name: String,
+        defaultValue: Int,
+    ): Flow<Int> {
+        return callbackFlow {
+                val observer =
+                    object : ContentObserver(null) {
+                        override fun onChange(selfChange: Boolean) {
+                            trySend(Unit)
+                        }
+                    }
+
+                contentResolver.registerContentObserver(
+                    Settings.System.getUriFor(name),
+                    /* notifyForDescendants= */ false,
+                    observer,
+                )
+                send(Unit)
+
+                awaitClose { contentResolver.unregisterContentObserver(observer) }
+            }
+            .map { Settings.System.getInt(contentResolver, name, defaultValue) }
+            // The above work is done on the background thread (which is important for accessing
+            // settings through the content resolver).
+            .flowOn(backgroundDispatcher)
+    }
+
+    override suspend fun setInt(name: String, value: Int) {
+        withContext(backgroundDispatcher) {
+            Settings.System.putInt(
+                contentResolver,
+                name,
+                value,
+            )
+        }
+    }
+
+    override suspend fun getInt(name: String, defaultValue: Int): Int {
+        return withContext(backgroundDispatcher) {
+            Settings.System.getInt(
+                contentResolver,
+                name,
+                defaultValue,
+            )
+        }
+    }
+
+    override suspend fun getString(name: String): String? {
+        return withContext(backgroundDispatcher) {
+            Settings.System.getString(
+                contentResolver,
+                name,
+            )
+        }
+    }
+}
diff --git a/packages/SystemUI/customization/tests/utils/src/com/android/systemui/shared/settings/data/repository/FakeSystemSettingsRepository.kt b/packages/SystemUI/customization/tests/utils/src/com/android/systemui/shared/settings/data/repository/FakeSystemSettingsRepository.kt
new file mode 100644
index 0000000..7da2b40
--- /dev/null
+++ b/packages/SystemUI/customization/tests/utils/src/com/android/systemui/shared/settings/data/repository/FakeSystemSettingsRepository.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2024 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.shared.settings.data.repository
+
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.map
+
+class FakeSystemSettingsRepository : SystemSettingsRepository {
+
+    private val settings = MutableStateFlow<Map<String, String>>(mutableMapOf())
+
+    override fun intSetting(name: String, defaultValue: Int): Flow<Int> {
+        return settings.map { it.getOrDefault(name, defaultValue.toString()) }.map { it.toInt() }
+    }
+
+    override suspend fun setInt(name: String, value: Int) {
+        settings.value = settings.value.toMutableMap().apply { this[name] = value.toString() }
+    }
+
+    override suspend fun getInt(name: String, defaultValue: Int): Int {
+        return settings.value[name]?.toInt() ?: defaultValue
+    }
+
+    override suspend fun getString(name: String): String? {
+        return settings.value[name]
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerOverlayTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerOverlayTest.kt
index 242e822..e2a6a55 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerOverlayTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerOverlayTest.kt
@@ -35,6 +35,8 @@
 import android.view.accessibility.AccessibilityManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.app.viewcapture.ViewCapture
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
@@ -70,6 +72,7 @@
 import com.android.systemui.testKosmos
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
 import com.google.common.truth.Truth.assertThat
+import dagger.Lazy
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
@@ -108,6 +111,7 @@
 
     @Mock private lateinit var inflater: LayoutInflater
     @Mock private lateinit var windowManager: WindowManager
+    @Mock private lateinit var lazyViewCapture: kotlin.Lazy<ViewCapture>
     @Mock private lateinit var accessibilityManager: AccessibilityManager
     @Mock private lateinit var statusBarStateController: StatusBarStateController
     @Mock private lateinit var statusBarKeyguardViewManager: StatusBarKeyguardViewManager
@@ -192,7 +196,8 @@
             UdfpsControllerOverlay(
                 context,
                 inflater,
-                windowManager,
+                ViewCaptureAwareWindowManager(windowManager, lazyViewCapture,
+                        isViewCaptureEnabled = false),
                 accessibilityManager,
                 statusBarStateController,
                 statusBarKeyguardViewManager,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerTest.java
index 54e0725..d86890b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/UdfpsControllerTest.java
@@ -65,12 +65,12 @@
 import android.view.View;
 import android.view.ViewGroup;
 import android.view.ViewRootImpl;
-import android.view.WindowManager;
 import android.view.accessibility.AccessibilityManager;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.internal.logging.InstanceIdSequence;
 import com.android.internal.util.LatencyTracker;
 import com.android.keyguard.KeyguardUpdateMonitor;
@@ -87,6 +87,7 @@
 import com.android.systemui.biometrics.ui.viewmodel.DeviceEntryUdfpsTouchOverlayViewModel;
 import com.android.systemui.bouncer.domain.interactor.AlternateBouncerInteractor;
 import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor;
+import com.android.systemui.camera.CameraGestureHelper;
 import com.android.systemui.classifier.FalsingCollector;
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryFaceAuthInteractor;
 import com.android.systemui.dump.DumpManager;
@@ -118,6 +119,8 @@
 
 import dagger.Lazy;
 
+import javax.inject.Provider;
+
 import kotlinx.coroutines.CoroutineScope;
 
 import org.junit.Before;
@@ -152,7 +155,7 @@
     @Mock
     private FingerprintManager mFingerprintManager;
     @Mock
-    private WindowManager mWindowManager;
+    private ViewCaptureAwareWindowManager mWindowManager;
     @Mock
     private StatusBarStateController mStatusBarStateController;
     @Mock
@@ -261,6 +264,8 @@
     private Lazy<DeviceEntryUdfpsTouchOverlayViewModel> mDeviceEntryUdfpsTouchOverlayViewModel;
     @Mock
     private Lazy<DefaultUdfpsTouchOverlayViewModel> mDefaultUdfpsTouchOverlayViewModel;
+    @Mock
+    private Provider<CameraGestureHelper> mCameraGestureHelper;
 
     @Before
     public void setUp() {
@@ -269,7 +274,8 @@
                 mPowerRepository,
                 mock(FalsingCollector.class),
                 mock(ScreenOffAnimationController.class),
-                mStatusBarStateController
+                mStatusBarStateController,
+                mCameraGestureHelper
         );
         mPowerRepository.updateWakefulness(
                 WakefulnessState.AWAKE,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/camera/CameraGestureHelperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/camera/CameraGestureHelperTest.kt
similarity index 67%
rename from packages/SystemUI/tests/src/com/android/systemui/camera/CameraGestureHelperTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/camera/CameraGestureHelperTest.kt
index bea0db6..a0928ad 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/camera/CameraGestureHelperTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/camera/CameraGestureHelperTest.kt
@@ -18,6 +18,7 @@
 
 import android.app.ActivityManager
 import android.app.IActivityTaskManager
+import android.app.admin.DevicePolicyManager
 import android.content.ComponentName
 import android.content.ContentResolver
 import android.content.Intent
@@ -29,82 +30,73 @@
 import com.android.systemui.ActivityIntentHelper
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.plugins.ActivityStarter
+import com.android.systemui.statusbar.NotificationLockscreenUserManager
 import com.android.systemui.statusbar.StatusBarState
-import com.android.systemui.statusbar.phone.CentralSurfaces
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
 import com.android.systemui.util.mockito.KotlinArgumentCaptor
+import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import com.google.common.util.concurrent.MoreExecutors
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.isNull
 import org.mockito.Mock
-import org.mockito.Mockito.any
 import org.mockito.Mockito.anyInt
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
-import org.mockito.Mockito.`when` as whenever
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class CameraGestureHelperTest : SysuiTestCase() {
 
-    @Mock
-    lateinit var centralSurfaces: CentralSurfaces
-    @Mock
-    lateinit var statusBarKeyguardViewManager: StatusBarKeyguardViewManager
-    @Mock
-    lateinit var keyguardStateController: KeyguardStateController
-    @Mock
-    lateinit var packageManager: PackageManager
-    @Mock
-    lateinit var activityManager: ActivityManager
-    @Mock
-    lateinit var activityStarter: ActivityStarter
-    @Mock
-    lateinit var activityIntentHelper: ActivityIntentHelper
-    @Mock
-    lateinit var activityTaskManager: IActivityTaskManager
-    @Mock
-    lateinit var cameraIntents: CameraIntentsWrapper
-    @Mock
-    lateinit var contentResolver: ContentResolver
-    @Mock
-    lateinit var mSelectedUserInteractor: SelectedUserInteractor
+    @Mock lateinit var statusBarKeyguardViewManager: StatusBarKeyguardViewManager
+    @Mock lateinit var keyguardStateController: KeyguardStateController
+    @Mock lateinit var packageManager: PackageManager
+    @Mock lateinit var activityManager: ActivityManager
+    @Mock lateinit var activityStarter: ActivityStarter
+    @Mock lateinit var activityIntentHelper: ActivityIntentHelper
+    @Mock lateinit var activityTaskManager: IActivityTaskManager
+    @Mock lateinit var cameraIntents: CameraIntentsWrapper
+    @Mock lateinit var contentResolver: ContentResolver
+    @Mock lateinit var mSelectedUserInteractor: SelectedUserInteractor
+    @Mock lateinit var devicePolicyManager: DevicePolicyManager
+    @Mock lateinit var lockscreenUserManager: NotificationLockscreenUserManager
 
     private lateinit var underTest: CameraGestureHelper
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
-        whenever(cameraIntents.getSecureCameraIntent(anyInt())).thenReturn(
-            Intent(CameraIntents.DEFAULT_SECURE_CAMERA_INTENT_ACTION)
-        )
-        whenever(cameraIntents.getInsecureCameraIntent(anyInt())).thenReturn(
-            Intent(CameraIntents.DEFAULT_INSECURE_CAMERA_INTENT_ACTION)
-        )
+        whenever(cameraIntents.getSecureCameraIntent(any()))
+            .thenReturn(Intent(CameraIntents.DEFAULT_SECURE_CAMERA_INTENT_ACTION))
+        whenever(cameraIntents.getInsecureCameraIntent(any()))
+            .thenReturn(Intent(CameraIntents.DEFAULT_INSECURE_CAMERA_INTENT_ACTION))
 
         prepare()
 
-        underTest = CameraGestureHelper(
-            context = mock(),
-            centralSurfaces = centralSurfaces,
-            keyguardStateController = keyguardStateController,
-            statusBarKeyguardViewManager = statusBarKeyguardViewManager,
-            packageManager = packageManager,
-            activityManager = activityManager,
-            activityStarter = activityStarter,
-            activityIntentHelper = activityIntentHelper,
-            activityTaskManager = activityTaskManager,
-            cameraIntents = cameraIntents,
-            contentResolver = contentResolver,
-            uiExecutor = MoreExecutors.directExecutor(),
-            selectedUserInteractor = mSelectedUserInteractor,
-        )
+        underTest =
+            CameraGestureHelper(
+                context = mock(),
+                keyguardStateController = keyguardStateController,
+                statusBarKeyguardViewManager = statusBarKeyguardViewManager,
+                packageManager = packageManager,
+                activityManager = activityManager,
+                activityStarter = activityStarter,
+                activityIntentHelper = activityIntentHelper,
+                activityTaskManager = activityTaskManager,
+                cameraIntents = cameraIntents,
+                contentResolver = contentResolver,
+                uiExecutor = MoreExecutors.directExecutor(),
+                selectedUserInteractor = mSelectedUserInteractor,
+                devicePolicyManager = devicePolicyManager,
+                lockscreenUserManager = lockscreenUserManager,
+            )
     }
 
     /**
@@ -116,13 +108,13 @@
      * @param isCameraAllowedByAdmin Whether the device administrator allows use of the camera app
      * @param installedCameraAppCount The number of installed camera apps on the device
      * @param isUsingSecureScreenLockOption Whether the user-controlled setting for Screen Lock is
-     * set with a "secure" option that requires the user to provide some secret/credentials to be
-     * able to unlock the device, for example "Face Unlock", "PIN", or "Password". Examples of
-     * non-secure options are "None" and "Swipe"
+     *   set with a "secure" option that requires the user to provide some secret/credentials to be
+     *   able to unlock the device, for example "Face Unlock", "PIN", or "Password". Examples of
+     *   non-secure options are "None" and "Swipe"
      * @param isCameraActivityRunningOnTop Whether the camera activity is running at the top of the
-     * most recent/current task of activities
+     *   most recent/current task of activities
      * @param isTaskListEmpty Whether there are no active activity tasks at all. Note that this is
-     * treated as `false` if [isCameraActivityRunningOnTop] is set to `true`
+     *   treated as `false` if [isCameraActivityRunningOnTop] is set to `true`
      */
     private fun prepare(
         isCameraAllowedByAdmin: Boolean = true,
@@ -131,7 +123,13 @@
         isCameraActivityRunningOnTop: Boolean = false,
         isTaskListEmpty: Boolean = false,
     ) {
-        whenever(centralSurfaces.isCameraAllowedByAdmin).thenReturn(isCameraAllowedByAdmin)
+        whenever(lockscreenUserManager.getCurrentUserId()).thenReturn(1)
+        if (isCameraAllowedByAdmin) {
+            whenever(devicePolicyManager.getCameraDisabled(isNull(), any())).thenReturn(false)
+            whenever(keyguardStateController.isMethodSecure).thenReturn(false)
+        } else {
+            whenever(devicePolicyManager.getCameraDisabled(isNull(), any())).thenReturn(true)
+        }
 
         whenever(activityIntentHelper.wouldLaunchResolverActivity(any(), anyInt()))
             .thenReturn(installedCameraAppCount > 1)
@@ -141,30 +139,26 @@
             .thenReturn(!isUsingSecureScreenLockOption)
 
         if (installedCameraAppCount >= 1) {
-            val resolveInfo = ResolveInfo().apply {
-                this.activityInfo = ActivityInfo().apply {
-                    packageName = CAMERA_APP_PACKAGE_NAME
+            val resolveInfo =
+                ResolveInfo().apply {
+                    this.activityInfo =
+                        ActivityInfo().apply { packageName = CAMERA_APP_PACKAGE_NAME }
                 }
-            }
-            whenever(packageManager.resolveActivityAsUser(any(), anyInt(), anyInt())).thenReturn(
-                resolveInfo
-            )
+            whenever(packageManager.resolveActivityAsUser(any(), anyInt(), anyInt()))
+                .thenReturn(resolveInfo)
         } else {
-            whenever(packageManager.resolveActivityAsUser(any(), anyInt(), anyInt())).thenReturn(
-                null
-            )
+            whenever(packageManager.resolveActivityAsUser(any(), anyInt(), anyInt()))
+                .thenReturn(null)
         }
 
         when {
             isCameraActivityRunningOnTop -> {
-                val runningTaskInfo = ActivityManager.RunningTaskInfo().apply {
-                    topActivity = ComponentName(CAMERA_APP_PACKAGE_NAME, "cameraActivity")
-                }
-                whenever(activityManager.getRunningTasks(anyInt())).thenReturn(
-                    listOf(
-                        runningTaskInfo
-                    )
-                )
+                val runningTaskInfo =
+                    ActivityManager.RunningTaskInfo().apply {
+                        topActivity = ComponentName(CAMERA_APP_PACKAGE_NAME, "cameraActivity")
+                    }
+                whenever(activityManager.getRunningTasks(anyInt()))
+                    .thenReturn(listOf(runningTaskInfo))
             }
             isTaskListEmpty -> {
                 whenever(activityManager.getRunningTasks(anyInt())).thenReturn(emptyList())
@@ -289,28 +283,28 @@
     ) {
         val intentCaptor = KotlinArgumentCaptor(Intent::class.java)
         if (isSecure && !moreThanOneCameraAppInstalled) {
-            verify(activityTaskManager).startActivityAsUser(
-                any(),
-                any(),
-                any(),
-                intentCaptor.capture(),
-                any(),
-                any(),
-                any(),
-                anyInt(),
-                anyInt(),
-                any(),
-                any(),
-                anyInt()
-            )
+            verify(activityTaskManager)
+                .startActivityAsUser(
+                    isNull(),
+                    isNull(),
+                    isNull(),
+                    intentCaptor.capture(),
+                    isNull(),
+                    isNull(),
+                    isNull(),
+                    anyInt(),
+                    anyInt(),
+                    isNull(),
+                    any(),
+                    anyInt()
+                )
         } else {
             verify(activityStarter).startActivity(intentCaptor.capture(), eq(false))
         }
         val intent = intentCaptor.value
 
         assertThat(CameraIntents.isSecureCameraIntent(intent)).isEqualTo(isSecure)
-        assertThat(intent.getIntExtra(CameraIntents.EXTRA_LAUNCH_SOURCE, -1))
-            .isEqualTo(source)
+        assertThat(intent.getIntExtra(CameraIntents.EXTRA_LAUNCH_SOURCE, -1)).isEqualTo(source)
     }
 
     companion object {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt
index c1816ed..d251585 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt
@@ -22,6 +22,7 @@
 import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.systemui.Flags.FLAG_COMMUNAL_TIMER_FLICKER_FIX
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.communal.smartspace.CommunalSmartspaceController
 import com.android.systemui.concurrency.fakeExecutor
@@ -97,6 +98,7 @@
         }
 
     @EnableFlags(FLAG_REMOTE_VIEWS)
+    @DisableFlags(FLAG_COMMUNAL_TIMER_FLICKER_FIX)
     @Test
     fun communalTimers_onlyShowTimersWithRemoteViews() =
         testScope.runTest {
@@ -138,6 +140,48 @@
             assertThat(communalTimers?.first()?.smartspaceTargetId).isEqualTo("timer-1-started")
         }
 
+    @EnableFlags(FLAG_REMOTE_VIEWS, FLAG_COMMUNAL_TIMER_FLICKER_FIX)
+    @Test
+    fun communalTimers_onlyShowTimersWithRemoteViews_timerFlickerFix() =
+        testScope.runTest {
+            underTest.startListening()
+
+            val communalTimers by collectLastValue(underTest.timers)
+            runCurrent()
+            fakeExecutor.runAllReady()
+
+            with(captureSmartspaceTargetListener()) {
+                onSmartspaceTargetsUpdated(
+                    listOf(
+                        // Invalid. Not a timer
+                        mock<SmartspaceTarget> {
+                            on { smartspaceTargetId }.doReturn("weather")
+                            on { featureType }.doReturn(SmartspaceTarget.FEATURE_WEATHER)
+                        },
+                        // Invalid. RemoteViews absent
+                        mock<SmartspaceTarget> {
+                            on { smartspaceTargetId }.doReturn("timer-0-started")
+                            on { featureType }.doReturn(SmartspaceTarget.FEATURE_TIMER)
+                            on { remoteViews }.doReturn(null)
+                            on { creationTimeMillis }.doReturn(1000)
+                        },
+                        // Valid
+                        mock<SmartspaceTarget> {
+                            on { smartspaceTargetId }.doReturn("timer-1-started")
+                            on { featureType }.doReturn(SmartspaceTarget.FEATURE_TIMER)
+                            on { remoteViews }.doReturn(mock())
+                            on { creationTimeMillis }.doReturn(2000)
+                        },
+                    )
+                )
+            }
+            runCurrent()
+
+            // Verify that only the valid target is listed
+            assertThat(communalTimers).hasSize(1)
+            assertThat(communalTimers?.first()?.smartspaceTargetId).isEqualTo("timer-1")
+        }
+
     @EnableFlags(FLAG_REMOTE_VIEWS)
     @Test
     fun communalTimers_cacheCreationTime() =
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
index 61487b0..57ce9de 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.communal.view.viewmodel
 
+import android.appwidget.AppWidgetProviderInfo
 import android.content.ActivityNotFoundException
 import android.content.ComponentName
 import android.content.Intent
@@ -24,6 +25,9 @@
 import android.content.pm.ResolveInfo
 import android.content.pm.UserInfo
 import android.provider.Settings
+import android.view.accessibility.AccessibilityEvent
+import android.view.accessibility.AccessibilityManager
+import android.view.accessibility.accessibilityManager
 import android.widget.RemoteViews
 import androidx.activity.result.ActivityResultLauncher
 import androidx.test.ext.junit.runners.AndroidJUnit4
@@ -42,7 +46,6 @@
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
 import com.android.systemui.communal.domain.interactor.communalInteractor
-import com.android.systemui.communal.domain.interactor.communalPrefsInteractor
 import com.android.systemui.communal.domain.interactor.communalSceneInteractor
 import com.android.systemui.communal.domain.interactor.communalSettingsInteractor
 import com.android.systemui.communal.domain.model.CommunalContentModel
@@ -61,8 +64,6 @@
 import com.android.systemui.settings.fakeUserTracker
 import com.android.systemui.testKosmos
 import com.android.systemui.user.data.repository.fakeUserRepository
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runTest
@@ -77,8 +78,12 @@
 import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -98,6 +103,7 @@
     private lateinit var mediaRepository: FakeCommunalMediaRepository
     private lateinit var communalSceneInteractor: CommunalSceneInteractor
     private lateinit var communalInteractor: CommunalInteractor
+    private lateinit var accessibilityManager: AccessibilityManager
 
     private val testableResources = context.orCreateTestableResources
 
@@ -119,6 +125,7 @@
             selectedUserIndex = 0,
         )
         kosmos.fakeFeatureFlagsClassic.set(Flags.COMMUNAL_SERVICE_ENABLED, true)
+        accessibilityManager = kosmos.accessibilityManager
 
         underTest =
             CommunalEditModeViewModel(
@@ -130,8 +137,10 @@
                 uiEventLogger,
                 logcatLogBuffer("CommunalEditModeViewModelTest"),
                 kosmos.testDispatcher,
-                kosmos.communalPrefsInteractor,
                 metricsLogger,
+                context,
+                accessibilityManager,
+                packageManager,
             )
     }
 
@@ -356,6 +365,37 @@
         verify(communalInteractor).setScrollPosition(eq(index), eq(offset))
     }
 
+    @Test
+    fun onNewWidgetAdded_accessibilityDisabled_doNothing() {
+        whenever(accessibilityManager.isEnabled).thenReturn(false)
+
+        val provider =
+            mock<AppWidgetProviderInfo> {
+                on { loadLabel(packageManager) }.thenReturn("Test Clock")
+            }
+        underTest.onNewWidgetAdded(provider)
+
+        verify(accessibilityManager, never()).sendAccessibilityEvent(any())
+    }
+
+    @Test
+    fun onNewWidgetAdded_accessibilityEnabled_sendAccessibilityAnnouncement() {
+        whenever(accessibilityManager.isEnabled).thenReturn(true)
+
+        val provider =
+            mock<AppWidgetProviderInfo> {
+                on { loadLabel(packageManager) }.thenReturn("Test Clock")
+            }
+        underTest.onNewWidgetAdded(provider)
+
+        val captor = argumentCaptor<AccessibilityEvent>()
+        verify(accessibilityManager).sendAccessibilityEvent(captor.capture())
+
+        val event = captor.firstValue
+        assertThat(event.eventType).isEqualTo(AccessibilityEvent.TYPE_ANNOUNCEMENT)
+        assertThat(event.contentDescription).isEqualTo("Test Clock widget added to lock screen")
+    }
+
     private companion object {
         val MAIN_USER_INFO = UserInfo(0, "primary", UserInfo.FLAG_MAIN)
         const val WIDGET_PICKER_PACKAGE_NAME = "widget_picker_package_name"
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
index 2bf50b3..91259a6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/data/repository/DeviceEntryFaceAuthRepositoryTest.kt
@@ -58,7 +58,6 @@
 import com.android.systemui.flags.FakeFeatureFlags
 import com.android.systemui.keyguard.data.repository.BiometricType
 import com.android.systemui.keyguard.data.repository.fakeBiometricSettingsRepository
-import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFingerprintAuthRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
@@ -78,7 +77,6 @@
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
 import com.android.systemui.power.domain.interactor.powerInteractor
-import com.android.systemui.statusbar.commandQueue
 import com.android.systemui.statusbar.phone.KeyguardBypassController
 import com.android.systemui.testKosmos
 import com.android.systemui.user.data.model.SelectionStatus
@@ -116,7 +114,7 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class DeviceEntryFaceAuthRepositoryTest : SysuiTestCase() {
-    private val kosmos = testKosmos().apply { this.commandQueue = this.fakeCommandQueue }
+    private val kosmos = testKosmos()
     private lateinit var underTest: DeviceEntryFaceAuthRepositoryImpl
 
     @Mock private lateinit var faceManager: FaceManager
@@ -162,7 +160,6 @@
     private val displayStateInteractor = kosmos.displayStateInteractor
     private val bouncerRepository = kosmos.fakeKeyguardBouncerRepository
     private val displayRepository = kosmos.displayRepository
-    private val fakeCommandQueue = kosmos.fakeCommandQueue
     private val keyguardTransitionInteractor = kosmos.keyguardTransitionInteractor
     private lateinit var featureFlags: FakeFeatureFlags
 
@@ -572,9 +569,7 @@
                 bouncerRepository.setAlternateVisible(false)
                 // Keyguard is occluded when secure camera is active.
                 keyguardRepository.setKeyguardOccluded(true)
-                fakeCommandQueue.doForEachCallback {
-                    it.onCameraLaunchGestureDetected(CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
-                }
+                keyguardInteractor.onCameraLaunchDetected(CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
             }
         }
 
@@ -589,9 +584,7 @@
             assertThat(canFaceAuthRun()).isTrue()
 
             // launch secure camera
-            fakeCommandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
-            }
+            keyguardInteractor.onCameraLaunchDetected(CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
             keyguardRepository.setKeyguardOccluded(true)
             runCurrent()
             assertThat(canFaceAuthRun()).isFalse()
@@ -870,9 +863,7 @@
                 bouncerRepository.setAlternateVisible(false)
                 // Keyguard is occluded when secure camera is active.
                 keyguardRepository.setKeyguardOccluded(true)
-                fakeCommandQueue.doForEachCallback {
-                    it.onCameraLaunchGestureDetected(CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
-                }
+                keyguardInteractor.onCameraLaunchDetected(CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
             }
         }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
index 4c24ce2..5c09777 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
@@ -62,12 +62,13 @@
 import com.android.systemui.dreams.complication.HideComplicationTouchHandler
 import com.android.systemui.dreams.dagger.DreamOverlayComponent
 import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
+import com.android.systemui.keyguard.gesture.domain.gestureInteractor
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.navigationbar.gestural.domain.GestureInteractor
 import com.android.systemui.testKosmos
 import com.android.systemui.touch.TouchInsetManager
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
@@ -84,9 +85,11 @@
 import org.mockito.Mockito
 import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.isNull
-import org.mockito.Mockito.spy
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.spy
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -166,6 +169,7 @@
     private lateinit var bouncerRepository: FakeKeyguardBouncerRepository
     private lateinit var communalRepository: FakeCommunalSceneRepository
     private var viewCaptureSpy = spy(ViewCaptureFactory.getInstance(context))
+    private lateinit var gestureInteractor: GestureInteractor
 
     @Captor var mViewCaptor: ArgumentCaptor<View>? = null
     private lateinit var mService: DreamOverlayService
@@ -177,6 +181,7 @@
         lifecycleRegistry = FakeLifecycleRegistry(mLifecycleOwner)
         bouncerRepository = kosmos.fakeKeyguardBouncerRepository
         communalRepository = kosmos.fakeCommunalSceneRepository
+        gestureInteractor = spy(kosmos.gestureInteractor)
 
         whenever(mDreamOverlayComponent.getDreamOverlayContainerViewController())
             .thenReturn(mDreamOverlayContainerViewController)
@@ -231,6 +236,7 @@
                 HOME_CONTROL_PANEL_DREAM_COMPONENT,
                 mDreamOverlayCallbackController,
                 kosmos.keyguardInteractor,
+                gestureInteractor,
                 WINDOW_NAME
             )
     }
@@ -955,6 +961,47 @@
         assertThat(lifecycleRegistry.currentState).isEqualTo(Lifecycle.State.RESUMED)
     }
 
+    @Test
+    fun testDreamActivityGesturesBlockedOnStart() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        val captor = argumentCaptor<ComponentName>()
+        verify(gestureInteractor)
+            .addGestureBlockedActivity(captor.capture(), eq(GestureInteractor.Scope.Global))
+        assertThat(captor.firstValue.packageName)
+            .isEqualTo(ComponentName.unflattenFromString(DREAM_COMPONENT)?.packageName)
+    }
+
+    @Test
+    fun testDreamActivityGesturesUnblockedOnEnd() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        client.endDream()
+        mMainExecutor.runAllReady()
+        val captor = argumentCaptor<ComponentName>()
+        verify(gestureInteractor)
+            .removeGestureBlockedActivity(captor.capture(), eq(GestureInteractor.Scope.Global))
+        assertThat(captor.firstValue.packageName)
+            .isEqualTo(ComponentName.unflattenFromString(DREAM_COMPONENT)?.packageName)
+    }
+
     internal class FakeLifecycleRegistry(provider: LifecycleOwner) : LifecycleRegistry(provider) {
         val mLifecycles: MutableList<State> = ArrayList()
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/data/GestureRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/data/GestureRepositoryTest.kt
new file mode 100644
index 0000000..91d37cf
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/data/GestureRepositoryTest.kt
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2024 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.gesture.data
+
+import android.content.ComponentName
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.navigationbar.gestural.data.respository.GestureRepositoryImpl
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class GestureRepositoryTest : SysuiTestCase() {
+    private val testDispatcher = StandardTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
+    private val underTest by lazy { GestureRepositoryImpl(testDispatcher) }
+
+    @Test
+    fun addRemoveComponentToBlock_updatesBlockedComponentSet() =
+        testScope.runTest {
+            val component = mock<ComponentName>()
+
+            underTest.addGestureBlockedActivity(component)
+            val addedBlockedComponents by collectLastValue(underTest.gestureBlockedActivities)
+            assertThat(addedBlockedComponents).contains(component)
+
+            underTest.removeGestureBlockedActivity(component)
+            val removedBlockedComponents by collectLastValue(underTest.gestureBlockedActivities)
+            assertThat(removedBlockedComponents).isEmpty()
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/GestureInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/GestureInteractorTest.kt
new file mode 100644
index 0000000..bc142e6
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/GestureInteractorTest.kt
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2024 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.gesture.domain
+
+import android.content.ComponentName
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.navigationbar.gestural.data.respository.GestureRepository
+import com.android.systemui.navigationbar.gestural.domain.GestureInteractor
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.junit.MockitoJUnit
+import org.mockito.junit.MockitoRule
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class GestureInteractorTest : SysuiTestCase() {
+    @Rule @JvmField val mockitoRule: MockitoRule = MockitoJUnit.rule()
+
+    val dispatcher = StandardTestDispatcher()
+    val testScope = TestScope(dispatcher)
+
+    @Mock private lateinit var gestureRepository: GestureRepository
+
+    private val underTest by lazy {
+        GestureInteractor(gestureRepository, testScope.backgroundScope)
+    }
+
+    @Before
+    fun setup() {
+        Dispatchers.setMain(dispatcher)
+        whenever(gestureRepository.gestureBlockedActivities).thenReturn(MutableStateFlow(setOf()))
+    }
+
+    @After
+    fun tearDown() {
+        Dispatchers.resetMain()
+    }
+
+    @Test
+    fun addBlockedActivity_testCombination() =
+        testScope.runTest {
+            val globalComponent = mock<ComponentName>()
+            whenever(gestureRepository.gestureBlockedActivities)
+                .thenReturn(MutableStateFlow(setOf(globalComponent)))
+            val localComponent = mock<ComponentName>()
+            underTest.addGestureBlockedActivity(localComponent, GestureInteractor.Scope.Local)
+            val lastSeen by collectLastValue(underTest.gestureBlockedActivities)
+            testScope.runCurrent()
+            verify(gestureRepository, never()).addGestureBlockedActivity(any())
+            assertThat(lastSeen).hasSize(2)
+            assertThat(lastSeen).containsExactly(globalComponent, localComponent)
+        }
+
+    @Test
+    fun addBlockedActivityLocally_onlyAffectsLocalInteractor() =
+        testScope.runTest {
+            val component = mock<ComponentName>()
+            underTest.addGestureBlockedActivity(component, GestureInteractor.Scope.Local)
+            val lastSeen by collectLastValue(underTest.gestureBlockedActivities)
+            testScope.runCurrent()
+            verify(gestureRepository, never()).addGestureBlockedActivity(any())
+            assertThat(lastSeen).contains(component)
+        }
+
+    @Test
+    fun removeBlockedActivityLocally_onlyAffectsLocalInteractor() =
+        testScope.runTest {
+            val component = mock<ComponentName>()
+            underTest.addGestureBlockedActivity(component, GestureInteractor.Scope.Local)
+            val lastSeen by collectLastValue(underTest.gestureBlockedActivities)
+            testScope.runCurrent()
+            underTest.removeGestureBlockedActivity(component, GestureInteractor.Scope.Local)
+            testScope.runCurrent()
+            verify(gestureRepository, never()).removeGestureBlockedActivity(any())
+            assertThat(lastSeen).isEmpty()
+        }
+
+    @Test
+    fun addBlockedActivity_invokesRepository() =
+        testScope.runTest {
+            val component = mock<ComponentName>()
+            underTest.addGestureBlockedActivity(component, GestureInteractor.Scope.Global)
+            runCurrent()
+            val captor = argumentCaptor<ComponentName>()
+            verify(gestureRepository).addGestureBlockedActivity(captor.capture())
+            assertThat(captor.firstValue).isEqualTo(component)
+        }
+
+    @Test
+    fun removeBlockedActivity_invokesRepository() =
+        testScope.runTest {
+            val component = mock<ComponentName>()
+            underTest.removeGestureBlockedActivity(component, GestureInteractor.Scope.Global)
+            runCurrent()
+            val captor = argumentCaptor<ComponentName>()
+            verify(gestureRepository).removeGestureBlockedActivity(captor.capture())
+            assertThat(captor.firstValue).isEqualTo(component)
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt
index 18839e6..273e3cb 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt
@@ -23,9 +23,9 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.ActivityTransitionAnimator
-import com.android.systemui.classifier.falsingManager
 import com.android.systemui.haptics.vibratorHelper
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.log.core.FakeLogBuffer
 import com.android.systemui.qs.qsTileFactory
 import com.android.systemui.statusbar.policy.keyguardStateController
 import com.android.systemui.testKosmos
@@ -73,7 +73,7 @@
             QSLongPressEffect(
                 vibratorHelper,
                 kosmos.keyguardStateController,
-                kosmos.falsingManager,
+                FakeLogBuffer.Factory.create(),
             )
         longPressEffect.callback = callback
         longPressEffect.qsTile = qsTile
@@ -304,13 +304,12 @@
     }
 
     @Test
-    fun getStateForClick_withFalseTapWhenLocked_returnsIdle() {
+    fun getStateForClick_whenKeyguardsIsShowing_returnsIdle() {
         // GIVEN an active tile
         qsTile.state?.state = Tile.STATE_ACTIVE
 
-        // GIVEN that the device is locked and a false tap is detected
-        whenever(kosmos.keyguardStateController.isUnlocked).thenReturn(false)
-        kosmos.falsingManager.setFalseTap(true)
+        // GIVEN that the keyguard is showing
+        whenever(kosmos.keyguardStateController.isShowing).thenReturn(true)
 
         // WHEN determining the state of a click action
         val clickState = longPressEffect.getStateForClick()
@@ -324,9 +323,8 @@
         // GIVEN an active tile
         qsTile.state?.state = Tile.STATE_ACTIVE
 
-        // GIVEN that the device is locked and a false tap is not detected
-        whenever(kosmos.keyguardStateController.isUnlocked).thenReturn(false)
-        kosmos.falsingManager.setFalseTap(false)
+        // GIVEN that the keyguard is not showing
+        whenever(kosmos.keyguardStateController.isShowing).thenReturn(false)
 
         // WHEN determining the state of a click action
         val clickState = longPressEffect.getStateForClick()
@@ -340,9 +338,8 @@
         // GIVEN that the tile is null
         longPressEffect.qsTile = null
 
-        // GIVEN that the device is locked and a false tap is not detected
-        whenever(kosmos.keyguardStateController.isUnlocked).thenReturn(false)
-        kosmos.falsingManager.setFalseTap(false)
+        // GIVEN that the keyguard is not showing
+        whenever(kosmos.keyguardStateController.isShowing).thenReturn(false)
 
         // WHEN determining the state of a click action
         val clickState = longPressEffect.getStateForClick()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt
index 5115f5a..3b8ffcd 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt
@@ -173,6 +173,7 @@
         }
 
     @Test
+    @DisableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun transitionToGone_whenOpeningGlanceableHubEditMode() =
         testScope.runTest {
             kosmos.fakeKeyguardBouncerRepository.setAlternateVisible(true)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
index 7906a82..fc827a14 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
@@ -32,7 +32,7 @@
 import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
-import com.android.systemui.keyguard.shared.model.CameraLaunchSourceModel
+import com.android.systemui.keyguard.shared.model.CameraLaunchType
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.TransitionState
@@ -47,7 +47,6 @@
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -87,31 +86,17 @@
             val cameraLaunchSource = collectLastValue(flow)
             runCurrent()
 
-            commandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_WIGGLE)
-            }
-            assertThat(cameraLaunchSource()).isEqualTo(CameraLaunchSourceModel.WIGGLE)
+            underTest.onCameraLaunchDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
+            assertThat(cameraLaunchSource()!!.type).isEqualTo(CameraLaunchType.POWER_DOUBLE_TAP)
 
-            commandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(
-                    StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP
-                )
-            }
-            assertThat(cameraLaunchSource()).isEqualTo(CameraLaunchSourceModel.POWER_DOUBLE_TAP)
+            underTest.onCameraLaunchDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_WIGGLE)
+            assertThat(cameraLaunchSource()!!.type).isEqualTo(CameraLaunchType.WIGGLE)
 
-            commandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_LIFT_TRIGGER)
-            }
-            assertThat(cameraLaunchSource()).isEqualTo(CameraLaunchSourceModel.LIFT_TRIGGER)
+            underTest.onCameraLaunchDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_LIFT_TRIGGER)
+            assertThat(cameraLaunchSource()!!.type).isEqualTo(CameraLaunchType.LIFT_TRIGGER)
 
-            commandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(
-                    StatusBarManager.CAMERA_LAUNCH_SOURCE_QUICK_AFFORDANCE
-                )
-            }
-            assertThat(cameraLaunchSource()).isEqualTo(CameraLaunchSourceModel.QUICK_AFFORDANCE)
-
-            flow.onCompletion { assertThat(commandQueue.callbackCount()).isEqualTo(0) }
+            underTest.onCameraLaunchDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_QUICK_AFFORDANCE)
+            assertThat(cameraLaunchSource()!!.type).isEqualTo(CameraLaunchType.QUICK_AFFORDANCE)
         }
 
     @Test
@@ -121,11 +106,7 @@
             val secureCameraActive = collectLastValue(underTest.isSecureCameraActive)
             runCurrent()
 
-            commandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(
-                    StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP
-                )
-            }
+            underTest.onCameraLaunchDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
 
             assertThat(secureCameraActive()).isTrue()
 
@@ -146,11 +127,7 @@
             val secureCameraActive = collectLastValue(underTest.isSecureCameraActive)
             runCurrent()
 
-            commandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(
-                    StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP
-                )
-            }
+            underTest.onCameraLaunchDetected(StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
             assertThat(secureCameraActive()).isTrue()
 
             // Keyguard is showing and not occluded
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
index 2d77f4f..75c0d3b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
@@ -22,6 +22,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.internal.widget.LockPatternUtils
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.common.shared.model.ContentDescription
@@ -86,7 +87,8 @@
     @Mock private lateinit var launchAnimator: DialogTransitionAnimator
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
     @Mock private lateinit var shadeInteractor: ShadeInteractor
-    @Mock private lateinit var logger: KeyguardQuickAffordancesMetricsLogger
+    @Mock private lateinit var logger: KeyguardQuickAffordancesLogger
+    @Mock private lateinit var metricsLogger: KeyguardQuickAffordancesMetricsLogger
 
     private val kosmos = testKosmos()
 
@@ -194,6 +196,7 @@
                 repository = { quickAffordanceRepository },
                 launchAnimator = launchAnimator,
                 logger = logger,
+                metricsLogger = metricsLogger,
                 devicePolicyManager = devicePolicyManager,
                 dockManager = dockManager,
                 biometricSettingsRepository = biometricSettingsRepository,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
index d9708a4..9762fd8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionScenariosTest.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.keyguard.domain.interactor
 
-import android.app.StatusBarManager
+import android.app.StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.FlagsParameterization
@@ -42,7 +42,6 @@
 import com.android.systemui.flags.andSceneContainer
 import com.android.systemui.flags.fakeFeatureFlagsClassic
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
-import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
@@ -59,7 +58,6 @@
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.shade.shadeTestUtil
-import com.android.systemui.statusbar.commandQueue
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.whenever
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -93,13 +91,12 @@
     private val kosmos =
         testKosmos().apply {
             fakeKeyguardTransitionRepository = spy(FakeKeyguardTransitionRepository())
-            this.commandQueue = fakeCommandQueue
         }
     private val testScope = kosmos.testScope
 
     private val keyguardRepository by lazy { kosmos.fakeKeyguardRepository }
+    private val keyguardInteractor by lazy { kosmos.keyguardInteractor }
     private val bouncerRepository by lazy { kosmos.fakeKeyguardBouncerRepository }
-    private var commandQueue = kosmos.fakeCommandQueue
     private val shadeTestUtil by lazy { kosmos.shadeTestUtil }
     private val transitionRepository by lazy { kosmos.fakeKeyguardTransitionRepository }
     private lateinit var featureFlags: FakeFeatureFlags
@@ -1724,11 +1721,7 @@
             reset(transitionRepository)
 
             // ...AND WHEN the camera gesture is detected quickly afterwards
-            commandQueue.doForEachCallback {
-                it.onCameraLaunchGestureDetected(
-                    StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP
-                )
-            }
+            keyguardInteractor.onCameraLaunchDetected(CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP)
             runCurrent()
 
             // THEN a transition from DOZING => OCCLUDED should occur
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModelTest.kt
new file mode 100644
index 0000000..f74b74a
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModelTest.kt
@@ -0,0 +1,125 @@
+/*
+ * Copyright (C) 2024 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.keyguard.ui.viewmodel
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.biometrics.data.repository.fingerprintPropertyRepository
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.coroutines.collectValues
+import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.TransitionState
+import com.android.systemui.keyguard.shared.model.TransitionState.RUNNING
+import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@ExperimentalCoroutinesApi
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class AlternateBouncerToLockscreenTransitionViewModelTest : SysuiTestCase() {
+    val kosmos = testKosmos()
+    val testScope = kosmos.testScope
+
+    val keyguardTransitionRepository = kosmos.fakeKeyguardTransitionRepository
+    val fingerprintPropertyRepository = kosmos.fingerprintPropertyRepository
+
+    val underTest = kosmos.alternateBouncerToLockscreenTransitionViewModel
+
+    @Test
+    fun lockscreenAlpha_zeroInitialAlpha() =
+        testScope.runTest {
+            // ViewState starts at 0 alpha.
+            val viewState = ViewStateAccessor(alpha = { 0f })
+            val alpha by collectValues(underTest.lockscreenAlpha(viewState))
+
+            keyguardTransitionRepository.sendTransitionSteps(
+                from = KeyguardState.ALTERNATE_BOUNCER,
+                to = KeyguardState.LOCKSCREEN,
+                testScope
+            )
+
+            assertThat(alpha[0]).isEqualTo(0f)
+            // alpha duration is 250ms of the 300ms total, so 0.5f of the total is 0.6
+            assertThat(alpha[1]).isEqualTo(0.6f)
+            assertThat(alpha[2]).isEqualTo(1f)
+        }
+
+    @Test
+    fun deviceEntryParentViewAlpha() =
+        testScope.runTest {
+            val deviceEntryParentViewAlpha by collectLastValue(underTest.deviceEntryParentViewAlpha)
+
+            // immediately 1f
+            keyguardTransitionRepository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            assertThat(deviceEntryParentViewAlpha).isEqualTo(1f)
+
+            keyguardTransitionRepository.sendTransitionStep(step(0.4f))
+            assertThat(deviceEntryParentViewAlpha).isEqualTo(1f)
+
+            keyguardTransitionRepository.sendTransitionStep(step(.85f))
+            assertThat(deviceEntryParentViewAlpha).isEqualTo(1f)
+
+            keyguardTransitionRepository.sendTransitionStep(step(1f))
+            assertThat(deviceEntryParentViewAlpha).isEqualTo(1f)
+        }
+
+    @Test
+    fun deviceEntryBackgroundViewAlpha_udfpsEnrolled_show() =
+        testScope.runTest {
+            fingerprintPropertyRepository.supportsUdfps()
+            val bgViewAlpha by collectLastValue(underTest.deviceEntryBackgroundViewAlpha)
+            runCurrent()
+
+            // immediately 1f
+            keyguardTransitionRepository.sendTransitionStep(step(0f, TransitionState.STARTED))
+            assertThat(bgViewAlpha).isEqualTo(1f)
+
+            keyguardTransitionRepository.sendTransitionStep(step(0.1f))
+            assertThat(bgViewAlpha).isEqualTo(1f)
+
+            keyguardTransitionRepository.sendTransitionStep(step(.3f))
+            assertThat(bgViewAlpha).isEqualTo(1f)
+
+            keyguardTransitionRepository.sendTransitionStep(step(.5f))
+            assertThat(bgViewAlpha).isEqualTo(1f)
+
+            keyguardTransitionRepository.sendTransitionStep(step(1f, TransitionState.FINISHED))
+            assertThat(bgViewAlpha).isEqualTo(1f)
+        }
+
+    private fun step(
+        value: Float,
+        state: TransitionState = TransitionState.RUNNING
+    ): TransitionStep {
+        return TransitionStep(
+            from = KeyguardState.ALTERNATE_BOUNCER,
+            to = KeyguardState.LOCKSCREEN,
+            value = value,
+            transitionState = state,
+            ownerName = "AlternateBouncerToLockscreenTransitionViewModelTest"
+        )
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModelTest.kt
index aba21c9..cd0a11c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModelTest.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.content.res.Configuration
+import android.util.LayoutDirection
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -33,6 +35,8 @@
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -73,6 +77,9 @@
                 R.dimen.dreaming_to_hub_transition_dream_overlay_translation_x,
                 -100
             )
+            val configuration: Configuration = mock()
+            whenever(configuration.layoutDirection).thenReturn(LayoutDirection.LTR)
+            configurationRepository.onConfigurationChange(configuration)
 
             val values by collectValues(underTest.dreamOverlayTranslationX)
             assertThat(values).isEmpty()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModelTest.kt
index 11890c7..69361ef 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModelTest.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.content.res.Configuration
+import android.util.LayoutDirection
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -33,6 +35,8 @@
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -69,6 +73,9 @@
     @Test
     fun dreamOverlayTranslationX() =
         testScope.runTest {
+            val config: Configuration = mock()
+            whenever(config.layoutDirection).thenReturn(LayoutDirection.LTR)
+            configurationRepository.onConfigurationChange(config)
             configurationRepository.setDimensionPixelSize(
                 R.dimen.hub_to_dreaming_transition_dream_overlay_translation_x,
                 100
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModelTest.kt
index 1aa1ec4..d2be649 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModelTest.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.content.res.Configuration
+import android.util.LayoutDirection
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -34,6 +36,8 @@
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -77,6 +81,10 @@
     @Test
     fun lockscreenTranslationX() =
         testScope.runTest {
+            val config: Configuration = mock()
+            whenever(config.layoutDirection).thenReturn(LayoutDirection.LTR)
+            configurationRepository.onConfigurationChange(config)
+
             configurationRepository.setDimensionPixelSize(
                 R.dimen.hub_to_lockscreen_transition_lockscreen_translation_x,
                 100
@@ -102,6 +110,10 @@
     @Test
     fun lockscreenTranslationX_resetsAfterCancellation() =
         testScope.runTest {
+            val config: Configuration = mock()
+            whenever(config.layoutDirection).thenReturn(LayoutDirection.LTR)
+            configurationRepository.onConfigurationChange(config)
+
             configurationRepository.setDimensionPixelSize(
                 R.dimen.hub_to_lockscreen_transition_lockscreen_translation_x,
                 100
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModelTest.kt
index 68a7b7e..a60a486 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModelTest.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.content.res.Configuration
+import android.util.LayoutDirection
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -34,6 +36,9 @@
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mockito.mock
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -82,6 +87,9 @@
                 R.dimen.lockscreen_to_hub_transition_lockscreen_translation_x,
                 -100
             )
+            val configuration = mock<Configuration>()
+            whenever(configuration.layoutDirection).thenReturn(LayoutDirection.LTR)
+            configurationRepository.onConfigurationChange(configuration)
             val values by collectValues(underTest.keyguardTranslationX)
             assertThat(values).isEmpty()
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/power/domain/interactor/PowerInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/power/domain/interactor/PowerInteractorTest.kt
similarity index 76%
rename from packages/SystemUI/tests/src/com/android/systemui/power/domain/interactor/PowerInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/power/domain/interactor/PowerInteractorTest.kt
index 12c9eb9..53dec69 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/power/domain/interactor/PowerInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/power/domain/interactor/PowerInteractorTest.kt
@@ -21,22 +21,23 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.camera.cameraGestureHelper
 import com.android.systemui.classifier.FalsingCollector
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
-import com.android.systemui.power.shared.model.WakeSleepReason
-import com.android.systemui.power.shared.model.WakefulnessState
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.power.data.repository.FakePowerRepository
+import com.android.systemui.power.shared.model.WakeSleepReason
+import com.android.systemui.power.shared.model.WakefulnessState
 import com.android.systemui.statusbar.phone.ScreenOffAnimationController
+import com.android.systemui.testKosmos
+import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertTrue
-import kotlin.test.assertEquals
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -47,6 +48,9 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class PowerInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val cameraGestureHelper = kosmos.cameraGestureHelper
 
     private lateinit var underTest: PowerInteractor
     private lateinit var repository: FakePowerRepository
@@ -66,33 +70,30 @@
                 falsingCollector,
                 screenOffAnimationController,
                 statusBarStateController,
+                { cameraGestureHelper },
             )
+
+        whenever(cameraGestureHelper.canCameraGestureBeLaunched(any())).thenReturn(true)
     }
 
     @Test
     fun isInteractive_screenTurnsOff() =
-        runBlocking(IMMEDIATE) {
+        testScope.runTest {
             repository.setInteractive(true)
-            var value: Boolean? = null
-            val job = underTest.isInteractive.onEach { value = it }.launchIn(this)
+            val isInteractive by collectLastValue(underTest.isInteractive)
 
             repository.setInteractive(false)
-
-            assertThat(value).isFalse()
-            job.cancel()
+            assertThat(isInteractive).isFalse()
         }
 
     @Test
     fun isInteractive_becomesInteractive() =
-        runBlocking(IMMEDIATE) {
+        testScope.runTest {
             repository.setInteractive(false)
-            var value: Boolean? = null
-            val job = underTest.isInteractive.onEach { value = it }.launchIn(this)
+            val isInteractive by collectLastValue(underTest.isInteractive)
 
             repository.setInteractive(true)
-
-            assertThat(value).isTrue()
-            job.cancel()
+            assertThat(isInteractive).isTrue()
         }
 
     @Test
@@ -203,6 +204,23 @@
     }
 
     @Test
+    fun onCameraLaunchGestureDetected_isNotTrueWhenCannotLaunch() {
+        whenever(cameraGestureHelper.canCameraGestureBeLaunched(any())).thenReturn(false)
+        underTest.onStartedWakingUp(
+            PowerManager.WAKE_REASON_POWER_BUTTON,
+            /*powerButtonLaunchGestureTriggeredDuringSleep= */ false
+        )
+        underTest.onFinishedWakingUp()
+        underTest.onCameraLaunchGestureDetected()
+
+        assertThat(repository.wakefulness.value.internalWakefulnessState)
+            .isEqualTo(WakefulnessState.AWAKE)
+        assertThat(repository.wakefulness.value.lastWakeReason)
+            .isEqualTo(WakeSleepReason.POWER_BUTTON)
+        assertFalse(repository.wakefulness.value.powerButtonLaunchGestureTriggered)
+    }
+
+    @Test
     fun onCameraLaunchGestureDetected_maintainsAllOtherState() {
         underTest.onStartedWakingUp(
             PowerManager.WAKE_REASON_POWER_BUTTON,
@@ -211,8 +229,10 @@
         underTest.onFinishedWakingUp()
         underTest.onCameraLaunchGestureDetected()
 
-        assertEquals(WakefulnessState.AWAKE, repository.wakefulness.value.internalWakefulnessState)
-        assertEquals(WakeSleepReason.POWER_BUTTON, repository.wakefulness.value.lastWakeReason)
+        assertThat(repository.wakefulness.value.internalWakefulnessState)
+            .isEqualTo(WakefulnessState.AWAKE)
+        assertThat(repository.wakefulness.value.lastWakeReason)
+            .isEqualTo(WakeSleepReason.POWER_BUTTON)
         assertTrue(repository.wakefulness.value.powerButtonLaunchGestureTriggered)
     }
 
@@ -221,21 +241,23 @@
         underTest.onCameraLaunchGestureDetected()
         // Ensure that the 'false' here does not clear the direct launch detection call earlier.
         // This state should only be reset onStartedGoingToSleep.
-        underTest.onFinishedGoingToSleep(/*powerButtonLaunchGestureTriggeredDuringSleep= */ false)
+        underTest.onFinishedGoingToSleep(/* powerButtonLaunchGestureTriggeredDuringSleep= */ false)
         underTest.onStartedWakingUp(
             PowerManager.WAKE_REASON_POWER_BUTTON,
             /*powerButtonLaunchGestureTriggeredDuringSleep= */ false
         )
         underTest.onFinishedWakingUp()
 
-        assertEquals(WakefulnessState.AWAKE, repository.wakefulness.value.internalWakefulnessState)
-        assertEquals(WakeSleepReason.POWER_BUTTON, repository.wakefulness.value.lastWakeReason)
+        assertThat(repository.wakefulness.value.internalWakefulnessState)
+            .isEqualTo(WakefulnessState.AWAKE)
+        assertThat(repository.wakefulness.value.lastWakeReason)
+            .isEqualTo(WakeSleepReason.POWER_BUTTON)
         assertTrue(repository.wakefulness.value.powerButtonLaunchGestureTriggered)
     }
 
     @Test
     fun cameraLaunchDetectedOnGoingToSleep_stillTrue_ifGestureNotDetectedOnWakingUp() {
-        underTest.onFinishedGoingToSleep(/*powerButtonLaunchGestureTriggeredDuringSleep= */ true)
+        underTest.onFinishedGoingToSleep(/* powerButtonLaunchGestureTriggeredDuringSleep= */ true)
         // Ensure that the 'false' here does not clear the direct launch detection call earlier.
         // This state should only be reset onStartedGoingToSleep.
         underTest.onStartedWakingUp(
@@ -244,12 +266,10 @@
         )
         underTest.onFinishedWakingUp()
 
-        assertEquals(WakefulnessState.AWAKE, repository.wakefulness.value.internalWakefulnessState)
-        assertEquals(WakeSleepReason.POWER_BUTTON, repository.wakefulness.value.lastWakeReason)
+        assertThat(repository.wakefulness.value.internalWakefulnessState)
+            .isEqualTo(WakefulnessState.AWAKE)
+        assertThat(repository.wakefulness.value.lastWakeReason)
+            .isEqualTo(WakeSleepReason.POWER_BUTTON)
         assertTrue(repository.wakefulness.value.powerButtonLaunchGestureTriggered)
     }
-
-    companion object {
-        private val IMMEDIATE = Dispatchers.Main.immediate
-    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/GridConsistencyInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/GridConsistencyInteractorTest.kt
index 56b3679..42db96e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/GridConsistencyInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/GridConsistencyInteractorTest.kt
@@ -23,8 +23,8 @@
 import com.android.systemui.qs.panels.data.repository.DefaultLargeTilesRepository
 import com.android.systemui.qs.panels.data.repository.defaultLargeTilesRepository
 import com.android.systemui.qs.panels.data.repository.gridLayoutTypeRepository
+import com.android.systemui.qs.panels.shared.model.GridLayoutType
 import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
-import com.android.systemui.qs.panels.shared.model.PartitionedGridLayoutType
 import com.android.systemui.qs.pipeline.data.repository.tileSpecRepository
 import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
 import com.android.systemui.qs.pipeline.shared.TileSpec
@@ -42,6 +42,8 @@
 @RunWith(AndroidJUnit4::class)
 class GridConsistencyInteractorTest : SysuiTestCase() {
 
+    data object NoopGridLayoutType : GridLayoutType
+
     private val kosmos =
         testKosmos().apply {
             defaultLargeTilesRepository =
@@ -54,6 +56,11 @@
                             TileSpec.create("largeD"),
                         )
                 }
+            gridConsistencyInteractorsMap =
+                mapOf(
+                    Pair(NoopGridLayoutType, noopGridConsistencyInteractor),
+                    Pair(InfiniteGridLayoutType, infiniteGridConsistencyInteractor)
+                )
         }
 
     private val underTest = with(kosmos) { gridConsistencyInteractor }
@@ -71,7 +78,7 @@
         with(kosmos) {
             testScope.runTest {
                 // Using the no-op grid consistency interactor
-                gridLayoutTypeRepository.setLayout(PartitionedGridLayoutType)
+                gridLayoutTypeRepository.setLayout(NoopGridLayoutType)
 
                 // Setting an invalid layout with holes
                 // [ Large A ] [ sa ]
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/compose/PartitionedGridLayoutTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/compose/PartitionedGridLayoutTest.kt
deleted file mode 100644
index 55b7454..0000000
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/compose/PartitionedGridLayoutTest.kt
+++ /dev/null
@@ -1,123 +0,0 @@
-/*
- * Copyright (C) 2024 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.qs.panels.ui.compose
-
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.kosmos.testScope
-import com.android.systemui.qs.panels.data.repository.DefaultLargeTilesRepository
-import com.android.systemui.qs.panels.data.repository.defaultLargeTilesRepository
-import com.android.systemui.qs.panels.ui.viewmodel.MockTileViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.iconTilesViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
-import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.android.systemui.testKosmos
-import com.google.common.truth.Truth
-import kotlinx.coroutines.test.runTest
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@SmallTest
-@RunWith(AndroidJUnit4::class)
-class PartitionedGridLayoutTest : SysuiTestCase() {
-    private val kosmos =
-        testKosmos().apply {
-            defaultLargeTilesRepository =
-                object : DefaultLargeTilesRepository {
-                    override val defaultLargeTiles: Set<TileSpec> = setOf(TileSpec.create("large"))
-                }
-        }
-
-    private val underTest = with(kosmos) { PartitionedGridLayout(partitionedGridViewModel) }
-
-    @Test
-    fun correctPagination_underOnePage_partitioned_sameRelativeOrder() =
-        with(kosmos) {
-            testScope.runTest {
-                val rows = 3
-                val columns = 4
-
-                val tiles =
-                    listOf(
-                        largeTile(),
-                        smallTile(),
-                        smallTile(),
-                        largeTile(),
-                        largeTile(),
-                        smallTile()
-                    )
-                val (smallTiles, largeTiles) =
-                    tiles.partition { iconTilesViewModel.isIconTile(it.spec) }
-
-                // [L L] [L L]
-                // [L L]
-                // [S] [S] [S]
-
-                val pages = underTest.splitIntoPages(tiles, rows = rows, columns = columns)
-
-                Truth.assertThat(pages).hasSize(1)
-                Truth.assertThat(pages[0]).isEqualTo(largeTiles + smallTiles)
-            }
-        }
-
-    @Test
-    fun correctPagination_twoPages_partitioned_sameRelativeOrder() =
-        with(kosmos) {
-            testScope.runTest {
-                val rows = 3
-                val columns = 4
-
-                val tiles =
-                    listOf(
-                        largeTile(),
-                        smallTile(),
-                        smallTile(),
-                        largeTile(),
-                        smallTile(),
-                        smallTile(),
-                        largeTile(),
-                        smallTile(),
-                        smallTile(),
-                    )
-                // --- Page 1 ---
-                // [L L] [L L]
-                // [L L]
-                // [S] [S] [S] [S]
-                // --- Page 2 ---
-                // [S] [S]
-
-                val (smallTiles, largeTiles) =
-                    tiles.partition { iconTilesViewModel.isIconTile(it.spec) }
-
-                val pages = underTest.splitIntoPages(tiles, rows = rows, columns = columns)
-
-                val expectedPage0 = largeTiles + smallTiles.take(4)
-                val expectedPage1 = smallTiles.drop(4)
-
-                Truth.assertThat(pages).hasSize(2)
-                Truth.assertThat(pages[0]).isEqualTo(expectedPage0)
-                Truth.assertThat(pages[1]).isEqualTo(expectedPage1)
-            }
-        }
-
-    companion object {
-        fun largeTile() = MockTileViewModel(TileSpec.create("large"))
-
-        fun smallTile() = MockTileViewModel(TileSpec.create("small"))
-    }
-}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
index 09dca25..69b8ee1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
@@ -22,11 +22,14 @@
 import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.settingslib.notification.data.repository.FakeZenModeRepository
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectValues
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.qs.tiles.base.interactor.DataUpdateTrigger
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
+import com.android.systemui.statusbar.policy.data.repository.fakeZenModeRepository
+import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.flowOf
@@ -40,51 +43,72 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class ModesTileDataInteractorTest : SysuiTestCase() {
-    private val zenModeRepository = FakeZenModeRepository()
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val dispatcher = kosmos.testDispatcher
+    private val zenModeRepository = kosmos.fakeZenModeRepository
 
-    private val underTest = ModesTileDataInteractor(zenModeRepository)
+    private val underTest = ModesTileDataInteractor(zenModeRepository, dispatcher)
 
     @EnableFlags(Flags.FLAG_MODES_UI)
     @Test
-    fun availableWhenFlagIsOn() = runTest {
-        val availability = underTest.availability(TEST_USER).toCollection(mutableListOf())
+    fun availableWhenFlagIsOn() =
+        testScope.runTest {
+            val availability = underTest.availability(TEST_USER).toCollection(mutableListOf())
 
-        assertThat(availability).containsExactly(true)
-    }
+            assertThat(availability).containsExactly(true)
+        }
 
     @DisableFlags(Flags.FLAG_MODES_UI)
     @Test
-    fun unavailableWhenFlagIsOff() = runTest {
-        val availability = underTest.availability(TEST_USER).toCollection(mutableListOf())
+    fun unavailableWhenFlagIsOff() =
+        testScope.runTest {
+            val availability = underTest.availability(TEST_USER).toCollection(mutableListOf())
 
-        assertThat(availability).containsExactly(false)
-    }
+            assertThat(availability).containsExactly(false)
+        }
 
     @EnableFlags(Flags.FLAG_MODES_UI)
     @Test
-    fun isActivatedWhenModesChange() = runTest {
-        val dataList: List<ModesTileModel> by
-            collectValues(underTest.tileData(TEST_USER, flowOf(DataUpdateTrigger.InitialRequest)))
-        runCurrent()
-        assertThat(dataList.map { it.isActivated }).containsExactly(false).inOrder()
+    fun isActivatedWhenModesChange() =
+        testScope.runTest {
+            val dataList: List<ModesTileModel> by
+                collectValues(
+                    underTest.tileData(TEST_USER, flowOf(DataUpdateTrigger.InitialRequest))
+                )
+            runCurrent()
+            assertThat(dataList.map { it.isActivated }).containsExactly(false).inOrder()
 
-        // Add active mode
-        zenModeRepository.addMode(id = "One", active = true)
-        runCurrent()
-        assertThat(dataList.map { it.isActivated }).containsExactly(false, true).inOrder()
+            // Add active mode
+            zenModeRepository.addMode(id = "One", active = true)
+            runCurrent()
+            assertThat(dataList.map { it.isActivated }).containsExactly(false, true).inOrder()
+            assertThat(dataList.map { it.activeModes }.last()).containsExactly("Mode One")
 
-        // Add another mode: state hasn't changed, so this shouldn't cause another emission
-        zenModeRepository.addMode(id = "Two", active = true)
-        runCurrent()
-        assertThat(dataList.map { it.isActivated }).containsExactly(false, true).inOrder()
+            // Add an inactive mode: state hasn't changed, so this shouldn't cause another emission
+            zenModeRepository.addMode(id = "Two", active = false)
+            runCurrent()
+            assertThat(dataList.map { it.isActivated }).containsExactly(false, true).inOrder()
+            assertThat(dataList.map { it.activeModes }.last()).containsExactly("Mode One")
 
-        // Remove a mode and disable the other
-        zenModeRepository.removeMode("One")
-        runCurrent()
-        zenModeRepository.deactivateMode("Two")
-        runCurrent()
-        assertThat(dataList.map { it.isActivated }).containsExactly(false, true, false).inOrder()
-    }
+            // Add another active mode
+            zenModeRepository.addMode(id = "Three", active = true)
+            runCurrent()
+            assertThat(dataList.map { it.isActivated }).containsExactly(false, true, true).inOrder()
+            assertThat(dataList.map { it.activeModes }.last())
+                .containsExactly("Mode One", "Mode Three")
+                .inOrder()
+
+            // Remove a mode and deactivate the other
+            zenModeRepository.removeMode("One")
+            runCurrent()
+            zenModeRepository.deactivateMode("Three")
+            runCurrent()
+            assertThat(dataList.map { it.isActivated })
+                .containsExactly(false, true, true, true, false)
+                .inOrder()
+            assertThat(dataList.map { it.activeModes }.last()).isEmpty()
+        }
 
     private companion object {
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt
index d5c9102..4b75649 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt
@@ -21,64 +21,76 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.animation.DialogTransitionAnimator
-import com.android.systemui.qs.tiles.base.actions.FakeQSTileIntentUserInputHandler
+import com.android.systemui.animation.Expandable
 import com.android.systemui.qs.tiles.base.actions.QSTileIntentUserInputHandlerSubject
+import com.android.systemui.qs.tiles.base.actions.qsTileIntentUserInputHandler
 import com.android.systemui.qs.tiles.base.interactor.QSTileInputTestKtx
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
-import com.android.systemui.statusbar.phone.SystemUIDialog
-import com.android.systemui.statusbar.policy.ui.dialog.ModesDialogDelegate
-import com.google.common.truth.Truth
-import kotlin.coroutines.EmptyCoroutineContext
+import com.android.systemui.statusbar.policy.ui.dialog.mockModesDialogDelegate
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
-import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 @EnableFlags(android.app.Flags.FLAG_MODES_UI)
 class ModesTileUserActionInteractorTest : SysuiTestCase() {
-    private val inputHandler = FakeQSTileIntentUserInputHandler()
+    private val kosmos = testKosmos()
+    private val inputHandler = kosmos.qsTileIntentUserInputHandler
+    private val mockDialogDelegate = kosmos.mockModesDialogDelegate
 
-    @Mock private lateinit var dialogTransitionAnimator: DialogTransitionAnimator
-    @Mock private lateinit var dialogDelegate: ModesDialogDelegate
-    @Mock private lateinit var mockDialog: SystemUIDialog
+    private val underTest =
+        ModesTileUserActionInteractor(
+            inputHandler,
+            mockDialogDelegate,
+        )
 
-    private lateinit var underTest: ModesTileUserActionInteractor
-
-    @Before
-    fun setup() {
-        MockitoAnnotations.initMocks(this)
-
-        whenever(dialogDelegate.createDialog()).thenReturn(mockDialog)
-
-        underTest =
-            ModesTileUserActionInteractor(
-                EmptyCoroutineContext,
-                inputHandler,
-                dialogTransitionAnimator,
-                dialogDelegate,
+    @Test
+    fun handleClick_active() = runTest {
+        val expandable = mock<Expandable>()
+        underTest.handleInput(
+            QSTileInputTestKtx.click(
+                data = ModesTileModel(true, listOf("DND")),
+                expandable = expandable
             )
+        )
+
+        verify(mockDialogDelegate).showDialog(eq(expandable))
     }
 
     @Test
-    fun handleClick() = runTest {
-        underTest.handleInput(QSTileInputTestKtx.click(ModesTileModel(false)))
+    fun handleClick_inactive() = runTest {
+        val expandable = mock<Expandable>()
+        underTest.handleInput(
+            QSTileInputTestKtx.click(
+                data = ModesTileModel(false, emptyList()),
+                expandable = expandable
+            )
+        )
 
-        verify(mockDialog).show()
+        verify(mockDialogDelegate).showDialog(eq(expandable))
     }
 
     @Test
-    fun handleLongClick() = runTest {
-        underTest.handleInput(QSTileInputTestKtx.longClick(ModesTileModel(false)))
+    fun handleLongClick_active() = runTest {
+        underTest.handleInput(QSTileInputTestKtx.longClick(ModesTileModel(true, listOf("DND"))))
 
         QSTileIntentUserInputHandlerSubject.assertThat(inputHandler).handledOneIntentInput {
-            Truth.assertThat(it.intent.action).isEqualTo(Settings.ACTION_ZEN_MODE_SETTINGS)
+            assertThat(it.intent.action).isEqualTo(Settings.ACTION_ZEN_MODE_SETTINGS)
+        }
+    }
+
+    @Test
+    fun handleLongClick_inactive() = runTest {
+        underTest.handleInput(QSTileInputTestKtx.longClick(ModesTileModel(false, emptyList())))
+
+        QSTileIntentUserInputHandlerSubject.assertThat(inputHandler).handledOneIntentInput {
+            assertThat(it.intent.action).isEqualTo(Settings.ACTION_ZEN_MODE_SETTINGS)
         }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
index 3baf2f4..dd9711e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
@@ -54,21 +54,35 @@
 
     @Test
     fun inactiveState() {
-        val model = ModesTileModel(isActivated = false)
+        val model = ModesTileModel(isActivated = false, activeModes = emptyList())
 
         val state = underTest.map(config, model)
 
         assertThat(state.activationState).isEqualTo(QSTileState.ActivationState.INACTIVE)
         assertThat(state.iconRes).isEqualTo(R.drawable.qs_dnd_icon_off)
+        assertThat(state.secondaryLabel).isEqualTo("No active modes")
     }
 
     @Test
-    fun activeState() {
-        val model = ModesTileModel(isActivated = true)
+    fun activeState_oneMode() {
+        val model = ModesTileModel(isActivated = true, activeModes = listOf("DND"))
 
         val state = underTest.map(config, model)
 
         assertThat(state.activationState).isEqualTo(QSTileState.ActivationState.ACTIVE)
         assertThat(state.iconRes).isEqualTo(R.drawable.qs_dnd_icon_on)
+        assertThat(state.secondaryLabel).isEqualTo("DND is active")
+    }
+
+    @Test
+    fun activeState_multipleModes() {
+        val model =
+            ModesTileModel(isActivated = true, activeModes = listOf("Mode 1", "Mode 2", "Mode 3"))
+
+        val state = underTest.map(config, model)
+
+        assertThat(state.activationState).isEqualTo(QSTileState.ActivationState.ACTIVE)
+        assertThat(state.iconRes).isEqualTo(R.drawable.qs_dnd_icon_on)
+        assertThat(state.secondaryLabel).isEqualTo("3 modes are active")
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
index 636d5a7..4a7b887 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
@@ -46,6 +46,7 @@
 
 import androidx.test.filters.SmallTest;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.internal.colorextraction.ColorExtractor;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.biometrics.AuthController;
@@ -89,7 +90,7 @@
 @RunWithLooper(setAsMainLooper = true)
 @SmallTest
 public class NotificationShadeWindowControllerImplTest extends SysuiTestCase {
-    @Mock private WindowManager mWindowManager;
+    @Mock private ViewCaptureAwareWindowManager mWindowManager;
     @Mock private DozeParameters mDozeParameters;
     @Spy private final NotificationShadeWindowView mNotificationShadeWindowView = spy(
             new NotificationShadeWindowView(mContext, null));
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
index 6f09931..9fea7a2 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
@@ -19,12 +19,10 @@
 
 package com.android.systemui.statusbar.notification.stack.ui.viewmodel
 
-import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.FlagsParameterization
 import androidx.test.filters.SmallTest
 import com.android.compose.animation.scene.ObservableTransitionState
-import com.android.systemui.Flags.FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX
 import com.android.systemui.Flags.FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.bouncer.data.repository.keyguardBouncerRepository
@@ -38,8 +36,8 @@
 import com.android.systemui.flags.DisableSceneContainer
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.Flags
-import com.android.systemui.flags.andSceneContainer
 import com.android.systemui.flags.fakeFeatureFlagsClassic
+import com.android.systemui.flags.parameterizeSceneContainerFlag
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
@@ -60,7 +58,6 @@
 import com.android.systemui.res.R
 import com.android.systemui.shade.mockLargeScreenHeaderHelper
 import com.android.systemui.shade.shadeTestUtil
-import com.android.systemui.statusbar.notification.NotificationUtils.interpolate
 import com.android.systemui.statusbar.notification.stack.domain.interactor.sharedNotificationContainerInteractor
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
@@ -90,10 +87,7 @@
         @JvmStatic
         @Parameters(name = "{0}")
         fun getParams(): List<FlagsParameterization> {
-            return FlagsParameterization.allCombinationsOf(
-                    FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX,
-                )
-                .andSceneContainer()
+            return parameterizeSceneContainerFlag()
         }
     }
 
@@ -178,25 +172,6 @@
         }
 
     @Test
-    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
-    fun validatePaddingTopInSplitShade_refactorFlagOff_usesLargeHeaderResource() =
-        testScope.runTest {
-            whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight()).thenReturn(5)
-            overrideResource(R.bool.config_use_split_notification_shade, true)
-            overrideResource(R.bool.config_use_large_screen_shade_header, true)
-            overrideResource(R.dimen.large_screen_shade_header_height, 10)
-            overrideResource(R.dimen.keyguard_split_shade_top_margin, 50)
-
-            val paddingTop by collectLastValue(underTest.paddingTopDimen)
-
-            configurationRepository.onAnyConfigurationChange()
-
-            // Should directly use the header height (flagged off value)
-            assertThat(paddingTop).isEqualTo(10)
-        }
-
-    @Test
-    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     fun validatePaddingTopInSplitShade_refactorFlagOn_usesLargeHeaderHelper() =
         testScope.runTest {
             whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight()).thenReturn(5)
@@ -268,49 +243,8 @@
         }
 
     @Test
-    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     @DisableSceneContainer
-    fun validateMarginTopWithLargeScreenHeader_refactorFlagOff_usesResource() =
-        testScope.runTest {
-            val headerResourceHeight = 50
-            val headerHelperHeight = 100
-            whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight())
-                .thenReturn(headerHelperHeight)
-            overrideResource(R.bool.config_use_large_screen_shade_header, true)
-            overrideResource(R.dimen.large_screen_shade_header_height, headerResourceHeight)
-            overrideResource(R.dimen.notification_panel_margin_top, 0)
-
-            val dimens by collectLastValue(underTest.configurationBasedDimensions)
-
-            configurationRepository.onAnyConfigurationChange()
-
-            assertThat(dimens!!.marginTop).isEqualTo(headerResourceHeight)
-        }
-
-    @Test
-    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
-    @EnableSceneContainer
-    fun validateMarginTopWithLargeScreenHeader_refactorFlagOff_sceneContainerFlagOn_stillZero() =
-        testScope.runTest {
-            val headerResourceHeight = 50
-            val headerHelperHeight = 100
-            whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight())
-                .thenReturn(headerHelperHeight)
-            overrideResource(R.bool.config_use_large_screen_shade_header, true)
-            overrideResource(R.dimen.large_screen_shade_header_height, headerResourceHeight)
-            overrideResource(R.dimen.notification_panel_margin_top, 0)
-
-            val dimens by collectLastValue(underTest.configurationBasedDimensions)
-
-            configurationRepository.onAnyConfigurationChange()
-
-            assertThat(dimens!!.marginTop).isEqualTo(0)
-        }
-
-    @Test
-    @DisableSceneContainer
-    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
-    fun validateMarginTopWithLargeScreenHeader_refactorFlagOn_usesHelper() =
+    fun validateMarginTopWithLargeScreenHeader_usesHelper() =
         testScope.runTest {
             val headerResourceHeight = 50
             val headerHelperHeight = 100
@@ -329,7 +263,6 @@
 
     @Test
     @EnableSceneContainer
-    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     fun validateMarginTopWithLargeScreenHeader_sceneContainerFlagOn_stillZero() =
         testScope.runTest {
             val headerResourceHeight = 50
@@ -590,44 +523,45 @@
         }
 
     @Test
-    @DisableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     @DisableSceneContainer
-    fun boundsOnLockscreenInSplitShade_refactorFlagOff_usesLargeHeaderResource() =
+    fun boundsDoNotChangeWhileLockscreenToAodTransitionIsActive() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
 
-            // When in split shade
-            whenever(largeScreenHeaderHelper.getLargeScreenHeaderHeight()).thenReturn(5)
-            overrideResource(R.bool.config_use_split_notification_shade, true)
-            overrideResource(R.bool.config_use_large_screen_shade_header, true)
-            overrideResource(R.dimen.large_screen_shade_header_height, 10)
-            overrideResource(R.dimen.keyguard_split_shade_top_margin, 50)
-
-            configurationRepository.onAnyConfigurationChange()
-            runCurrent()
-
             // Start on lockscreen
             showLockscreen()
 
             keyguardInteractor.setNotificationContainerBounds(
-                NotificationContainerBounds(top = 1f, bottom = 52f)
+                NotificationContainerBounds(top = 1f, bottom = 1f)
+            )
+            assertThat(bounds).isEqualTo(NotificationContainerBounds(top = 1f, bottom = 1f))
+
+            // Begin transition to AOD
+            keyguardTransitionRepository.sendTransitionStep(
+                TransitionStep(LOCKSCREEN, AOD, 0f, TransitionState.STARTED)
             )
             runCurrent()
+            keyguardTransitionRepository.sendTransitionStep(
+                TransitionStep(LOCKSCREEN, AOD, 0.5f, TransitionState.RUNNING)
+            )
 
-            // Top should be equal to bounds (1) - padding adjustment (10)
-            assertThat(bounds)
-                .isEqualTo(
-                    NotificationContainerBounds(
-                        top = -9f,
-                        bottom = 2f,
-                    )
-                )
+            // Attempt to update bounds
+            keyguardInteractor.setNotificationContainerBounds(
+                NotificationContainerBounds(top = 5f, bottom = 5f)
+            )
+            // Bounds should not have moved
+            assertThat(bounds).isEqualTo(NotificationContainerBounds(top = 1f, bottom = 1f))
+
+            // Transition is over, now move
+            keyguardTransitionRepository.sendTransitionStep(
+                TransitionStep(LOCKSCREEN, AOD, 1f, TransitionState.FINISHED)
+            )
+            assertThat(bounds).isEqualTo(NotificationContainerBounds(top = 5f, bottom = 5f))
         }
 
     @Test
-    @EnableFlags(FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
     @DisableSceneContainer
-    fun boundsOnLockscreenInSplitShade_refactorFlagOn_usesLargeHeaderHelper() =
+    fun boundsOnLockscreenInSplitShade_usesLargeHeaderHelper() =
         testScope.runTest {
             val bounds by collectLastValue(underTest.bounds)
 
@@ -820,54 +754,6 @@
 
     @Test
     @DisableSceneContainer
-    fun updateBounds_fromKeyguardRoot() =
-        testScope.runTest {
-            val startProgress = 0f
-            val startStep = TransitionStep(LOCKSCREEN, AOD, startProgress, TransitionState.STARTED)
-            val boundsChangingProgress = 0.2f
-            val boundsChangingStep =
-                TransitionStep(LOCKSCREEN, AOD, boundsChangingProgress, TransitionState.RUNNING)
-            val boundsInterpolatingProgress = 0.6f
-            val boundsInterpolatingStep =
-                TransitionStep(
-                    LOCKSCREEN,
-                    AOD,
-                    boundsInterpolatingProgress,
-                    TransitionState.RUNNING
-                )
-            val finishProgress = 1.0f
-            val finishStep =
-                TransitionStep(LOCKSCREEN, AOD, finishProgress, TransitionState.FINISHED)
-
-            val bounds by collectLastValue(underTest.bounds)
-            val top = 123f
-            val bottom = 456f
-
-            kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(startStep)
-            runCurrent()
-            kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(boundsChangingStep)
-            runCurrent()
-            keyguardRootViewModel.onNotificationContainerBoundsChanged(top, bottom)
-
-            kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(boundsInterpolatingStep)
-            runCurrent()
-            val adjustedProgress =
-                (boundsInterpolatingProgress - boundsChangingProgress) /
-                    (1 - boundsChangingProgress)
-            val interpolatedTop = interpolate(0f, top, adjustedProgress)
-            val interpolatedBottom = interpolate(0f, bottom, adjustedProgress)
-            assertThat(bounds)
-                .isEqualTo(
-                    NotificationContainerBounds(top = interpolatedTop, bottom = interpolatedBottom)
-                )
-
-            kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(finishStep)
-            runCurrent()
-            assertThat(bounds).isEqualTo(NotificationContainerBounds(top = top, bottom = bottom))
-        }
-
-    @Test
-    @DisableSceneContainer
     fun updateBounds_fromGone_withoutTransitions() =
         testScope.runTest {
             // Start step is already at 1.0
@@ -878,9 +764,9 @@
             val top = 123f
             val bottom = 456f
 
-            kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(runningStep)
+            keyguardTransitionRepository.sendTransitionStep(runningStep)
             runCurrent()
-            kosmos.fakeKeyguardTransitionRepository.sendTransitionStep(finishStep)
+            keyguardTransitionRepository.sendTransitionStep(finishStep)
             runCurrent()
             keyguardRootViewModel.onNotificationContainerBoundsChanged(top, bottom)
             runCurrent()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
index fdfc7f1..62161bf 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
@@ -18,6 +18,8 @@
 
 package com.android.systemui.statusbar.policy.ui.dialog.viewmodel
 
+import android.content.Intent
+import android.provider.Settings
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.settingslib.notification.modes.TestModeBuilder
@@ -27,32 +29,46 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.statusbar.policy.data.repository.fakeZenModeRepository
 import com.android.systemui.statusbar.policy.domain.interactor.zenModeInteractor
+import com.android.systemui.statusbar.policy.ui.dialog.mockModesDialogDelegate
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mockito.clearInvocations
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.verify
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class ModesDialogViewModelTest : SysuiTestCase() {
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
-    val repository = kosmos.fakeZenModeRepository
-    val interactor = kosmos.zenModeInteractor
+    private val repository = kosmos.fakeZenModeRepository
+    private val interactor = kosmos.zenModeInteractor
+    private val mockDialogDelegate = kosmos.mockModesDialogDelegate
 
-    val underTest = ModesDialogViewModel(context, interactor, kosmos.testDispatcher)
+    private val underTest =
+        ModesDialogViewModel(context, interactor, kosmos.testDispatcher, mockDialogDelegate)
 
     @Test
-    fun tiles_filtersOutDisabledModes() =
+    fun tiles_filtersOutUserDisabledModes() =
         testScope.runTest {
             val tiles by collectLastValue(underTest.tiles)
 
             repository.addModes(
                 listOf(
-                    TestModeBuilder().setName("Disabled").setEnabled(false).build(),
+                    TestModeBuilder()
+                        .setName("Disabled by user")
+                        .setEnabled(false, /* byUser= */ true)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("Disabled by other")
+                        .setEnabled(false, /* byUser= */ false)
+                        .build(),
                     TestModeBuilder.MANUAL_DND,
                     TestModeBuilder()
                         .setName("Enabled")
@@ -61,19 +77,25 @@
                         .build(),
                     TestModeBuilder()
                         .setName("Disabled with manual")
-                        .setEnabled(false)
+                        .setEnabled(false, /* byUser= */ true)
                         .setManualInvocationAllowed(true)
                         .build(),
-                ))
+                )
+            )
             runCurrent()
 
-            assertThat(tiles?.size).isEqualTo(2)
+            assertThat(tiles?.size).isEqualTo(3)
             with(tiles?.elementAt(0)!!) {
+                assertThat(this.text).isEqualTo("Disabled by other")
+                assertThat(this.subtext).isEqualTo("Set up")
+                assertThat(this.enabled).isEqualTo(false)
+            }
+            with(tiles?.elementAt(1)!!) {
                 assertThat(this.text).isEqualTo("Manual DND")
                 assertThat(this.subtext).isEqualTo("On")
                 assertThat(this.enabled).isEqualTo(true)
             }
-            with(tiles?.elementAt(1)!!) {
+            with(tiles?.elementAt(2)!!) {
                 assertThat(this.text).isEqualTo("Enabled")
                 assertThat(this.subtext).isEqualTo("Off")
                 assertThat(this.enabled).isEqualTo(false)
@@ -108,7 +130,8 @@
                         .setActive(false)
                         .setManualInvocationAllowed(false)
                         .build(),
-                ))
+                )
+            )
             runCurrent()
 
             assertThat(tiles?.size).isEqualTo(3)
@@ -130,6 +153,117 @@
         }
 
     @Test
+    fun tiles_stableWhileCollecting() =
+        testScope.runTest {
+            val job = Job()
+            val tiles by collectLastValue(underTest.tiles, context = job)
+
+            repository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setName("Active without manual")
+                        .setActive(true)
+                        .setManualInvocationAllowed(false)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("Active with manual")
+                        .setActive(true)
+                        .setManualInvocationAllowed(true)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("Inactive with manual")
+                        .setActive(false)
+                        .setManualInvocationAllowed(true)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("Inactive without manual")
+                        .setActive(false)
+                        .setManualInvocationAllowed(false)
+                        .build(),
+                )
+            )
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(3)
+
+            // Check that tile is initially present
+            with(tiles?.elementAt(0)!!) {
+                assertThat(this.text).isEqualTo("Active without manual")
+                assertThat(this.subtext).isEqualTo("On")
+                assertThat(this.enabled).isEqualTo(true)
+
+                // Click tile to toggle it
+                this.onClick()
+                runCurrent()
+            }
+            // Check that tile is still present at the same location, but turned off
+            assertThat(tiles?.size).isEqualTo(3)
+            with(tiles?.elementAt(0)!!) {
+                assertThat(this.text).isEqualTo("Active without manual")
+                assertThat(this.subtext).isEqualTo("Manage in settings")
+                assertThat(this.enabled).isEqualTo(false)
+            }
+
+            // Stop collecting, then start again
+            job.cancel()
+            val tiles2 by collectLastValue(underTest.tiles)
+            runCurrent()
+
+            // Check that tile is now gone
+            assertThat(tiles2?.size).isEqualTo(2)
+            assertThat(tiles2?.elementAt(0)!!.text).isEqualTo("Active with manual")
+            assertThat(tiles2?.elementAt(1)!!.text).isEqualTo("Inactive with manual")
+        }
+
+    @Test
+    fun tiles_filtersOutRemovedModes() =
+        testScope.runTest {
+            val job = Job()
+            val tiles by collectLastValue(underTest.tiles, context = job)
+
+            repository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setId("A")
+                        .setName("Active without manual")
+                        .setActive(true)
+                        .setManualInvocationAllowed(false)
+                        .build(),
+                    TestModeBuilder()
+                        .setId("B")
+                        .setName("Active with manual")
+                        .setActive(true)
+                        .setManualInvocationAllowed(true)
+                        .build(),
+                    TestModeBuilder()
+                        .setId("C")
+                        .setName("Inactive with manual")
+                        .setActive(false)
+                        .setManualInvocationAllowed(true)
+                        .build(),
+                )
+            )
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(3)
+
+            repository.removeMode("A")
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(2)
+
+            repository.removeMode("B")
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(1)
+
+            repository.removeMode("C")
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(0)
+        }
+
+    @Test
     fun onClick_togglesTileState() =
         testScope.runTest {
             val tiles by collectLastValue(underTest.tiles)
@@ -161,4 +295,141 @@
 
             assertThat(tiles?.first()?.enabled).isFalse()
         }
+
+    @Test
+    fun onClick_noManualActivation() =
+        testScope.runTest {
+            val job = Job()
+            val tiles by collectLastValue(underTest.tiles, context = job)
+
+            repository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setName("Active without manual")
+                        .setActive(true)
+                        .setManualInvocationAllowed(false)
+                        .build(),
+                )
+            )
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(1)
+
+            // Click tile to toggle it off
+            tiles?.elementAt(0)!!.onClick()
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(1)
+            with(tiles?.elementAt(0)!!) {
+                assertThat(this.text).isEqualTo("Active without manual")
+                assertThat(this.subtext).isEqualTo("Manage in settings")
+                assertThat(this.enabled).isEqualTo(false)
+
+                // Press the tile again
+                this.onClick()
+                runCurrent()
+            }
+
+            // Check that nothing happened
+            with(tiles?.elementAt(0)!!) {
+                assertThat(this.text).isEqualTo("Active without manual")
+                assertThat(this.subtext).isEqualTo("Manage in settings")
+                assertThat(this.enabled).isEqualTo(false)
+            }
+        }
+
+    @Test
+    fun onClick_setUp() =
+        testScope.runTest {
+            val tiles by collectLastValue(underTest.tiles)
+
+            repository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setId("ID")
+                        .setName("Disabled by other")
+                        .setEnabled(false, /* byUser= */ false)
+                        .build(),
+                )
+            )
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(1)
+            with(tiles?.elementAt(0)!!) {
+                assertThat(this.text).isEqualTo("Disabled by other")
+                assertThat(this.subtext).isEqualTo("Set up")
+                assertThat(this.enabled).isEqualTo(false)
+
+                // Click the tile
+                this.onClick()
+                runCurrent()
+            }
+
+            // Check that it launched the correct intent
+            val intentCaptor = argumentCaptor<Intent>()
+            verify(mockDialogDelegate).launchFromDialog(intentCaptor.capture())
+            val intent = intentCaptor.lastValue
+            assertThat(intent.action).isEqualTo(Settings.ACTION_AUTOMATIC_ZEN_RULE_SETTINGS)
+            assertThat(intent.extras?.getString(Settings.EXTRA_AUTOMATIC_ZEN_RULE_ID))
+                .isEqualTo("ID")
+
+            // Check that nothing happened to the tile
+            with(tiles?.elementAt(0)!!) {
+                assertThat(this.text).isEqualTo("Disabled by other")
+                assertThat(this.subtext).isEqualTo("Set up")
+                assertThat(this.enabled).isEqualTo(false)
+            }
+        }
+
+    @Test
+    fun onLongClick_launchesIntent() =
+        testScope.runTest {
+            val tiles by collectLastValue(underTest.tiles)
+            val intentCaptor = argumentCaptor<Intent>()
+
+            val modeId = "id"
+            repository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setId(modeId)
+                        .setId("A")
+                        .setActive(true)
+                        .setManualInvocationAllowed(true)
+                        .build(),
+                    TestModeBuilder()
+                        .setId(modeId)
+                        .setId("B")
+                        .setActive(false)
+                        .setManualInvocationAllowed(true)
+                        .build(),
+                )
+            )
+            runCurrent()
+
+            assertThat(tiles?.size).isEqualTo(2)
+
+            // Trigger onLongClick for A
+            tiles?.first()?.onLongClick?.let { it() }
+            runCurrent()
+
+            // Check that it launched the correct intent
+            verify(mockDialogDelegate).launchFromDialog(intentCaptor.capture())
+            var intent = intentCaptor.lastValue
+            assertThat(intent.action).isEqualTo(Settings.ACTION_AUTOMATIC_ZEN_RULE_SETTINGS)
+            assertThat(intent.extras?.getString(Settings.EXTRA_AUTOMATIC_ZEN_RULE_ID))
+                .isEqualTo("A")
+
+            clearInvocations(mockDialogDelegate)
+
+            // Trigger onLongClick for B
+            tiles?.last()?.onLongClick?.let { it() }
+            runCurrent()
+
+            // Check that it launched the correct intent
+            verify(mockDialogDelegate).launchFromDialog(intentCaptor.capture())
+            intent = intentCaptor.lastValue
+            assertThat(intent.action).isEqualTo(Settings.ACTION_AUTOMATIC_ZEN_RULE_SETTINGS)
+            assertThat(intent.extras?.getString(Settings.EXTRA_AUTOMATIC_ZEN_RULE_ID))
+                .isEqualTo("B")
+        }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorTest.kt
index 142631e..a1fcfcd 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorTest.kt
@@ -16,8 +16,10 @@
 
 package com.android.systemui.volume.domain.interactor
 
+import android.media.AudioManager.STREAM_MUSIC
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.settingslib.volume.shared.model.AudioStream
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testScope
@@ -40,17 +42,57 @@
 
     @Before
     fun setUp() {
-        with(kosmos) { underTest = audioSharingInteractor }
+        with(kosmos) {
+            with(audioSharingRepository) { setVolumeMap(mapOf(TEST_GROUP_ID to TEST_VOLUME)) }
+            underTest = audioSharingInteractor
+        }
     }
 
     @Test
-    fun volumeChanges_returnVolume() {
+    fun handlePrimaryGroupChange_nullVolume() {
         with(kosmos) {
             testScope.runTest {
-                with(audioSharingRepository) {
-                    setSecondaryGroupId(TEST_GROUP_ID)
-                    setVolumeMap(mapOf(TEST_GROUP_ID to TEST_VOLUME))
-                }
+                with(audioSharingRepository) { setPrimaryGroupId(TEST_GROUP_ID_INVALID) }
+                val preMusicStream by
+                    collectLastValue(
+                        audioVolumeInteractor.getAudioStream(AudioStream(STREAM_MUSIC))
+                    )
+                val preVolume = preMusicStream?.volume
+                runCurrent()
+                underTest.handlePrimaryGroupChange()
+                val musicStream by
+                    collectLastValue(
+                        audioVolumeInteractor.getAudioStream(AudioStream(STREAM_MUSIC))
+                    )
+                runCurrent()
+
+                Truth.assertThat(musicStream?.volume).isEqualTo(preVolume)
+            }
+        }
+    }
+
+    @Test
+    fun handlePrimaryGroupChange_setStreamVolume() {
+        with(kosmos) {
+            testScope.runTest {
+                with(audioSharingRepository) { setPrimaryGroupId(TEST_GROUP_ID) }
+                underTest.handlePrimaryGroupChange()
+                val musicStream by
+                    collectLastValue(
+                        audioVolumeInteractor.getAudioStream(AudioStream(STREAM_MUSIC))
+                    )
+                runCurrent()
+
+                Truth.assertThat(musicStream?.volume).isEqualTo(TEST_MUSIC_VOLUME)
+            }
+        }
+    }
+
+    @Test
+    fun secondaryGroupVolumeChanges_returnVolume() {
+        with(kosmos) {
+            testScope.runTest {
+                with(audioSharingRepository) { setSecondaryGroupId(TEST_GROUP_ID) }
                 val volume by collectLastValue(underTest.volume)
                 runCurrent()
 
@@ -60,13 +102,10 @@
     }
 
     @Test
-    fun volumeChanges_returnNull() {
+    fun secondaryGroupVolumeChanges_returnNull() {
         with(kosmos) {
             testScope.runTest {
-                with(audioSharingRepository) {
-                    setSecondaryGroupId(TEST_GROUP_ID_INVALID)
-                    setVolumeMap(mapOf(TEST_GROUP_ID to TEST_VOLUME))
-                }
+                with(audioSharingRepository) { setSecondaryGroupId(TEST_GROUP_ID_INVALID) }
                 val volume by collectLastValue(underTest.volume)
                 runCurrent()
 
@@ -76,7 +115,7 @@
     }
 
     @Test
-    fun volumeChanges_returnDefaultVolume() {
+    fun secondaryGroupVolumeChanges_returnDefaultVolume() {
         with(kosmos) {
             testScope.runTest {
                 with(audioSharingRepository) {
@@ -94,7 +133,8 @@
     private companion object {
         const val TEST_GROUP_ID = 1
         const val TEST_GROUP_ID_INVALID = -1
-        const val TEST_VOLUME = 10
+        const val TEST_MUSIC_VOLUME = 10
+        const val TEST_VOLUME = 255
         const val TEST_VOLUME_DEFAULT = 20
     }
 }
diff --git a/packages/SystemUI/res/layout/app_clips_backlinks_drop_down_entry.xml b/packages/SystemUI/res/layout/app_clips_backlinks_drop_down_entry.xml
new file mode 100644
index 0000000..7eab340
--- /dev/null
+++ b/packages/SystemUI/res/layout/app_clips_backlinks_drop_down_entry.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 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.
+  -->
+<TextView xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="48dp"
+    android:drawablePadding="4dp"
+    android:ellipsize="end"
+    android:gravity="center_vertical"
+    android:paddingHorizontal="8dp"
+    android:textColor="?android:textColorSecondary" />
diff --git a/packages/SystemUI/res/values/config.xml b/packages/SystemUI/res/values/config.xml
index 30f23bf..c29c236 100644
--- a/packages/SystemUI/res/values/config.xml
+++ b/packages/SystemUI/res/values/config.xml
@@ -104,7 +104,7 @@
 
     <!-- Tiles native to System UI. Order should match "quick_settings_tiles_default" -->
     <string name="quick_settings_tiles_stock" translatable="false">
-        internet,bt,flashlight,dnd,modes,alarm,airplane,controls,wallet,rotation,battery,cast,screenrecord,mictoggle,cameratoggle,location,hotspot,inversion,saver,dark,work,night,reverse,reduce_brightness,qr_code_scanner,onehanded,color_correction,dream,font_scaling,record_issue,hearing_devices
+        internet,bt,flashlight,dnd,alarm,airplane,controls,wallet,rotation,battery,cast,screenrecord,mictoggle,cameratoggle,location,hotspot,inversion,saver,dark,work,night,reverse,reduce_brightness,qr_code_scanner,onehanded,color_correction,dream,font_scaling,record_issue,hearing_devices
     </string>
 
     <!-- The tiles to display in QuickSettings -->
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 8146cc5..d7c3527 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -1103,6 +1103,20 @@
     <!-- Priority modes: label for an inactive mode [CHAR LIMIT=35] -->
     <string name="zen_mode_off">Off</string>
 
+    <!-- Priority modes: label for a mode that needs to be set up [CHAR LIMIT=35] -->
+    <string name="zen_mode_set_up">Set up</string>
+
+    <!-- Priority modes: label for a mode that cannot be manually turned on [CHAR LIMIT=35] -->
+    <string name="zen_mode_no_manual_invocation">Manage in settings</string>
+
+    <string name="zen_mode_active_modes">
+        {count, plural,
+            =0 {No active modes}
+            =1 {{mode} is active}
+            other {# modes are active}
+        }
+    </string>
+
     <!-- Zen mode: Priority only introduction message on first use -->
     <string name="zen_priority_introduction">You won\'t be disturbed by sounds and vibrations, except from alarms, reminders, events, and callers you specify. You\'ll still hear anything you choose to play including music, videos, and games.</string>
 
@@ -1211,6 +1225,9 @@
     <!-- Label for accessibility action that shows widgets on lock screen on click. [CHAR LIMIT=NONE] -->
     <string name="accessibility_action_open_communal_hub">Widgets on lock screen</string>
 
+    <!-- Label for an accessibility announcement when a widget has been added to the lock screen. [CHAR LIMIT=NONE] -->
+    <string name="accessibility_announcement_communal_widget_added"><xliff:g id="widget_name" example="Calendar month view">%1$s</xliff:g> widget added to lock screen</string>
+
     <!-- Indicator on keyguard to start the communal tutorial. [CHAR LIMIT=100] -->
     <string name="communal_tutorial_indicator_text">Swipe left to start the communal tutorial</string>
 
@@ -3208,9 +3225,6 @@
 
     <!-- Provider Model: Default title of the mobile network in the mobile layout. [CHAR LIMIT=50] -->
     <string name="mobile_data_settings_title">Mobile data</string>
-    <!-- Provider Model: Summary text separator for preferences including a short description
-         (eg. "Connected / 5G"). [CHAR LIMIT=50] -->
-    <string name="preference_summary_default_combination"><xliff:g id="state" example="Connected">%1$s</xliff:g> / <xliff:g id="networkMode" example="LTE">%2$s</xliff:g></string>
     <!-- Provider Model:
          Summary indicating that a SIM has an active mobile data connection [CHAR LIMIT=50] -->
     <string name="mobile_data_connection_active">Connected</string>
diff --git a/packages/SystemUI/res/values/tiles_states_strings.xml b/packages/SystemUI/res/values/tiles_states_strings.xml
index c702927..ad09b46 100644
--- a/packages/SystemUI/res/values/tiles_states_strings.xml
+++ b/packages/SystemUI/res/values/tiles_states_strings.xml
@@ -85,16 +85,6 @@
         <item>On</item>
     </string-array>
 
-    <!-- State names for modes (Priority modes) tile: unavailable, off, on.
-         This subtitle is shown when the tile is in that particular state but does not set its own
-         subtitle, so some of these may never appear on screen. They should still be translated as
-         if they could appear. [CHAR LIMIT=32] -->
-    <string-array name="tile_states_modes">
-        <item>Unavailable</item>
-        <item>Off</item>
-        <item>On</item>
-    </string-array>
-
     <!-- State names for flashlight tile: unavailable, off, on.
          This subtitle is shown when the tile is in that particular state but does not set its own
          subtitle, so some of these may never appear on screen. They should still be translated as
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
index 428cd0e..93ee179 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
@@ -724,7 +724,10 @@
     @Override
     public void onResume(int reason) {
         if (DEBUG) Log.d(TAG, "screen on, instance " + Integer.toHexString(hashCode()));
+        mView.clearFocus();
+        mView.clearAccessibilityFocus();
         mView.requestFocus();
+        mView.requestAccessibilityFocus();
         if (mCurrentSecurityMode != SecurityMode.None) {
             int state = SysUiStatsLog.KEYGUARD_BOUNCER_STATE_CHANGED__STATE__SHOWN;
             if (mView.isSidedSecurityMode()) {
diff --git a/packages/SystemUI/src/com/android/keyguard/logging/KeyguardQuickAffordancesLogger.kt b/packages/SystemUI/src/com/android/keyguard/logging/KeyguardQuickAffordancesLogger.kt
new file mode 100644
index 0000000..c11cf55
--- /dev/null
+++ b/packages/SystemUI/src/com/android/keyguard/logging/KeyguardQuickAffordancesLogger.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2024 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.keyguard.logging
+
+import com.android.systemui.log.LogBuffer
+import com.android.systemui.log.core.LogLevel
+import com.android.systemui.log.dagger.KeyguardQuickAffordancesLog
+import javax.inject.Inject
+
+class KeyguardQuickAffordancesLogger
+@Inject
+constructor(
+    @KeyguardQuickAffordancesLog val buffer: LogBuffer,
+) {
+    fun logQuickAffordanceTapped(configKey: String?) {
+        val (slotId, affordanceId) = configKey?.decode() ?: ("" to "")
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = affordanceId
+                str2 = slotId
+            },
+            { "QuickAffordance tapped with id: $str1, in slot: $str2" }
+        )
+    }
+
+    fun logQuickAffordanceTriggered(slotId: String, affordanceId: String) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = affordanceId
+                str2 = slotId
+            },
+            { "QuickAffordance triggered with id: $str1, in slot: $str2" }
+        )
+    }
+
+    fun logQuickAffordanceSelected(slotId: String, affordanceId: String) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = affordanceId
+                str2 = slotId
+            },
+            { "QuickAffordance selected with id: $str1, in slot: $str2" }
+        )
+    }
+
+    private fun String.decode(): Pair<String, String> {
+        val splitUp = this.split(DELIMITER)
+        return Pair(splitUp[0], splitUp[1])
+    }
+
+    companion object {
+        private const val TAG = "KeyguardQuickAffordancesLogger"
+        private const val DELIMITER = "::"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java b/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java
index 3c0ac9a..394f8dd 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java
@@ -30,6 +30,8 @@
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.graphics.Region;
+import android.graphics.drawable.GradientDrawable;
+import android.hardware.display.DisplayManager;
 import android.os.Handler;
 import android.util.Log;
 import android.view.AttachedSurfaceControl;
@@ -49,6 +51,8 @@
 import androidx.annotation.UiThread;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.policy.ScreenDecorationsUtils;
+import com.android.systemui.Flags;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.res.R;
 import com.android.systemui.util.leak.RotationUtils;
@@ -70,6 +74,7 @@
     private SurfaceControl.Transaction mTransaction;
     private View mFullscreenBorder = null;
     private int mBorderOffset;
+    private int mBorderStoke;
     private final int mDisplayId;
     private static final Region sEmptyRegion = new Region();
     private ValueAnimator mShowHideBorderAnimator;
@@ -86,16 +91,20 @@
         }
     };
     private final long mLongAnimationTimeMs;
+    private final DisplayManager mDisplayManager;
+    private final DisplayManager.DisplayListener mDisplayListener;
+    private String mCurrentDisplayUniqueId;
 
     public FullscreenMagnificationController(
             @UiContext Context context,
             @Main Handler handler,
             @Main Executor executor,
+            DisplayManager displayManager,
             AccessibilityManager accessibilityManager,
             WindowManager windowManager,
             IWindowManager iWindowManager,
             Supplier<SurfaceControlViewHost> scvhSupplier) {
-        this(context, handler, executor, accessibilityManager,
+        this(context, handler, executor, displayManager, accessibilityManager,
                 windowManager, iWindowManager, scvhSupplier,
                 new SurfaceControl.Transaction(), null);
     }
@@ -105,6 +114,7 @@
             @UiContext Context context,
             @Main Handler handler,
             @Main Executor executor,
+            DisplayManager displayManager,
             AccessibilityManager accessibilityManager,
             WindowManager windowManager,
             IWindowManager iWindowManager,
@@ -120,10 +130,7 @@
         mWindowBounds = mWindowManager.getCurrentWindowMetrics().getBounds();
         mTransaction = transaction;
         mScvhSupplier = scvhSupplier;
-        mBorderOffset = mContext.getResources().getDimensionPixelSize(
-                R.dimen.magnifier_border_width_fullscreen_with_offset)
-                - mContext.getResources().getDimensionPixelSize(
-                R.dimen.magnifier_border_width_fullscreen);
+        updateDimensions();
         mDisplayId = mContext.getDisplayId();
         mConfiguration = new Configuration(context.getResources().getConfiguration());
         mLongAnimationTimeMs = mContext.getResources().getInteger(
@@ -140,6 +147,31 @@
                 }
             }
         });
+        mCurrentDisplayUniqueId = mContext.getDisplayNoVerify().getUniqueId();
+        mDisplayManager = displayManager;
+        mDisplayListener = new DisplayManager.DisplayListener() {
+            @Override
+            public void onDisplayAdded(int displayId) {
+                // Do nothing
+            }
+
+            @Override
+            public void onDisplayRemoved(int displayId) {
+                // Do nothing
+            }
+
+            @Override
+            public void onDisplayChanged(int displayId) {
+                final String uniqueId = mContext.getDisplayNoVerify().getUniqueId();
+                if (uniqueId.equals(mCurrentDisplayUniqueId)) {
+                    // Same unique ID means the physical display doesn't change. Early return.
+                    return;
+                }
+
+                mCurrentDisplayUniqueId = uniqueId;
+                applyCornerRadiusToBorder();
+            }
+        };
     }
 
     private ValueAnimator createNullTargetObjectAnimator() {
@@ -180,10 +212,15 @@
         }
         mContext.unregisterComponentCallbacks(this);
 
+
         mShowHideBorderAnimator.reverse();
     }
 
     private void cleanUpBorder() {
+        if (Flags.updateCornerRadiusOnDisplayChanged()) {
+            mDisplayManager.unregisterDisplayListener(mDisplayListener);
+        }
+
         if (mSurfaceControlViewHost != null) {
             mSurfaceControlViewHost.release();
             mSurfaceControlViewHost = null;
@@ -226,6 +263,9 @@
             } catch (Exception e) {
                 Log.w(TAG, "Failed to register rotation watcher", e);
             }
+            if (Flags.updateCornerRadiusOnDisplayChanged()) {
+                mHandler.post(this::applyCornerRadiusToBorder);
+            }
         }
 
         mTransaction
@@ -247,6 +287,9 @@
 
         mAccessibilityManager.attachAccessibilityOverlayToDisplay(
                 mDisplayId, mBorderSurfaceControl);
+        if (Flags.updateCornerRadiusOnDisplayChanged()) {
+            mDisplayManager.registerDisplayListener(mDisplayListener, mHandler);
+        }
 
         applyTouchableRegion();
     }
@@ -304,6 +347,11 @@
             final int newWidth = mWindowBounds.width() + 2 * mBorderOffset;
             final int newHeight = mWindowBounds.height() + 2 * mBorderOffset;
             mSurfaceControlViewHost.relayout(newWidth, newHeight);
+            if (Flags.updateCornerRadiusOnDisplayChanged()) {
+                // Recenter the border
+                mTransaction.setPosition(
+                        mBorderSurfaceControl, -mBorderOffset, -mBorderOffset).apply();
+            }
         }
 
         // Rotating from Landscape to ReverseLandscape will not trigger the config changes in
@@ -352,6 +400,22 @@
                 R.dimen.magnifier_border_width_fullscreen_with_offset)
                 - mContext.getResources().getDimensionPixelSize(
                         R.dimen.magnifier_border_width_fullscreen);
+        mBorderStoke = mContext.getResources().getDimensionPixelSize(
+                R.dimen.magnifier_border_width_fullscreen_with_offset);
+    }
+
+    private void applyCornerRadiusToBorder() {
+        if (!isActivated()) {
+            return;
+        }
+
+        float cornerRadius = ScreenDecorationsUtils.getWindowCornerRadius(mContext);
+        GradientDrawable backgroundDrawable = (GradientDrawable) mFullscreenBorder.getBackground();
+        backgroundDrawable.setStroke(
+                mBorderStoke,
+                mContext.getResources().getColor(
+                        R.color.magnification_border_color, mContext.getTheme()));
+        backgroundDrawable.setCornerRadius(cornerRadius);
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/MagnificationImpl.java b/packages/SystemUI/src/com/android/systemui/accessibility/MagnificationImpl.java
index e9c9bc7..93c4630 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/MagnificationImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/MagnificationImpl.java
@@ -149,6 +149,7 @@
         private final Context mContext;
         private final Handler mHandler;
         private final Executor mExecutor;
+        private final DisplayManager mDisplayManager;
         private final IWindowManager mIWindowManager;
 
         FullscreenMagnificationControllerSupplier(Context context,
@@ -159,6 +160,7 @@
             mContext = context;
             mHandler = handler;
             mExecutor = executor;
+            mDisplayManager = displayManager;
             mIWindowManager = iWindowManager;
         }
 
@@ -173,6 +175,7 @@
                     windowContext,
                     mHandler,
                     mExecutor,
+                    mDisplayManager,
                     windowContext.getSystemService(AccessibilityManager.class),
                     windowContext.getSystemService(WindowManager.class),
                     mIWindowManager,
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsController.java b/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsController.java
index 3dd3758..5ffb9ab2 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsController.java
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsController.java
@@ -56,13 +56,13 @@
 import android.view.LayoutInflater;
 import android.view.MotionEvent;
 import android.view.View;
-import android.view.WindowManager;
 import android.view.accessibility.AccessibilityManager;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.OptIn;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.InstanceId;
@@ -147,7 +147,7 @@
     private final Execution mExecution;
     private final FingerprintManager mFingerprintManager;
     @NonNull private final LayoutInflater mInflater;
-    private final WindowManager mWindowManager;
+    private final ViewCaptureAwareWindowManager mWindowManager;
     private final DelayableExecutor mFgExecutor;
     @NonNull private final Executor mBiometricExecutor;
     @NonNull private final StatusBarStateController mStatusBarStateController;
@@ -693,7 +693,7 @@
             @NonNull Execution execution,
             @NonNull LayoutInflater inflater,
             @Nullable FingerprintManager fingerprintManager,
-            @NonNull WindowManager windowManager,
+            @NonNull ViewCaptureAwareWindowManager viewCaptureAwareWindowManager,
             @NonNull StatusBarStateController statusBarStateController,
             @Main DelayableExecutor fgExecutor,
             @NonNull StatusBarKeyguardViewManager statusBarKeyguardViewManager,
@@ -741,7 +741,7 @@
         // The fingerprint manager is queried for UDFPS before this class is constructed, so the
         // fingerprint manager should never be null.
         mFingerprintManager = checkNotNull(fingerprintManager);
-        mWindowManager = windowManager;
+        mWindowManager = viewCaptureAwareWindowManager;
         mFgExecutor = fgExecutor;
         mStatusBarStateController = statusBarStateController;
         mKeyguardStateController = keyguardStateController;
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsControllerOverlay.kt b/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsControllerOverlay.kt
index e03d160..1bac0bc 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsControllerOverlay.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/UdfpsControllerOverlay.kt
@@ -44,6 +44,7 @@
 import android.view.accessibility.AccessibilityManager.TouchExplorationStateChangeListener
 import androidx.annotation.LayoutRes
 import androidx.annotation.VisibleForTesting
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.systemui.animation.ActivityTransitionAnimator
 import com.android.systemui.biometrics.domain.interactor.UdfpsOverlayInteractor
@@ -94,7 +95,7 @@
 constructor(
     private val context: Context,
     private val inflater: LayoutInflater,
-    private val windowManager: WindowManager,
+    private val windowManager: ViewCaptureAwareWindowManager,
     private val accessibilityManager: AccessibilityManager,
     private val statusBarStateController: StatusBarStateController,
     private val statusBarKeyguardViewManager: StatusBarKeyguardViewManager,
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt b/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt
index 3904ee1..b1cba2f 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt
@@ -915,7 +915,12 @@
         event: MotionEvent,
         touchExplorationEnabled: Boolean,
     ): Boolean {
-        if (bpTalkback() && modalities.first().hasUdfps && touchExplorationEnabled) {
+        if (
+            bpTalkback() &&
+                modalities.first().hasUdfps &&
+                touchExplorationEnabled &&
+                !isAuthenticated.first().isAuthenticated
+        ) {
             // TODO(b/315184924): Remove uses of UdfpsUtils
             val scaledTouch =
                 udfpsUtils.getTouchInNativeCoordinates(
diff --git a/packages/SystemUI/src/com/android/systemui/camera/CameraGestureHelper.kt b/packages/SystemUI/src/com/android/systemui/camera/CameraGestureHelper.kt
index ecbd3f9..6757edb 100644
--- a/packages/SystemUI/src/com/android/systemui/camera/CameraGestureHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/camera/CameraGestureHelper.kt
@@ -19,6 +19,7 @@
 import android.app.ActivityManager
 import android.app.ActivityOptions
 import android.app.IActivityTaskManager
+import android.app.admin.DevicePolicyManager
 import android.content.ContentResolver
 import android.content.Context
 import android.content.Intent
@@ -32,8 +33,8 @@
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.shared.system.ActivityManagerKt.isInForeground
+import com.android.systemui.statusbar.NotificationLockscreenUserManager
 import com.android.systemui.statusbar.StatusBarState
-import com.android.systemui.statusbar.phone.CentralSurfaces
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
@@ -45,9 +46,10 @@
  * the camera).
  */
 @SysUISingleton
-class CameraGestureHelper @Inject constructor(
+class CameraGestureHelper
+@Inject
+constructor(
     private val context: Context,
-    private val centralSurfaces: CentralSurfaces,
     private val keyguardStateController: KeyguardStateController,
     private val statusBarKeyguardViewManager: StatusBarKeyguardViewManager,
     private val packageManager: PackageManager,
@@ -59,24 +61,25 @@
     private val contentResolver: ContentResolver,
     @Main private val uiExecutor: Executor,
     private val selectedUserInteractor: SelectedUserInteractor,
+    private val devicePolicyManager: DevicePolicyManager,
+    private val lockscreenUserManager: NotificationLockscreenUserManager,
 ) {
-    /**
-     * Whether the camera application can be launched for the camera launch gesture.
-     */
+    /** Whether the camera application can be launched for the camera launch gesture. */
     fun canCameraGestureBeLaunched(statusBarState: Int): Boolean {
-        if (!centralSurfaces.isCameraAllowedByAdmin) {
+        if (!isCameraAllowedByAdmin()) {
             return false
         }
 
-        val resolveInfo: ResolveInfo? = packageManager.resolveActivityAsUser(
-            getStartCameraIntent(selectedUserInteractor.getSelectedUserId()),
-            PackageManager.MATCH_DEFAULT_ONLY,
-            selectedUserInteractor.getSelectedUserId()
-        )
+        val resolveInfo: ResolveInfo? =
+            packageManager.resolveActivityAsUser(
+                getStartCameraIntent(selectedUserInteractor.getSelectedUserId()),
+                PackageManager.MATCH_DEFAULT_ONLY,
+                selectedUserInteractor.getSelectedUserId()
+            )
         val resolvedPackage = resolveInfo?.activityInfo?.packageName
         return (resolvedPackage != null &&
-                (statusBarState != StatusBarState.SHADE ||
-                        !activityManager.isInForeground(resolvedPackage)))
+            (statusBarState != StatusBarState.SHADE ||
+                !activityManager.isInForeground(resolvedPackage)))
     }
 
     /**
@@ -87,9 +90,11 @@
     fun launchCamera(source: Int) {
         val intent: Intent = getStartCameraIntent(selectedUserInteractor.getSelectedUserId())
         intent.putExtra(CameraIntents.EXTRA_LAUNCH_SOURCE, source)
-        val wouldLaunchResolverActivity = activityIntentHelper.wouldLaunchResolverActivity(
-            intent, selectedUserInteractor.getSelectedUserId()
-        )
+        val wouldLaunchResolverActivity =
+            activityIntentHelper.wouldLaunchResolverActivity(
+                intent,
+                selectedUserInteractor.getSelectedUserId()
+            )
         if (CameraIntents.isSecureCameraIntent(intent) && !wouldLaunchResolverActivity) {
             uiExecutor.execute {
                 // Normally an activity will set its requested rotation animation on its window.
@@ -101,7 +106,7 @@
                 val activityOptions = ActivityOptions.makeBasic()
                 activityOptions.setDisallowEnterPictureInPictureWhileLaunching(true)
                 activityOptions.rotationAnimationHint =
-                        WindowManager.LayoutParams.ROTATION_ANIMATION_SEAMLESS
+                    WindowManager.LayoutParams.ROTATION_ANIMATION_SEAMLESS
                 try {
                     activityTaskManager.startActivityAsUser(
                         null,
@@ -118,11 +123,7 @@
                         selectedUserInteractor.getSelectedUserId(true),
                     )
                 } catch (e: RemoteException) {
-                    Log.w(
-                        "CameraGestureHelper",
-                        "Unable to start camera activity",
-                        e
-                    )
+                    Log.w("CameraGestureHelper", "Unable to start camera activity", e)
                 }
             }
         } else {
@@ -131,9 +132,6 @@
             activityStarter.startActivity(intent, false /* dismissShade */)
         }
 
-        // Call this to make sure that the keyguard returns if the app that is being launched
-        // crashes after a timeout.
-        centralSurfaces.startLaunchTransitionTimeout()
         // Call this to make sure the keyguard is ready to be dismissed once the next intent is
         // handled by the OS (in our case it is the activity we started right above)
         statusBarKeyguardViewManager.readyForKeyguardDone()
@@ -152,4 +150,17 @@
             cameraIntents.getInsecureCameraIntent(userId)
         }
     }
+
+    private fun isCameraAllowedByAdmin(): Boolean {
+        if (devicePolicyManager.getCameraDisabled(null, lockscreenUserManager.getCurrentUserId())) {
+            return false
+        } else if (keyguardStateController.isShowing() && statusBarKeyguardViewManager.isSecure()) {
+            // Check if the admin has disabled the camera specifically for the keyguard
+            return (devicePolicyManager.getKeyguardDisabledFeatures(
+                null,
+                lockscreenUserManager.getCurrentUserId()
+            ) and DevicePolicyManager.KEYGUARD_DISABLE_SECURE_CAMERA) == 0
+        }
+        return true
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt b/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt
index e0e1971..adb1ee2 100644
--- a/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt
@@ -25,6 +25,7 @@
 import javax.inject.Inject
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -74,6 +75,13 @@
         return onAnyConfigurationChange.mapLatest { repository.getDimensionPixelSize(resourceId) }
     }
 
+    /** Emits the dimensional pixel size of the given resource, inverting it for RTL if necessary */
+    fun directionalDimensionPixelSize(originLayoutDirection: Int, resourceId: Int): Flow<Int> {
+        return dimensionPixelSize(resourceId).combine(layoutDirection) { size, direction ->
+            if (originLayoutDirection == direction) size else -size
+        }
+    }
+
     /** Given a set of [resourceId]s, emit Map<ResourceId, DimensionPixelSize> on config change */
     fun dimensionPixelSize(resourceIds: Set<Int>): Flow<Map<Int, Int>> {
         return onAnyConfigurationChange.mapLatest {
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt
index e1d9bef..86241a5 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt
@@ -19,6 +19,7 @@
 import android.app.smartspace.SmartspaceTarget
 import android.os.Parcelable
 import androidx.annotation.VisibleForTesting
+import com.android.systemui.Flags.communalTimerFlickerFix
 import com.android.systemui.communal.data.model.CommunalSmartspaceTimer
 import com.android.systemui.communal.smartspace.CommunalSmartspaceController
 import com.android.systemui.dagger.SysUISingleton
@@ -80,7 +81,8 @@
                     // The view layer should have the instance based smartspaceTargetId instead of
                     // stable id, so that when a new instance of the timer is created, for example,
                     // when it is paused, the view should re-render its remote views.
-                    smartspaceTargetId = target.smartspaceTargetId,
+                    smartspaceTargetId =
+                        if (communalTimerFlickerFix()) stableId else target.smartspaceTargetId,
                     createdTimestampMillis = targetCreationTimes[stableId]!!,
                     remoteViews = target.remoteViews!!,
                 )
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
index 4be93cc..d1a5a4b 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/BaseCommunalViewModel.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.communal.ui.viewmodel
 
+import android.appwidget.AppWidgetProviderInfo
 import android.content.ComponentName
 import android.os.UserHandle
 import android.view.View
@@ -197,6 +198,9 @@
     /** Called as the user request to show the customize widget button. */
     open fun onLongClick() {}
 
+    /** Called as the UI determines that a new widget has been added to the grid. */
+    open fun onNewWidgetAdded(provider: AppWidgetProviderInfo) {}
+
     /** Called when the grid scroll position has been updated. */
     open fun onScrollPositionUpdated(firstVisibleItemIndex: Int, firstVisibleItemScroll: Int) {
         currentScrollIndex = firstVisibleItemIndex
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
index 5b825d8..1a86c71 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
@@ -19,16 +19,18 @@
 import android.appwidget.AppWidgetManager
 import android.appwidget.AppWidgetProviderInfo
 import android.content.ComponentName
+import android.content.Context
 import android.content.Intent
 import android.content.pm.PackageManager
 import android.content.res.Resources
 import android.os.UserHandle
 import android.util.Log
+import android.view.accessibility.AccessibilityEvent
+import android.view.accessibility.AccessibilityManager
 import androidx.activity.result.ActivityResultLauncher
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.communal.data.model.CommunalWidgetCategories
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
-import com.android.systemui.communal.domain.interactor.CommunalPrefsInteractor
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
 import com.android.systemui.communal.domain.interactor.CommunalSettingsInteractor
 import com.android.systemui.communal.domain.model.CommunalContentModel
@@ -37,6 +39,7 @@
 import com.android.systemui.communal.shared.model.EditModeState
 import com.android.systemui.communal.widgets.WidgetConfigurator
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
@@ -74,8 +77,10 @@
     private val uiEventLogger: UiEventLogger,
     @CommunalLog logBuffer: LogBuffer,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
-    private val communalPrefsInteractor: CommunalPrefsInteractor,
     private val metricsLogger: CommunalMetricsLogger,
+    @Application private val context: Context,
+    private val accessibilityManager: AccessibilityManager,
+    private val packageManager: PackageManager,
 ) : BaseCommunalViewModel(communalSceneInteractor, communalInteractor, mediaHost) {
 
     private val logger = Logger(logBuffer, "CommunalEditModeViewModel")
@@ -156,6 +161,25 @@
         uiEventLogger.log(CommunalUiEvent.COMMUNAL_HUB_REORDER_WIDGET_CANCEL)
     }
 
+    override fun onNewWidgetAdded(provider: AppWidgetProviderInfo) {
+        if (!accessibilityManager.isEnabled) {
+            return
+        }
+
+        // Send an accessibility announcement for the newly added widget
+        val widgetLabel = provider.loadLabel(packageManager)
+        val announcementText =
+            context.getString(
+                R.string.accessibility_announcement_communal_widget_added,
+                widgetLabel
+            )
+        accessibilityManager.sendAccessibilityEvent(
+            AccessibilityEvent(AccessibilityEvent.TYPE_ANNOUNCEMENT).apply {
+                contentDescription = announcementText
+            }
+        )
+    }
+
     val isIdleOnCommunal: StateFlow<Boolean> = communalInteractor.isIdleOnCommunal
 
     /** Launch the widget picker activity using the given {@link ActivityResultLauncher}. */
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index 1771f4d..15ddf5b 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -80,6 +80,7 @@
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.motiontool.MotionToolModule;
 import com.android.systemui.navigationbar.NavigationBarComponent;
+import com.android.systemui.navigationbar.gestural.dagger.GestureModule;
 import com.android.systemui.notetask.NoteTaskModule;
 import com.android.systemui.people.PeopleModule;
 import com.android.systemui.plugins.BcSmartspaceConfigPlugin;
@@ -215,6 +216,7 @@
         FlagsModule.class,
         FlagDependenciesModule.class,
         FooterActionsModule.class,
+        GestureModule.class,
         InputMethodModule.class,
         KeyEventRepositoryModule.class,
         KeyboardModule.class,
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
index 9823985..931066d 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
@@ -27,7 +27,9 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.content.pm.ActivityInfo;
 import android.graphics.drawable.ColorDrawable;
+import android.service.dreams.DreamActivity;
 import android.util.Log;
 import android.view.View;
 import android.view.ViewGroup;
@@ -63,6 +65,7 @@
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dreams.dagger.DreamOverlayComponent;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
+import com.android.systemui.navigationbar.gestural.domain.GestureInteractor;
 import com.android.systemui.shade.ShadeExpansionChangeEvent;
 import com.android.systemui.touch.TouchInsetManager;
 import com.android.systemui.util.concurrency.DelayableExecutor;
@@ -135,6 +138,8 @@
 
     private final DreamOverlayComponent mDreamOverlayComponent;
 
+    private ComponentName mCurrentBlockedGestureDreamActivityComponent;
+
     /**
      * This {@link LifecycleRegistry} controls when dream overlay functionality, like touch
      * handling, should be active. It will automatically be paused when the dream overlay is hidden
@@ -222,6 +227,8 @@
 
     private final DreamOverlayStateController mStateController;
 
+    private final GestureInteractor mGestureInteractor;
+
     @VisibleForTesting
     public enum DreamOverlayEvent implements UiEventLogger.UiEventEnum {
         @UiEvent(doc = "The dream overlay has entered start.")
@@ -265,6 +272,7 @@
             ComponentName homeControlPanelDreamComponent,
             DreamOverlayCallbackController dreamOverlayCallbackController,
             KeyguardInteractor keyguardInteractor,
+            GestureInteractor gestureInteractor,
             @Named(DREAM_OVERLAY_WINDOW_TITLE) String windowTitle) {
         super(executor);
         mContext = context;
@@ -281,6 +289,7 @@
         mWindowTitle = windowTitle;
         mCommunalInteractor = communalInteractor;
         mSystemDialogsCloser = systemDialogsCloser;
+        mGestureInteractor = gestureInteractor;
 
         final ViewModelStore viewModelStore = new ViewModelStore();
         final Complication.Host host =
@@ -391,6 +400,7 @@
         mStarted = true;
 
         updateRedirectWakeup();
+        updateBlockedGestureDreamActivityComponent();
     }
 
     private void updateRedirectWakeup() {
@@ -401,6 +411,18 @@
         redirectWake(mCommunalAvailable && !glanceableHubAllowKeyguardWhenDreaming());
     }
 
+    private void updateBlockedGestureDreamActivityComponent() {
+        // TODO(b/343815446): We should not be crafting this ActivityInfo ourselves. It should be
+        // in a common place, Such as DreamActivity itself.
+        final ActivityInfo info = new ActivityInfo();
+        info.name = DreamActivity.class.getName();
+        info.packageName = getDreamComponent().getPackageName();
+        mCurrentBlockedGestureDreamActivityComponent = info.getComponentName();
+
+        mGestureInteractor.addGestureBlockedActivity(mCurrentBlockedGestureDreamActivityComponent,
+                GestureInteractor.Scope.Global);
+    }
+
     @Override
     public void onEndDream() {
         resetCurrentDreamOverlayLocked();
@@ -472,6 +494,7 @@
      *                     into the dream window.
      */
     private boolean addOverlayWindowLocked(WindowManager.LayoutParams layoutParams) {
+
         mWindow = new PhoneWindow(mContext);
         // Default to SystemUI name for TalkBack.
         mWindow.setTitle(mWindowTitle);
@@ -554,6 +577,14 @@
         }
 
         mWindow = null;
+
+        // Always unregister the any set DreamActivity from being blocked from gestures.
+        if (mCurrentBlockedGestureDreamActivityComponent != null) {
+            mGestureInteractor.removeGestureBlockedActivity(
+                    mCurrentBlockedGestureDreamActivityComponent, GestureInteractor.Scope.Global);
+            mCurrentBlockedGestureDreamActivityComponent = null;
+        }
+
         mStarted = false;
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/touch/CommunalTouchHandler.java b/packages/SystemUI/src/com/android/systemui/dreams/touch/CommunalTouchHandler.java
index 04fda33..ee7b6f5 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/touch/CommunalTouchHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/touch/CommunalTouchHandler.java
@@ -20,6 +20,7 @@
 
 import android.graphics.Rect;
 import android.graphics.Region;
+import android.util.LayoutDirection;
 import android.view.GestureDetector;
 import android.view.MotionEvent;
 
@@ -27,6 +28,7 @@
 import androidx.lifecycle.Lifecycle;
 
 import com.android.systemui.ambient.touch.TouchHandler;
+import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor;
 import com.android.systemui.communal.domain.interactor.CommunalInteractor;
 import com.android.systemui.dreams.touch.dagger.CommunalTouchModule;
 import com.android.systemui.statusbar.phone.CentralSurfaces;
@@ -43,30 +45,42 @@
     private final Optional<CentralSurfaces> mCentralSurfaces;
     private final Lifecycle mLifecycle;
     private final CommunalInteractor mCommunalInteractor;
+
+    private final ConfigurationInteractor mConfigurationInteractor;
     private Boolean mIsEnabled = false;
 
+    private int mLayoutDirection = LayoutDirection.LTR;
+
     @VisibleForTesting
-    final Consumer<Boolean> mIsCommunalAvailableCallback =
-            isAvailable -> {
-                setIsEnabled(isAvailable);
-            };
+    final Consumer<Boolean> mIsCommunalAvailableCallback = isAvailable -> setIsEnabled(isAvailable);
+
+    @VisibleForTesting
+    final Consumer<Integer> mLayoutDirectionCallback = direction -> mLayoutDirection = direction;
 
     @Inject
     public CommunalTouchHandler(
             Optional<CentralSurfaces> centralSurfaces,
             @Named(CommunalTouchModule.COMMUNAL_GESTURE_INITIATION_WIDTH) int initiationWidth,
             CommunalInteractor communalInteractor,
+            ConfigurationInteractor configurationInteractor,
             Lifecycle lifecycle) {
         mInitiationWidth = initiationWidth;
         mCentralSurfaces = centralSurfaces;
         mLifecycle = lifecycle;
         mCommunalInteractor = communalInteractor;
+        mConfigurationInteractor = configurationInteractor;
 
         collectFlow(
                 mLifecycle,
                 mCommunalInteractor.isCommunalAvailable(),
                 mIsCommunalAvailableCallback
         );
+
+        collectFlow(
+                mLifecycle,
+                mConfigurationInteractor.getLayoutDirection(),
+                mLayoutDirectionCallback
+        );
     }
 
     @Override
@@ -90,7 +104,15 @@
     @Override
     public void getTouchInitiationRegion(Rect bounds, Region region, Rect exclusionRect) {
         final Rect outBounds = new Rect(bounds);
-        outBounds.inset(outBounds.width() - mInitiationWidth, 0, 0, 0);
+        final int inset = outBounds.width() - mInitiationWidth;
+
+        // Touch initiation area is defined in terms of LTR. The insets must be flipped for RTL
+        if (mLayoutDirection == LayoutDirection.LTR) {
+            outBounds.inset(inset, 0, 0, 0);
+        } else {
+            outBounds.inset(0, 0, inset, 0);
+        }
+
         region.op(outBounds, Region.Op.UNION);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt b/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt
index 0e06117..6e4038d 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt
@@ -40,6 +40,7 @@
 import com.android.systemui.statusbar.notification.shared.NotificationMinimalismPrototype
 import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.notification.shared.NotificationsLiveDataStoreRefactor
+import com.android.systemui.statusbar.notification.shared.NotificationThrottleHun
 import com.android.systemui.statusbar.notification.shared.PriorityPeopleSection
 import javax.inject.Inject
 
@@ -51,6 +52,7 @@
         // Internal notification backend dependencies
         crossAppPoliteNotifications dependsOn politeNotifications
         vibrateWhileUnlockedToken dependsOn politeNotifications
+        modesUi dependsOn modesApi
 
         // Internal notification frontend dependencies
         NotificationsLiveDataStoreRefactor.token dependsOn NotificationIconContainerRefactor.token
@@ -58,6 +60,7 @@
         NotificationAvalancheSuppression.token dependsOn VisualInterruptionRefactor.token
         PriorityPeopleSection.token dependsOn SortBySectionTimeFlag.token
         NotificationMinimalismPrototype.token dependsOn NotificationsHeadsUpRefactor.token
+        NotificationsHeadsUpRefactor.token dependsOn NotificationThrottleHun.token
 
         // SceneContainer dependencies
         SceneContainerFlag.getFlagDependencies().forEach { (alpha, beta) -> alpha dependsOn beta }
@@ -85,6 +88,12 @@
     private inline val vibrateWhileUnlockedToken: FlagToken
         get() = FlagToken(FLAG_VIBRATE_WHILE_UNLOCKED, vibrateWhileUnlocked())
 
+    private inline val modesUi
+        get() = FlagToken(android.app.Flags.FLAG_MODES_UI, android.app.Flags.modesUi())
+
+    private inline val modesApi
+        get() = FlagToken(android.app.Flags.FLAG_MODES_API, android.app.Flags.modesApi())
+
     private inline val communalHub
         get() = FlagToken(FLAG_COMMUNAL_HUB, communalHub())
 }
diff --git a/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt b/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt
index 491c73d..4652b2a 100644
--- a/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt
+++ b/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt
@@ -26,7 +26,9 @@
 import com.android.systemui.animation.DialogCuj
 import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.animation.Expandable
-import com.android.systemui.plugins.FalsingManager
+import com.android.systemui.log.LogBuffer
+import com.android.systemui.log.core.LogLevel
+import com.android.systemui.log.dagger.QSLog
 import com.android.systemui.plugins.qs.QSTile
 import com.android.systemui.statusbar.VibratorHelper
 import com.android.systemui.statusbar.policy.KeyguardStateController
@@ -48,7 +50,7 @@
 constructor(
     private val vibratorHelper: VibratorHelper?,
     private val keyguardStateController: KeyguardStateController,
-    private val falsingManager: FalsingManager,
+    @QSLog private val logBuffer: LogBuffer,
 ) {
 
     var effectDuration = 0
@@ -103,6 +105,7 @@
     }
 
     fun handleActionDown() {
+        logEvent(qsTile?.tileSpec, state, "action down received")
         when (state) {
             State.IDLE -> {
                 setState(State.TIMEOUT_WAIT)
@@ -114,6 +117,7 @@
     }
 
     fun handleActionUp() {
+        logEvent(qsTile?.tileSpec, state, "action up received")
         if (state == State.RUNNING_FORWARD) {
             setState(State.RUNNING_BACKWARDS_FROM_UP)
             callback?.onReverseAnimator()
@@ -132,6 +136,7 @@
     }
 
     fun handleAnimationStart() {
+        logEvent(qsTile?.tileSpec, state, "animation started")
         if (state == State.TIMEOUT_WAIT) {
             vibrate(longPressHint)
             setState(State.RUNNING_FORWARD)
@@ -140,6 +145,7 @@
 
     /** This function is called both when an animator completes or gets cancelled */
     fun handleAnimationComplete() {
+        logEvent(qsTile?.tileSpec, state, "animation completed")
         when (state) {
             State.RUNNING_FORWARD -> {
                 vibrate(snapEffect)
@@ -149,11 +155,13 @@
                     callback?.onResetProperties()
                     setState(State.IDLE)
                 }
+                logEvent(qsTile?.tileSpec, state, "long click action triggered")
                 qsTile?.longClick(expandable)
             }
             State.RUNNING_BACKWARDS_FROM_UP -> {
                 callback?.onEffectFinishedReversing()
                 setState(getStateForClick())
+                logEvent(qsTile?.tileSpec, state, "click action triggered")
                 qsTile?.click(expandable)
             }
             State.RUNNING_BACKWARDS_FROM_CANCEL -> {
@@ -181,6 +189,7 @@
         if (keyguardStateController.isPrimaryBouncerShowing || !isStateClickable) return false
 
         setState(getStateForClick())
+        logEvent(qsTile?.tileSpec, state, "click action triggered")
         qsTile?.click(expandable)
         return true
     }
@@ -195,11 +204,8 @@
     @VisibleForTesting
     fun getStateForClick(): State {
         val isTileUnavailable = qsTile?.state?.state == Tile.STATE_UNAVAILABLE
-        val isFalseTapWhileLocked =
-            !keyguardStateController.isUnlocked &&
-                falsingManager.isFalseTap(FalsingManager.LOW_PENALTY)
         val handlesLongClick = qsTile?.state?.handlesLongClick == true
-        return if (isTileUnavailable || isFalseTapWhileLocked || !handlesLongClick) {
+        return if (isTileUnavailable || !handlesLongClick || keyguardStateController.isShowing) {
             // The click event will not perform an action that resets the state. Therefore, this is
             // the last opportunity to reset the state back to IDLE.
             State.IDLE
@@ -278,6 +284,20 @@
         return delegated
     }
 
+    private fun logEvent(tileSpec: String?, state: State, event: String) {
+        if (!DEBUG) return
+        logBuffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = tileSpec
+                str2 = event
+                str3 = state.name
+            },
+            { "[long-press effect on $str1 tile] $str2 on state: $str3" }
+        )
+    }
+
     enum class State {
         IDLE, /* The effect is idle waiting for touch input */
         TIMEOUT_WAIT, /* The effect is waiting for a tap timeout period */
@@ -309,4 +329,9 @@
         /** Cancel the effect animator */
         fun onCancelAnimator()
     }
+
+    companion object {
+        private const val TAG = "QSLongPressEffect"
+        private const val DEBUG = true
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/docking/binder/KeyboardDockingIndicationViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyboard/docking/binder/KeyboardDockingIndicationViewBinder.kt
index f649be2..b859cdc 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/docking/binder/KeyboardDockingIndicationViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/docking/binder/KeyboardDockingIndicationViewBinder.kt
@@ -20,6 +20,7 @@
 import android.graphics.Paint
 import android.graphics.PixelFormat
 import android.view.WindowManager
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyboard.docking.ui.KeyboardDockingIndicationView
@@ -37,7 +38,7 @@
     context: Context,
     @Application private val applicationScope: CoroutineScope,
     private val viewModel: KeyboardDockingIndicationViewModel,
-    private val windowManager: WindowManager
+    private val windowManager: ViewCaptureAwareWindowManager,
 ) {
 
     private val windowLayoutParams =
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
index be64ff6..af755d3 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
@@ -384,16 +384,7 @@
         if (index > 0) {
             HorizontalDivider()
         }
-        ShortcutSinglePane(searchQuery, shortcut)
-    }
-}
-
-@Composable
-private fun ShortcutSinglePane(searchQuery: String, shortcut: Shortcut) {
-    Column(Modifier.padding(vertical = 24.dp)) {
-        ShortcutDescriptionText(searchQuery = searchQuery, shortcut = shortcut)
-        Spacer(modifier = Modifier.height(12.dp))
-        ShortcutKeyCombinations(shortcut = shortcut)
+        ShortcutView(Modifier.padding(vertical = 24.dp), searchQuery, shortcut)
     }
 }
 
@@ -421,7 +412,7 @@
                 onCategoryClicked = { onCategorySelected(it.type) }
             )
             Spacer(modifier = Modifier.width(24.dp))
-            EndSidePanel(searchQuery, Modifier.fillMaxSize(), selectedCategory)
+            EndSidePanel(searchQuery, Modifier.fillMaxSize().padding(top = 8.dp), selectedCategory)
         }
     }
 }
@@ -447,14 +438,14 @@
         shape = RoundedCornerShape(28.dp),
         color = MaterialTheme.colorScheme.surfaceBright
     ) {
-        Column(Modifier.padding(horizontal = 32.dp, vertical = 24.dp)) {
+        Column(Modifier.padding(24.dp)) {
             SubCategoryTitle(subCategory.label)
-            Spacer(Modifier.height(24.dp))
+            Spacer(Modifier.height(8.dp))
             subCategory.shortcuts.fastForEachIndexed { index, shortcut ->
                 if (index > 0) {
                     HorizontalDivider()
                 }
-                ShortcutViewDualPane(searchQuery, shortcut)
+                ShortcutView(Modifier.padding(vertical = 16.dp), searchQuery, shortcut)
             }
         }
     }
@@ -470,17 +461,17 @@
 }
 
 @Composable
-private fun ShortcutViewDualPane(searchQuery: String, shortcut: Shortcut) {
-    Row(Modifier.padding(vertical = 16.dp)) {
+private fun ShortcutView(modifier: Modifier, searchQuery: String, shortcut: Shortcut) {
+    Row(modifier) {
         Row(
-            modifier = Modifier.width(160.dp).align(Alignment.CenterVertically),
+            modifier = Modifier.width(128.dp).align(Alignment.CenterVertically),
             horizontalArrangement = Arrangement.spacedBy(16.dp),
             verticalAlignment = Alignment.CenterVertically,
         ) {
             if (shortcut.icon != null) {
                 ShortcutIcon(
                     shortcut.icon,
-                    modifier = Modifier.size(36.dp),
+                    modifier = Modifier.size(24.dp),
                 )
             }
             ShortcutDescriptionText(
@@ -520,7 +511,11 @@
     modifier: Modifier = Modifier,
     shortcut: Shortcut,
 ) {
-    FlowRow(modifier = modifier, verticalArrangement = Arrangement.spacedBy(8.dp)) {
+    FlowRow(
+        modifier = modifier,
+        verticalArrangement = Arrangement.spacedBy(8.dp),
+        horizontalArrangement = Arrangement.End
+    ) {
         shortcut.commands.forEachIndexed { index, command ->
             if (index > 0) {
                 ShortcutOrSeparator(spacing = 16.dp)
@@ -641,7 +636,7 @@
 ) {
     Column(modifier) {
         ShortcutsSearchBar(onSearchQueryChanged)
-        Spacer(modifier = Modifier.heightIn(16.dp))
+        Spacer(modifier = Modifier.heightIn(8.dp))
         CategoriesPanelTwoPane(categories, selectedCategory, onCategoryClicked)
         Spacer(modifier = Modifier.weight(1f))
         KeyboardSettings(onKeyboardSettingsClicked)
@@ -678,7 +673,7 @@
     Surface(
         selected = selected,
         onClick = onClick,
-        modifier = Modifier.semantics { role = Role.Tab }.heightIn(min = 72.dp).fillMaxWidth(),
+        modifier = Modifier.semantics { role = Role.Tab }.heightIn(min = 64.dp).fillMaxWidth(),
         shape = RoundedCornerShape(28.dp),
         color = colors.containerColor(selected).value,
     ) {
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
index b8d0c23..2039743 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
@@ -26,8 +26,10 @@
 import androidx.activity.BackEventCompat
 import androidx.activity.ComponentActivity
 import androidx.activity.OnBackPressedCallback
+import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.platform.ComposeView
+import androidx.compose.ui.platform.LocalContext
 import androidx.core.view.updatePadding
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import androidx.lifecycle.flowWithLifecycle
@@ -36,6 +38,7 @@
 import com.android.systemui.keyboard.shortcut.ui.composable.ShortcutHelper
 import com.android.systemui.keyboard.shortcut.ui.viewmodel.ShortcutHelperViewModel
 import com.android.systemui.res.R
+import com.android.systemui.settings.UserTracker
 import com.google.android.material.bottomsheet.BottomSheetBehavior
 import com.google.android.material.bottomsheet.BottomSheetBehavior.BottomSheetCallback
 import com.google.android.material.bottomsheet.BottomSheetBehavior.STATE_HIDDEN
@@ -49,6 +52,7 @@
 class ShortcutHelperActivity
 @Inject
 constructor(
+    private val userTracker: UserTracker,
     private val viewModel: ShortcutHelperViewModel,
 ) : ComponentActivity() {
 
@@ -79,13 +83,16 @@
     private fun setUpComposeView() {
         requireViewById<ComposeView>(R.id.shortcut_helper_compose_container).apply {
             setContent {
-                PlatformTheme {
-                    val shortcutsUiState by viewModel.shortcutsUiState.collectAsStateWithLifecycle()
-                    ShortcutHelper(
-                        shortcutsUiState = shortcutsUiState,
-                        onKeyboardSettingsClicked = ::onKeyboardSettingsClicked,
-                        onSearchQueryChanged = { viewModel.onSearchQueryChanged(it) },
-                    )
+                CompositionLocalProvider(LocalContext provides userTracker.userContext) {
+                    PlatformTheme {
+                        val shortcutsUiState by
+                            viewModel.shortcutsUiState.collectAsStateWithLifecycle()
+                        ShortcutHelper(
+                            shortcutsUiState = shortcutsUiState,
+                            onKeyboardSettingsClicked = ::onKeyboardSettingsClicked,
+                            onSearchQueryChanged = { viewModel.onSearchQueryChanged(it) },
+                        )
+                    }
                 }
             }
         }
@@ -93,7 +100,10 @@
 
     private fun onKeyboardSettingsClicked() {
         try {
-            startActivity(Intent(Settings.ACTION_HARD_KEYBOARD_SETTINGS))
+            startActivityAsUser(
+                Intent(Settings.ACTION_HARD_KEYBOARD_SETTINGS),
+                userTracker.userHandle
+            )
         } catch (e: ActivityNotFoundException) {
             // From the Settings docs: In some cases, a matching Activity may not exist, so ensure
             // you safeguard against this.
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
index e64cc80..19b46e3 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyboard.shortcut.ui.viewmodel
 
+import android.app.role.RoleManager
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperCategoriesInteractor
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperStateInteractor
@@ -25,6 +26,7 @@
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType.CurrentApp
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutSubCategory
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutsUiState
+import com.android.systemui.settings.UserTracker
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
@@ -35,10 +37,13 @@
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.withContext
 
 class ShortcutHelperViewModel
 @Inject
 constructor(
+    private val roleManager: RoleManager,
+    private val userTracker: UserTracker,
     @Background private val backgroundScope: CoroutineScope,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
     private val stateInteractor: ShortcutHelperStateInteractor,
@@ -72,13 +77,22 @@
                 initialValue = ShortcutsUiState.Inactive
             )
 
-    private fun getDefaultSelectedCategory(
+    private suspend fun getDefaultSelectedCategory(
         categories: List<ShortcutCategory>
     ): ShortcutCategoryType? {
-        val currentAppShortcuts = categories.firstOrNull { it.type is CurrentApp }
+        val currentAppShortcuts =
+            categories.firstOrNull { it.type is CurrentApp && !isAppLauncher(it.type.packageName) }
         return currentAppShortcuts?.type ?: categories.firstOrNull()?.type
     }
 
+    private suspend fun isAppLauncher(packageName: String): Boolean {
+        return withContext(backgroundDispatcher) {
+            roleManager
+                .getRoleHoldersAsUser(RoleManager.ROLE_HOME, userTracker.userHandle)
+                .firstOrNull() == packageName
+        }
+    }
+
     private fun filterCategoriesBySearchQuery(
         query: String,
         categories: List<ShortcutCategory>
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/DismissCallbackRegistry.java b/packages/SystemUI/src/com/android/systemui/keyguard/DismissCallbackRegistry.java
index d1bbc33..0d01ee1 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/DismissCallbackRegistry.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/DismissCallbackRegistry.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard;
 
+import android.util.Log;
 import com.android.internal.policy.IKeyguardDismissCallback;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.UiBackground;
@@ -33,6 +34,7 @@
 
     private final ArrayList<DismissCallbackWrapper> mDismissCallbacks = new ArrayList<>();
     private final Executor mUiBgExecutor;
+    private final static String TAG = "DismissCallbackRegistry";
 
     @Inject
     public DismissCallbackRegistry(@UiBackground Executor uiBgExecutor) {
@@ -40,10 +42,12 @@
     }
 
     public void addCallback(IKeyguardDismissCallback callback) {
+        Log.d(TAG, "Adding callback: " + callback);
         mDismissCallbacks.add(new DismissCallbackWrapper(callback));
     }
 
     public void notifyDismissCancelled() {
+        Log.d(TAG, "notifyDismissCancelled(" + mDismissCallbacks.size() + ")");
         for (int i = mDismissCallbacks.size() - 1; i >= 0; i--) {
             DismissCallbackWrapper callback = mDismissCallbacks.get(i);
             mUiBgExecutor.execute(callback::notifyDismissCancelled);
@@ -52,6 +56,7 @@
     }
 
     public void notifyDismissSucceeded() {
+        Log.d(TAG, "notifyDismissSucceeded(" + mDismissCallbacks.size() + ")");
         for (int i = mDismissCallbacks.size() - 1; i >= 0; i--) {
             DismissCallbackWrapper callback = mDismissCallbacks.get(i);
             mUiBgExecutor.execute(callback::notifyDismissSucceeded);
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
index ae751db..edf17c1 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
@@ -34,6 +34,7 @@
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
 import com.android.systemui.keyguard.shared.model.BiometricUnlockModel
 import com.android.systemui.keyguard.shared.model.BiometricUnlockSource
+import com.android.systemui.keyguard.shared.model.CameraLaunchSourceModel
 import com.android.systemui.keyguard.shared.model.DismissAction
 import com.android.systemui.keyguard.shared.model.DozeStateModel
 import com.android.systemui.keyguard.shared.model.DozeTransitionModel
@@ -237,6 +238,9 @@
     /** Observable updated when keyguardDone should be called either now or soon. */
     val keyguardDone: Flow<KeyguardDone>
 
+    /** Last camera launch detection event */
+    val onCameraLaunchDetected: MutableStateFlow<CameraLaunchSourceModel>
+
     /**
      * Emits after the keyguard is done animating away.
      *
@@ -380,6 +384,8 @@
     private val _keyguardAlpha = MutableStateFlow(1f)
     override val keyguardAlpha = _keyguardAlpha.asStateFlow()
 
+    override val onCameraLaunchDetected = MutableStateFlow(CameraLaunchSourceModel())
+
     override val panelAlpha: MutableStateFlow<Float> = MutableStateFlow(1f)
 
     private val _clockShouldBeCentered = MutableStateFlow(true)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt
index a915241..ae830ee 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt
@@ -198,7 +198,12 @@
             interpolator = Interpolators.LINEAR
             duration =
                 when (toState) {
+                    KeyguardState.AOD -> TO_AOD_DURATION
+                    KeyguardState.DOZING -> TO_DOZING_DURATION
                     KeyguardState.GONE -> TO_GONE_DURATION
+                    KeyguardState.LOCKSCREEN -> TO_LOCKSCREEN_DURATION
+                    KeyguardState.OCCLUDED -> TO_OCCLUDED_DURATION
+                    KeyguardState.PRIMARY_BOUNCER -> TO_PRIMARY_BOUNCER_DURATION
                     else -> TRANSITION_DURATION_MS
                 }.inWholeMilliseconds
         }
@@ -211,10 +216,11 @@
     companion object {
         const val TAG = "FromAlternateBouncerTransitionInteractor"
         val TRANSITION_DURATION_MS = 300.milliseconds
-        val TO_GONE_DURATION = 500.milliseconds
         val TO_AOD_DURATION = TRANSITION_DURATION_MS
-        val TO_PRIMARY_BOUNCER_DURATION = TRANSITION_DURATION_MS
         val TO_DOZING_DURATION = TRANSITION_DURATION_MS
+        val TO_GONE_DURATION = 500.milliseconds
+        val TO_LOCKSCREEN_DURATION = 300.milliseconds
         val TO_OCCLUDED_DURATION = TRANSITION_DURATION_MS
+        val TO_PRIMARY_BOUNCER_DURATION = TRANSITION_DURATION_MS
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
index 046e79c..42490c4 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -23,10 +23,7 @@
 import android.graphics.Point
 import android.util.MathUtils
 import com.android.app.animation.Interpolators
-import com.android.app.tracing.FlowTracing.tracedAwaitClose
-import com.android.app.tracing.FlowTracing.tracedConflatedCallbackFlow
 import com.android.systemui.bouncer.data.repository.KeyguardBouncerRepository
-import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
 import com.android.systemui.common.shared.model.NotificationContainerBounds
 import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.dagger.SysUISingleton
@@ -35,9 +32,11 @@
 import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.keyguard.shared.model.BiometricUnlockModel
 import com.android.systemui.keyguard.shared.model.CameraLaunchSourceModel
+import com.android.systemui.keyguard.shared.model.CameraLaunchType
 import com.android.systemui.keyguard.shared.model.DozeStateModel
 import com.android.systemui.keyguard.shared.model.DozeStateModel.Companion.isDozeOff
 import com.android.systemui.keyguard.shared.model.DozeTransitionModel
+import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
 import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
@@ -48,11 +47,8 @@
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.data.repository.ShadeRepository
-import com.android.systemui.statusbar.CommandQueue
-import com.android.systemui.statusbar.notification.NotificationUtils.interpolate
 import com.android.systemui.statusbar.notification.stack.domain.interactor.SharedNotificationContainerInteractor
 import com.android.systemui.util.kotlin.Utils.Companion.sample as sampleCombine
-import com.android.systemui.util.kotlin.pairwise
 import com.android.systemui.util.kotlin.sample
 import javax.inject.Inject
 import javax.inject.Provider
@@ -69,9 +65,7 @@
 import kotlinx.coroutines.flow.debounce
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
-import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flatMapLatest
-import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.merge
@@ -87,7 +81,6 @@
 @Inject
 constructor(
     private val repository: KeyguardRepository,
-    private val commandQueue: CommandQueue,
     powerInteractor: PowerInteractor,
     bouncerRepository: KeyguardBouncerRepository,
     configurationInteractor: ConfigurationInteractor,
@@ -102,55 +95,34 @@
     // TODO(b/296118689): move to a repository
     private val _notificationPlaceholderBounds = MutableStateFlow(NotificationContainerBounds())
 
-    // When going to AOD, we interpolate bounds when receiving the new bounds
-    // When going back to LS, we'll apply new bounds directly
-    private val _nonSplitShadeNotifciationPlaceholderBounds =
-        _notificationPlaceholderBounds.pairwise().flatMapLatest { (oldBounds, newBounds) ->
-            val lastChangeStep = keyguardTransitionInteractor.transitionState.first()
-            if (lastChangeStep.to == AOD) {
-                keyguardTransitionInteractor.transitionState.map { step ->
-                    val startingProgress = lastChangeStep.value
-                    val progress = step.value
-                    if (step.to == AOD && progress >= startingProgress && startingProgress < 1f) {
-                        val adjustedProgress =
-                            ((progress - startingProgress) / (1F - startingProgress)).coerceIn(
-                                0F,
-                                1F
-                            )
-                        val top = interpolate(oldBounds.top, newBounds.top, adjustedProgress)
-                        val bottom =
-                            interpolate(
-                                oldBounds.bottom,
-                                newBounds.bottom,
-                                adjustedProgress.coerceIn(0F, 1F)
-                            )
-                        NotificationContainerBounds(top = top, bottom = bottom)
-                    } else {
-                        newBounds
-                    }
-                }
-            } else {
-                flow { emit(newBounds) }
-            }
-        }
-
     /** Bounds of the notification container. */
     val notificationContainerBounds: StateFlow<NotificationContainerBounds> by lazy {
         SceneContainerFlag.assertInLegacyMode()
-        combine(
+        combineTransform(
                 _notificationPlaceholderBounds,
-                _nonSplitShadeNotifciationPlaceholderBounds,
                 sharedNotificationContainerInteractor.get().configurationBasedDimensions,
-            ) { bounds, nonSplitShadeBounds: NotificationContainerBounds, cfg ->
+                keyguardTransitionInteractor.isInTransition(
+                    edge = Edge.create(from = LOCKSCREEN, to = AOD)
+                ),
+            ) { bounds, cfg, isTransitioningToAod ->
+                if (isTransitioningToAod) {
+                    // Keep bounds stable during this transition, to prevent cases like smartspace
+                    // popping in and adjusting the bounds. A prime example would be media playing,
+                    // which then updates smartspace on transition to AOD
+                    return@combineTransform
+                }
+
                 // We offset the placeholder bounds by the configured top margin to account for
                 // legacy placement behavior within notifications for splitshade.
-                if (MigrateClocksToBlueprint.isEnabled) {
-                    if (cfg.useSplitShade) {
-                        bounds.copy(bottom = bounds.bottom - cfg.keyguardSplitShadeTopMargin)
-                    } else {
-                        nonSplitShadeBounds
-                    }
-                } else bounds
+                emit(
+                    if (MigrateClocksToBlueprint.isEnabled) {
+                        if (cfg.useSplitShade) {
+                            bounds.copy(bottom = bounds.bottom - cfg.keyguardSplitShadeTopMargin)
+                        } else {
+                            bounds
+                        }
+                    } else bounds
+                )
             }
             .stateIn(
                 scope = applicationScope,
@@ -198,22 +170,7 @@
 
     /** Event for when the camera gesture is detected */
     val onCameraLaunchDetected: Flow<CameraLaunchSourceModel> =
-        tracedConflatedCallbackFlow("KeyguardInteractor#onCameraLaunchDetected") {
-            val callback =
-                object : CommandQueue.Callbacks {
-                    override fun onCameraLaunchGestureDetected(source: Int) {
-                        trySendWithFailureLogging(
-                            cameraLaunchSourceIntToModel(source),
-                            TAG,
-                            "updated onCameraLaunchGestureDetected"
-                        )
-                    }
-                }
-
-            commandQueue.addCallback(callback)
-
-            tracedAwaitClose("onCameraLaunchDetected") { commandQueue.removeCallback(callback) }
-        }
+        repository.onCameraLaunchDetected.filter { it.type != CameraLaunchType.IGNORE }
 
     /**
      * Dozing and dreaming have overlapping events. If the doze state remains in FINISH, it means
@@ -310,7 +267,7 @@
                 when {
                     isKeyguardVisible -> false
                     isPrimaryBouncerShowing -> false
-                    else -> cameraLaunchEvent == CameraLaunchSourceModel.POWER_DOUBLE_TAP
+                    else -> cameraLaunchEvent.type == CameraLaunchType.POWER_DOUBLE_TAP
                 }
             }
             .onStart { emit(false) }
@@ -440,16 +397,15 @@
         return repository.isKeyguardShowing()
     }
 
-    private fun cameraLaunchSourceIntToModel(value: Int): CameraLaunchSourceModel {
+    private fun cameraLaunchSourceIntToType(value: Int): CameraLaunchType {
         return when (value) {
-            StatusBarManager.CAMERA_LAUNCH_SOURCE_WIGGLE -> CameraLaunchSourceModel.WIGGLE
+            StatusBarManager.CAMERA_LAUNCH_SOURCE_WIGGLE -> CameraLaunchType.WIGGLE
             StatusBarManager.CAMERA_LAUNCH_SOURCE_POWER_DOUBLE_TAP ->
-                CameraLaunchSourceModel.POWER_DOUBLE_TAP
-            StatusBarManager.CAMERA_LAUNCH_SOURCE_LIFT_TRIGGER ->
-                CameraLaunchSourceModel.LIFT_TRIGGER
+                CameraLaunchType.POWER_DOUBLE_TAP
+            StatusBarManager.CAMERA_LAUNCH_SOURCE_LIFT_TRIGGER -> CameraLaunchType.LIFT_TRIGGER
             StatusBarManager.CAMERA_LAUNCH_SOURCE_QUICK_AFFORDANCE ->
-                CameraLaunchSourceModel.QUICK_AFFORDANCE
-            else -> throw IllegalArgumentException("Invalid CameraLaunchSourceModel value: $value")
+                CameraLaunchType.QUICK_AFFORDANCE
+            else -> throw IllegalArgumentException("Invalid CameraLaunchType value: $value")
         }
     }
 
@@ -508,6 +464,11 @@
         fromLockscreenTransitionInteractor.get().dismissKeyguard()
     }
 
+    fun onCameraLaunchDetected(source: Int) {
+        repository.onCameraLaunchDetected.value =
+            CameraLaunchSourceModel(type = cameraLaunchSourceIntToType(source))
+    }
+
     companion object {
         private const val TAG = "KeyguardInteractor"
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
index ccce3bf..31236a4 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
@@ -25,6 +25,7 @@
 import com.android.app.tracing.coroutines.withContext
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.internal.widget.LockPatternUtils
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.animation.Expandable
 import com.android.systemui.dagger.SysUISingleton
@@ -80,7 +81,8 @@
     private val featureFlags: FeatureFlags,
     private val repository: Lazy<KeyguardQuickAffordanceRepository>,
     private val launchAnimator: DialogTransitionAnimator,
-    private val logger: KeyguardQuickAffordancesMetricsLogger,
+    private val logger: KeyguardQuickAffordancesLogger,
+    private val metricsLogger: KeyguardQuickAffordancesMetricsLogger,
     private val devicePolicyManager: DevicePolicyManager,
     private val dockManager: DockManager,
     private val biometricSettingsRepository: BiometricSettingsRepository,
@@ -171,7 +173,8 @@
             Log.e(TAG, "Affordance config with key of \"$configKey\" not found!")
             return
         }
-        logger.logOnShortcutTriggered(slotId, configKey)
+        logger.logQuickAffordanceTriggered(decodedSlotId, decodedConfigKey)
+        metricsLogger.logOnShortcutTriggered(slotId, configKey)
 
         when (val result = config.onTriggered(expandable)) {
             is KeyguardQuickAffordanceConfig.OnTriggeredResult.StartActivity ->
@@ -223,7 +226,8 @@
                 affordanceIds = selections,
             )
 
-        logger.logOnShortcutSelected(slotId, affordanceId)
+        logger.logQuickAffordanceSelected(slotId, affordanceId)
+        metricsLogger.logOnShortcutSelected(slotId, affordanceId)
         return true
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/CameraLaunchSourceModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/CameraLaunchSourceModel.kt
index 19baf77..c017651 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/CameraLaunchSourceModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/CameraLaunchSourceModel.kt
@@ -15,14 +15,8 @@
  */
 package com.android.systemui.keyguard.shared.model
 
-/** Camera launch sources */
-enum class CameraLaunchSourceModel {
-    /** Device is wiggled */
-    WIGGLE,
-    /** Power button has been double tapped */
-    POWER_DOUBLE_TAP,
-    /** Device has been lifted */
-    LIFT_TRIGGER,
-    /** Quick affordance button has been pressed */
-    QUICK_AFFORDANCE,
-}
+/** Camera launch source, with type and time detected */
+data class CameraLaunchSourceModel(
+    val type: CameraLaunchType = CameraLaunchType.IGNORE,
+    val detectedTime: Long = System.currentTimeMillis(),
+)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/CameraLaunchType.kt b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/CameraLaunchType.kt
new file mode 100644
index 0000000..984abbb
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/CameraLaunchType.kt
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2024 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.keyguard.shared.model
+
+/** Camera launch sources */
+enum class CameraLaunchType {
+    /** Models no value */
+    IGNORE,
+    /** Device is wiggled */
+    WIGGLE,
+    /** Power button has been double tapped */
+    POWER_DOUBLE_TAP,
+    /** Device has been lifted */
+    LIFT_TRIGGER,
+    /** Quick affordance button has been pressed */
+    QUICK_AFFORDANCE,
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinder.kt
index db33acb..a250b22 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinder.kt
@@ -37,6 +37,7 @@
 import com.android.systemui.deviceentry.ui.binder.UdfpsAccessibilityOverlayBinder
 import com.android.systemui.deviceentry.ui.view.UdfpsAccessibilityOverlay
 import com.android.systemui.deviceentry.ui.viewmodel.AlternateBouncerUdfpsAccessibilityOverlayViewModel
+import com.android.systemui.keyguard.DismissCallbackRegistry
 import com.android.systemui.keyguard.ui.view.DeviceEntryIconView
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerDependencies
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerUdfpsIconViewModel
@@ -66,6 +67,7 @@
     private val alternateBouncerDependencies: Lazy<AlternateBouncerDependencies>,
     private val windowManager: Lazy<WindowManager>,
     private val layoutInflater: Lazy<LayoutInflater>,
+    private val dismissCallbackRegistry: DismissCallbackRegistry,
 ) : CoreStartable {
     private val layoutParams: WindowManager.LayoutParams
         get() =
@@ -162,6 +164,7 @@
 
             fun onBackRequested() {
                 alternateBouncerDependencies.get().viewModel.hideAlternateBouncer()
+                dismissCallbackRegistry.notifyDismissCancelled()
             }
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/DeviceEntryIconViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/DeviceEntryIconViewBinder.kt
index 1b9788f..4d6577c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/DeviceEntryIconViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/DeviceEntryIconViewBinder.kt
@@ -74,15 +74,26 @@
         val bgView = view.bgView
         longPressHandlingView.listener =
             object : LongPressHandlingView.Listener {
-                override fun onLongPressDetected(view: View, x: Int, y: Int, isA11yAction: Boolean) {
-                    if (!isA11yAction && falsingManager.isFalseLongTap(FalsingManager.LOW_PENALTY)) {
+                override fun onLongPressDetected(
+                    view: View,
+                    x: Int,
+                    y: Int,
+                    isA11yAction: Boolean
+                ) {
+                    if (
+                        !isA11yAction && falsingManager.isFalseLongTap(FalsingManager.LOW_PENALTY)
+                    ) {
                         return
                     }
                     vibratorHelper.performHapticFeedback(
                         view,
                         HapticFeedbackConstants.CONFIRM,
                     )
-                    applicationScope.launch { viewModel.onUserInteraction() }
+                    applicationScope.launch {
+                        view.clearFocus()
+                        view.clearAccessibilityFocus()
+                        viewModel.onUserInteraction()
+                    }
                 }
             }
 
@@ -95,6 +106,7 @@
                     launch("$TAG#viewModel.isVisible") {
                         viewModel.isVisible.collect { isVisible ->
                             longPressHandlingView.isInvisible = !isVisible
+                            view.isClickable = isVisible
                         }
                     }
                     launch("$TAG#viewModel.isLongPressEnabled") {
@@ -131,7 +143,11 @@
                                         view,
                                         HapticFeedbackConstants.CONFIRM,
                                     )
-                                    applicationScope.launch { viewModel.onUserInteraction() }
+                                    applicationScope.launch {
+                                        view.clearFocus()
+                                        view.clearAccessibilityFocus()
+                                        viewModel.onUserInteraction()
+                                    }
                                 }
                             } else {
                                 view.setOnClickListener(null)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardLongPressViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardLongPressViewBinder.kt
index b387855..830ef3b 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardLongPressViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardLongPressViewBinder.kt
@@ -46,6 +46,7 @@
         onSingleTap: () -> Unit,
         falsingManager: FalsingManager,
     ) {
+        view.contentDescription = view.resources.getString(R.string.accessibility_desc_lock_screen)
         view.accessibilityHintLongPressAction =
             AccessibilityNodeInfo.AccessibilityAction(
                 AccessibilityNodeInfoCompat.ACTION_LONG_CLICK,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardQuickAffordanceViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardQuickAffordanceViewBinder.kt
index b9a79dc..162a0d2 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardQuickAffordanceViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardQuickAffordanceViewBinder.kt
@@ -30,6 +30,7 @@
 import androidx.core.view.updateLayoutParams
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.repeatOnLifecycle
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.settingslib.Utils
 import com.android.systemui.animation.Expandable
 import com.android.systemui.animation.view.LaunchableImageView
@@ -74,6 +75,7 @@
         alpha: Flow<Float>,
         falsingManager: FalsingManager?,
         vibratorHelper: VibratorHelper?,
+        logger: KeyguardQuickAffordancesLogger,
         messageDisplayer: (Int) -> Unit,
     ): Binding {
         val button = view as ImageView
@@ -89,6 +91,7 @@
                                 falsingManager = falsingManager,
                                 messageDisplayer = messageDisplayer,
                                 vibratorHelper = vibratorHelper,
+                                logger = logger,
                             )
                         }
                     }
@@ -131,6 +134,7 @@
         falsingManager: FalsingManager?,
         messageDisplayer: (Int) -> Unit,
         vibratorHelper: VibratorHelper?,
+        logger: KeyguardQuickAffordancesLogger,
     ) {
         if (!viewModel.isVisible) {
             view.isInvisible = true
@@ -228,6 +232,7 @@
                     shakeAnimator.start()
 
                     vibratorHelper?.vibrate(KeyguardBottomAreaVibrations.Shake)
+                    logger.logQuickAffordanceTapped(viewModel.configKey)
                 }
                 view.onLongClickListener =
                     OnLongClickListener(falsingManager, viewModel, vibratorHelper, onTouchListener)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
index bc5b7b9..6faca1e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
@@ -50,6 +50,7 @@
 import com.android.internal.policy.SystemBarUtils
 import com.android.keyguard.ClockEventController
 import com.android.keyguard.KeyguardClockSwitch
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.animation.view.LaunchableImageView
 import com.android.systemui.biometrics.domain.interactor.UdfpsOverlayInteractor
 import com.android.systemui.broadcast.BroadcastDispatcher
@@ -147,6 +148,7 @@
     private val defaultShortcutsSection: DefaultShortcutsSection,
     private val keyguardClockInteractor: KeyguardClockInteractor,
     private val keyguardClockViewModel: KeyguardClockViewModel,
+    private val quickAffordancesLogger: KeyguardQuickAffordancesLogger,
 ) {
     val hostToken: IBinder? = bundle.getBinder(KEY_HOST_TOKEN)
     private val width: Int = bundle.getInt(KEY_VIEW_WIDTH)
@@ -462,6 +464,7 @@
                     alpha = flowOf(1f),
                     falsingManager = falsingManager,
                     vibratorHelper = vibratorHelper,
+                    logger = quickAffordancesLogger,
                 ) { message ->
                     indicationController.showTransientIndication(message)
                 }
@@ -476,6 +479,7 @@
                     alpha = flowOf(1f),
                     falsingManager = falsingManager,
                     vibratorHelper = vibratorHelper,
+                    logger = quickAffordancesLogger,
                 ) { message ->
                     indicationController.showTransientIndication(message)
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt
index dc7a649..0032c2f 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt
@@ -18,6 +18,7 @@
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerToAodTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerToDozingTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerToGoneTransitionViewModel
+import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerToLockscreenTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerToOccludedTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerToPrimaryBouncerTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.AodToGoneTransitionViewModel
@@ -80,6 +81,12 @@
 
     @Binds
     @IntoSet
+    abstract fun alternateBouncerToLockscreen(
+        impl: AlternateBouncerToLockscreenTransitionViewModel
+    ): DeviceEntryIconTransition
+
+    @Binds
+    @IntoSet
     abstract fun alternateBouncerToOccluded(
         impl: AlternateBouncerToOccludedTransitionViewModel
     ): DeviceEntryIconTransition
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/AlignShortcutsToUdfpsSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/AlignShortcutsToUdfpsSection.kt
index 2e96638..1ba830b 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/AlignShortcutsToUdfpsSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/AlignShortcutsToUdfpsSection.kt
@@ -25,6 +25,7 @@
 import androidx.constraintlayout.widget.ConstraintSet.PARENT_ID
 import androidx.constraintlayout.widget.ConstraintSet.RIGHT
 import androidx.constraintlayout.widget.ConstraintSet.TOP
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.deviceentry.shared.DeviceEntryUdfpsRefactor
 import com.android.systemui.keyguard.KeyguardBottomAreaRefactor
@@ -47,6 +48,7 @@
     private val falsingManager: FalsingManager,
     private val indicationController: KeyguardIndicationController,
     private val vibratorHelper: VibratorHelper,
+    private val shortcutsLogger: KeyguardQuickAffordancesLogger,
 ) : BaseShortcutSection() {
     override fun addViews(constraintLayout: ConstraintLayout) {
         if (KeyguardBottomAreaRefactor.isEnabled) {
@@ -64,6 +66,7 @@
                     keyguardQuickAffordancesCombinedViewModel.transitionAlpha,
                     falsingManager,
                     vibratorHelper,
+                    shortcutsLogger,
                 ) {
                     indicationController.showTransientIndication(it)
                 }
@@ -74,6 +77,7 @@
                     keyguardQuickAffordancesCombinedViewModel.transitionAlpha,
                     falsingManager,
                     vibratorHelper,
+                    shortcutsLogger,
                 ) {
                     indicationController.showTransientIndication(it)
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt
index 9146c60..64c46db 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt
@@ -26,6 +26,7 @@
 import androidx.constraintlayout.widget.ConstraintSet.PARENT_ID
 import androidx.constraintlayout.widget.ConstraintSet.RIGHT
 import androidx.constraintlayout.widget.ConstraintSet.VISIBILITY_MODE_IGNORE
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.animation.view.LaunchableImageView
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.keyguard.KeyguardBottomAreaRefactor
@@ -52,6 +53,7 @@
     private val indicationController: KeyguardIndicationController,
     private val vibratorHelper: VibratorHelper,
     private val keyguardBlueprintInteractor: Lazy<KeyguardBlueprintInteractor>,
+    private val shortcutsLogger: KeyguardQuickAffordancesLogger,
 ) : BaseShortcutSection() {
 
     // Amount to increase the bottom margin by to avoid colliding with inset
@@ -86,6 +88,7 @@
                     keyguardQuickAffordancesCombinedViewModel.transitionAlpha,
                     falsingManager,
                     vibratorHelper,
+                    shortcutsLogger,
                 ) {
                     indicationController.showTransientIndication(it)
                 }
@@ -96,6 +99,7 @@
                     keyguardQuickAffordancesCombinedViewModel.transitionAlpha,
                     falsingManager,
                     vibratorHelper,
+                    shortcutsLogger,
                 ) {
                     indicationController.showTransientIndication(it)
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModel.kt
new file mode 100644
index 0000000..b04521c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModel.kt
@@ -0,0 +1,62 @@
+/*
+ * Copyright (C) 2024 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.keyguard.ui.viewmodel
+
+import android.util.MathUtils
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.keyguard.domain.interactor.FromAlternateBouncerTransitionInteractor
+import com.android.systemui.keyguard.shared.model.Edge
+import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
+import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
+import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
+import com.android.systemui.keyguard.ui.transitions.DeviceEntryIconTransition
+import javax.inject.Inject
+import kotlin.time.Duration.Companion.milliseconds
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
+
+/**
+ * Breaks down ALTERNATE_BOUNCER->LOCKSCREEN transition into discrete steps for corresponding views
+ * to consume.
+ */
+@ExperimentalCoroutinesApi
+@SysUISingleton
+class AlternateBouncerToLockscreenTransitionViewModel
+@Inject
+constructor(
+    animationFlow: KeyguardTransitionAnimationFlow,
+) : DeviceEntryIconTransition {
+    private val transitionAnimation =
+        animationFlow.setup(
+            duration = FromAlternateBouncerTransitionInteractor.TO_LOCKSCREEN_DURATION,
+            edge = Edge.create(from = ALTERNATE_BOUNCER, to = LOCKSCREEN),
+        )
+
+    fun lockscreenAlpha(viewState: ViewStateAccessor): Flow<Float> {
+        var startAlpha = 1f
+        return transitionAnimation.sharedFlow(
+            duration = 250.milliseconds,
+            onStart = { startAlpha = viewState.alpha() },
+            onStep = { MathUtils.lerp(startAlpha, 1f, it) },
+        )
+    }
+
+    val deviceEntryBackgroundViewAlpha: Flow<Float> =
+        transitionAnimation.immediatelyTransitionTo(1f)
+    override val deviceEntryParentViewAlpha: Flow<Float> =
+        transitionAnimation.immediatelyTransitionTo(1f)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModel.kt
index 00aa102..ea8fe29 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DreamingToGlanceableHubTransitionViewModel.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.util.LayoutDirection
 import com.android.app.animation.Interpolators.EMPHASIZED
 import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.dagger.SysUISingleton
@@ -54,7 +55,10 @@
 
     val dreamOverlayTranslationX: Flow<Float> =
         configurationInteractor
-            .dimensionPixelSize(R.dimen.dreaming_to_hub_transition_dream_overlay_translation_x)
+            .directionalDimensionPixelSize(
+                LayoutDirection.LTR,
+                R.dimen.dreaming_to_hub_transition_dream_overlay_translation_x
+            )
             .flatMapLatest { translatePx ->
                 transitionAnimation.sharedFlow(
                     duration = TO_GLANCEABLE_HUB_DURATION,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModel.kt
index d594488..76d5a8d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToDreamingTransitionViewModel.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.util.LayoutDirection
 import com.android.app.animation.Interpolators
 import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.dagger.SysUISingleton
@@ -63,7 +64,10 @@
 
     val dreamOverlayTranslationX: Flow<Float> =
         configurationInteractor
-            .dimensionPixelSize(R.dimen.hub_to_dreaming_transition_dream_overlay_translation_x)
+            .directionalDimensionPixelSize(
+                LayoutDirection.LTR,
+                R.dimen.hub_to_dreaming_transition_dream_overlay_translation_x
+            )
             .flatMapLatest { translatePx: Int ->
                 transitionAnimation.sharedFlow(
                     duration = FROM_GLANCEABLE_HUB_DURATION,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModel.kt
index 046b95f..67b009e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/GlanceableHubToLockscreenTransitionViewModel.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.util.LayoutDirection
 import com.android.app.animation.Interpolators.EMPHASIZED
 import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.dagger.SysUISingleton
@@ -72,7 +73,10 @@
 
     val keyguardTranslationX: Flow<StateToValue> =
         configurationInteractor
-            .dimensionPixelSize(R.dimen.hub_to_lockscreen_transition_lockscreen_translation_x)
+            .directionalDimensionPixelSize(
+                LayoutDirection.LTR,
+                R.dimen.hub_to_lockscreen_transition_lockscreen_translation_x
+            )
             .flatMapLatest { translatePx: Int ->
                 transitionAnimation.sharedFlowWithState(
                     duration = TO_LOCKSCREEN_DURATION,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
index 350ceb4..11889c5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
@@ -85,6 +85,8 @@
     private val notificationsKeyguardInteractor: NotificationsKeyguardInteractor,
     private val alternateBouncerToGoneTransitionViewModel:
         AlternateBouncerToGoneTransitionViewModel,
+    private val alternateBouncerToLockscreenTransitionViewModel:
+        AlternateBouncerToLockscreenTransitionViewModel,
     private val aodToGoneTransitionViewModel: AodToGoneTransitionViewModel,
     private val aodToLockscreenTransitionViewModel: AodToLockscreenTransitionViewModel,
     private val aodToOccludedTransitionViewModel: AodToOccludedTransitionViewModel,
@@ -238,6 +240,7 @@
                         alphaOnShadeExpansion,
                         keyguardInteractor.dismissAlpha,
                         alternateBouncerToGoneTransitionViewModel.lockscreenAlpha(viewState),
+                        alternateBouncerToLockscreenTransitionViewModel.lockscreenAlpha(viewState),
                         aodToGoneTransitionViewModel.lockscreenAlpha(viewState),
                         aodToLockscreenTransitionViewModel.lockscreenAlpha(viewState),
                         aodToOccludedTransitionViewModel.lockscreenAlpha(viewState),
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModel.kt
index c7273b7..378374e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToGlanceableHubTransitionViewModel.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import android.util.LayoutDirection
 import com.android.app.animation.Interpolators.EMPHASIZED
 import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.dagger.SysUISingleton
@@ -71,7 +72,10 @@
 
     val keyguardTranslationX: Flow<StateToValue> =
         configurationInteractor
-            .dimensionPixelSize(R.dimen.lockscreen_to_hub_transition_lockscreen_translation_x)
+            .directionalDimensionPixelSize(
+                LayoutDirection.LTR,
+                R.dimen.lockscreen_to_hub_transition_lockscreen_translation_x
+            )
             .flatMapLatest { translatePx: Int ->
                 transitionAnimation.sharedFlowWithState(
                     duration = FromLockscreenTransitionInteractor.TO_GLANCEABLE_HUB_DURATION,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt b/packages/SystemUI/src/com/android/systemui/log/dagger/KeyguardQuickAffordancesLog.kt
similarity index 63%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
copy to packages/SystemUI/src/com/android/systemui/log/dagger/KeyguardQuickAffordancesLog.kt
index 37c9552..e9cf7e2 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/KeyguardQuickAffordancesLog.kt
@@ -14,11 +14,12 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.domain.interactor
+package com.android.systemui.log.dagger
 
-import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
+import javax.inject.Qualifier
 
-val Kosmos.partitionedGridLayout by
-    Kosmos.Fixture { PartitionedGridLayout(partitionedGridViewModel) }
+/** A [com.android.systemui.log.LogBuffer] for keyguard quick affordances related stuff. */
+@Qualifier
+@MustBeDocumented
+@Retention(AnnotationRetention.RUNTIME)
+annotation class KeyguardQuickAffordancesLog
diff --git a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
index b2ba0e1..40bb8e1 100644
--- a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
@@ -568,6 +568,16 @@
     }
 
     /**
+     * Provides a {@link LogBuffer} for keyguard quick affordances-related logs.
+     */
+    @Provides
+    @SysUISingleton
+    @KeyguardQuickAffordancesLog
+    public static LogBuffer provideKeyguardQuickAffordancesLogBuffer(LogBufferFactory factory) {
+        return factory.create("KeyguardQuickAffordancesLog", 25);
+    }
+
+    /**
      * Provides a {@link LogBuffer} for keyguard transition animation logs.
      */
     @Provides
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
index 947336d..9eca34f 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
@@ -76,6 +76,7 @@
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.navigationbar.NavigationModeController;
+import com.android.systemui.navigationbar.gestural.domain.GestureInteractor;
 import com.android.systemui.plugins.FalsingManager;
 import com.android.systemui.plugins.NavigationEdgeBackPlugin;
 import com.android.systemui.plugins.PluginListener;
@@ -95,15 +96,14 @@
 import com.android.systemui.statusbar.phone.LightBarController;
 import com.android.systemui.util.concurrency.BackPanelUiThread;
 import com.android.systemui.util.concurrency.UiThreadContext;
+import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.wm.shell.back.BackAnimation;
 import com.android.wm.shell.desktopmode.DesktopMode;
 import com.android.wm.shell.pip.Pip;
 
 import java.io.PrintWriter;
 import java.util.ArrayDeque;
-import java.util.ArrayList;
 import java.util.Date;
-import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
@@ -157,12 +157,7 @@
     private TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
         @Override
         public void onTaskStackChanged() {
-            if (edgebackGestureHandlerGetRunningTasksBackground()) {
-                mBackgroundExecutor.execute(() -> mGestureBlockingActivityRunning.set(
-                        isGestureBlockingActivityRunning()));
-            } else {
-                mGestureBlockingActivityRunning.set(isGestureBlockingActivityRunning());
-            }
+            updateRunningActivityGesturesBlocked();
         }
         @Override
         public void onTaskCreated(int taskId, ComponentName componentName) {
@@ -209,8 +204,6 @@
     private final Optional<DesktopMode> mDesktopModeOptional;
     private final FalsingManager mFalsingManager;
     private final Configuration mLastReportedConfig = new Configuration();
-    // Activities which should not trigger Back gesture.
-    private final List<ComponentName> mGestureBlockingActivities = new ArrayList<>();
 
     private final Point mDisplaySize = new Point();
     private final int mDisplayId;
@@ -227,6 +220,10 @@
             mBackGestureTfClassifierProviderProvider;
     private final Provider<LightBarController> mLightBarControllerProvider;
 
+    private final GestureInteractor mGestureInteractor;
+
+    private final JavaAdapter mJavaAdapter;
+
     // The left side edge width where touch down is allowed
     private int mEdgeWidthLeft;
     // The right side edge width where touch down is allowed
@@ -426,7 +423,9 @@
             FalsingManager falsingManager,
             Provider<BackGestureTfClassifierProvider> backGestureTfClassifierProviderProvider,
             Provider<LightBarController> lightBarControllerProvider,
-            NotificationShadeWindowController notificationShadeWindowController) {
+            NotificationShadeWindowController notificationShadeWindowController,
+            GestureInteractor gestureInteractor,
+            JavaAdapter javaAdapter) {
         mContext = context;
         mDisplayId = context.getDisplayId();
         mUiThreadContext = uiThreadContext;
@@ -446,7 +445,13 @@
         mFalsingManager = falsingManager;
         mBackGestureTfClassifierProviderProvider = backGestureTfClassifierProviderProvider;
         mLightBarControllerProvider = lightBarControllerProvider;
+        mGestureInteractor = gestureInteractor;
+        mJavaAdapter = javaAdapter;
         mLastReportedConfig.setTo(mContext.getResources().getConfiguration());
+
+        mJavaAdapter.alwaysCollectFlow(mGestureInteractor.getGestureBlockedActivities(),
+                componentNames -> updateRunningActivityGesturesBlocked());
+
         ComponentName recentsComponentName = ComponentName.unflattenFromString(
                 context.getString(com.android.internal.R.string.config_recentsComponentName));
         if (recentsComponentName != null) {
@@ -466,8 +471,9 @@
                 } else {
                     String[] gestureBlockingActivities = resources.getStringArray(resId);
                     for (String gestureBlockingActivity : gestureBlockingActivities) {
-                        mGestureBlockingActivities.add(
-                                ComponentName.unflattenFromString(gestureBlockingActivity));
+                        mGestureInteractor.addGestureBlockedActivity(
+                                ComponentName.unflattenFromString(gestureBlockingActivity),
+                                GestureInteractor.Scope.Local);
                     }
                 }
             } catch (NameNotFoundException e) {
@@ -561,6 +567,15 @@
         }
     }
 
+    private void updateRunningActivityGesturesBlocked() {
+        if (edgebackGestureHandlerGetRunningTasksBackground()) {
+            mBackgroundExecutor.execute(() -> mGestureBlockingActivityRunning.set(
+                    isGestureBlockingActivityRunning()));
+        } else {
+            mGestureBlockingActivityRunning.set(isGestureBlockingActivityRunning());
+        }
+    }
+
     /**
      * Called when the nav/task bar is attached.
      */
@@ -1293,7 +1308,8 @@
         } else {
             mPackageName = "_UNKNOWN";
         }
-        return topActivity != null && mGestureBlockingActivities.contains(topActivity);
+
+        return topActivity != null && mGestureInteractor.areGesturesBlocked(topActivity);
     }
 
     public void setBackAnimation(BackAnimation backAnimation) {
@@ -1342,6 +1358,10 @@
         private final Provider<LightBarController> mLightBarControllerProvider;
         private final NotificationShadeWindowController mNotificationShadeWindowController;
 
+        private final GestureInteractor mGestureInteractor;
+
+        private final JavaAdapter mJavaAdapter;
+
         @Inject
         public Factory(OverviewProxyService overviewProxyService,
                         SysUiState sysUiState,
@@ -1361,8 +1381,10 @@
                         FalsingManager falsingManager,
                         Provider<BackGestureTfClassifierProvider>
                                 backGestureTfClassifierProviderProvider,
-                Provider<LightBarController> lightBarControllerProvider,
-                NotificationShadeWindowController notificationShadeWindowController) {
+                        Provider<LightBarController> lightBarControllerProvider,
+                        NotificationShadeWindowController notificationShadeWindowController,
+                        GestureInteractor gestureInteractor,
+                        JavaAdapter javaAdapter) {
             mOverviewProxyService = overviewProxyService;
             mSysUiState = sysUiState;
             mPluginManager = pluginManager;
@@ -1382,6 +1404,8 @@
             mBackGestureTfClassifierProviderProvider = backGestureTfClassifierProviderProvider;
             mLightBarControllerProvider = lightBarControllerProvider;
             mNotificationShadeWindowController = notificationShadeWindowController;
+            mGestureInteractor = gestureInteractor;
+            mJavaAdapter = javaAdapter;
         }
 
         /** Construct a {@link EdgeBackGestureHandler}. */
@@ -1407,7 +1431,9 @@
                             mFalsingManager,
                             mBackGestureTfClassifierProviderProvider,
                             mLightBarControllerProvider,
-                            mNotificationShadeWindowController));
+                            mNotificationShadeWindowController,
+                            mGestureInteractor,
+                            mJavaAdapter));
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/dagger/GestureModule.kt b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/dagger/GestureModule.kt
new file mode 100644
index 0000000..72a84f5
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/dagger/GestureModule.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 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.navigationbar.gestural.dagger
+
+import com.android.systemui.navigationbar.gestural.data.respository.GestureRepository
+import com.android.systemui.navigationbar.gestural.data.respository.GestureRepositoryImpl
+import dagger.Binds
+import dagger.Module
+
+/** {@link Module} for gesture related dependencies */
+@Module
+interface GestureModule {
+    /**  */
+    @Binds fun gestureRespoitory(impl: GestureRepositoryImpl): GestureRepository
+}
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/data/respository/GestureRepository.kt b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/data/respository/GestureRepository.kt
new file mode 100644
index 0000000..8f35343
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/data/respository/GestureRepository.kt
@@ -0,0 +1,63 @@
+/*
+ * Copyright (C) 2024 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.navigationbar.gestural.data.respository
+
+import android.content.ComponentName
+import android.util.ArraySet
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Main
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.withContext
+
+/** A repository for storing gesture related information */
+interface GestureRepository {
+    /** A {@link StateFlow} tracking activities currently blocked from gestures. */
+    val gestureBlockedActivities: StateFlow<Set<ComponentName>>
+
+    /** Adds an activity to be blocked from gestures. */
+    suspend fun addGestureBlockedActivity(activity: ComponentName)
+
+    /** Removes an activity from being blocked from gestures. */
+    suspend fun removeGestureBlockedActivity(activity: ComponentName)
+}
+
+@SysUISingleton
+class GestureRepositoryImpl
+@Inject
+constructor(@Main private val mainDispatcher: CoroutineDispatcher) : GestureRepository {
+    private val _gestureBlockedActivities = MutableStateFlow<Set<ComponentName>>(ArraySet())
+
+    override val gestureBlockedActivities: StateFlow<Set<ComponentName>>
+        get() = _gestureBlockedActivities
+
+    override suspend fun addGestureBlockedActivity(activity: ComponentName) =
+        withContext(mainDispatcher) {
+            _gestureBlockedActivities.emit(
+                _gestureBlockedActivities.value.toMutableSet().apply { add(activity) }
+            )
+        }
+
+    override suspend fun removeGestureBlockedActivity(activity: ComponentName) =
+        withContext(mainDispatcher) {
+            _gestureBlockedActivities.emit(
+                _gestureBlockedActivities.value.toMutableSet().apply { remove(activity) }
+            )
+        }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/GestureInteractor.kt b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/GestureInteractor.kt
new file mode 100644
index 0000000..6dc5939
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/GestureInteractor.kt
@@ -0,0 +1,103 @@
+/*
+ * Copyright (C) 2024 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.navigationbar.gestural.domain
+
+import android.content.ComponentName
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.navigationbar.gestural.data.respository.GestureRepository
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+
+/**
+ * {@link GestureInteractor} helps interact with gesture-related logic, including accessing the
+ * underlying {@link GestureRepository}.
+ */
+class GestureInteractor
+@Inject
+constructor(
+    private val gestureRepository: GestureRepository,
+    @Application private val scope: CoroutineScope
+) {
+    enum class Scope {
+        Local,
+        Global
+    }
+
+    private val _localGestureBlockedActivities = MutableStateFlow<Set<ComponentName>>(setOf())
+    /** A {@link StateFlow} for listening to changes in Activities where gestures are blocked */
+    val gestureBlockedActivities: StateFlow<Set<ComponentName>>
+        get() =
+            combine(
+                    gestureRepository.gestureBlockedActivities,
+                    _localGestureBlockedActivities.asStateFlow()
+                ) { global, local ->
+                    global + local
+                }
+                .stateIn(scope, SharingStarted.WhileSubscribed(), setOf())
+
+    /**
+     * Adds an {@link Activity} to be blocked based on component when the topmost, focused {@link
+     * Activity}.
+     */
+    fun addGestureBlockedActivity(activity: ComponentName, gestureScope: Scope) {
+        scope.launch {
+            when (gestureScope) {
+                Scope.Local -> {
+                    _localGestureBlockedActivities.emit(
+                        _localGestureBlockedActivities.value.toMutableSet().apply { add(activity) }
+                    )
+                }
+                Scope.Global -> {
+                    gestureRepository.addGestureBlockedActivity(activity)
+                }
+            }
+        }
+    }
+
+    /** Removes an {@link Activity} from being blocked from gestures. */
+    fun removeGestureBlockedActivity(activity: ComponentName, gestureScope: Scope) {
+        scope.launch {
+            when (gestureScope) {
+                Scope.Local -> {
+                    _localGestureBlockedActivities.emit(
+                        _localGestureBlockedActivities.value.toMutableSet().apply {
+                            remove(activity)
+                        }
+                    )
+                }
+                Scope.Global -> {
+                    gestureRepository.removeGestureBlockedActivity(activity)
+                }
+            }
+        }
+    }
+
+    /**
+     * Checks whether the specified {@link Activity} {@link ComponentName} is being blocked from
+     * gestures.
+     */
+    fun areGesturesBlocked(activity: ComponentName): Boolean {
+        return gestureBlockedActivities.value.contains(activity)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/power/domain/interactor/PowerInteractor.kt b/packages/SystemUI/src/com/android/systemui/power/domain/interactor/PowerInteractor.kt
index 9380d44..8d48c1d 100644
--- a/packages/SystemUI/src/com/android/systemui/power/domain/interactor/PowerInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/power/domain/interactor/PowerInteractor.kt
@@ -18,6 +18,7 @@
 package com.android.systemui.power.domain.interactor
 
 import android.os.PowerManager
+import com.android.systemui.camera.CameraGestureHelper
 import com.android.systemui.classifier.FalsingCollector
 import com.android.systemui.classifier.FalsingCollectorActual
 import com.android.systemui.dagger.SysUISingleton
@@ -28,6 +29,7 @@
 import com.android.systemui.power.shared.model.WakefulnessState
 import com.android.systemui.statusbar.phone.ScreenOffAnimationController
 import javax.inject.Inject
+import javax.inject.Provider
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
@@ -41,6 +43,7 @@
     @FalsingCollectorActual private val falsingCollector: FalsingCollector,
     private val screenOffAnimationController: ScreenOffAnimationController,
     private val statusBarStateController: StatusBarStateController,
+    private val cameraGestureHelper: Provider<CameraGestureHelper>,
 ) {
     /** Whether the screen is on or off. */
     val isInteractive: Flow<Boolean> = repository.isInteractive
@@ -206,7 +209,13 @@
     }
 
     fun onCameraLaunchGestureDetected() {
-        repository.updateWakefulness(powerButtonLaunchGestureTriggered = true)
+        if (
+            cameraGestureHelper
+                .get()
+                .canCameraGestureBeLaunched(statusBarStateController.getState())
+        ) {
+            repository.updateWakefulness(powerButtonLaunchGestureTriggered = true)
+        }
     }
 
     companion object {
diff --git a/packages/SystemUI/src/com/android/systemui/power/shared/model/WakefulnessModel.kt b/packages/SystemUI/src/com/android/systemui/power/shared/model/WakefulnessModel.kt
index 5432793d..0f49c94 100644
--- a/packages/SystemUI/src/com/android/systemui/power/shared/model/WakefulnessModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/power/shared/model/WakefulnessModel.kt
@@ -19,7 +19,7 @@
      * all use cases. If you need more granular information about a waking/sleeping transition, use
      * the [KeyguardTransitionInteractor].
      */
-    internal val internalWakefulnessState: WakefulnessState = WakefulnessState.AWAKE,
+    val internalWakefulnessState: WakefulnessState = WakefulnessState.AWAKE,
     val lastWakeReason: WakeSleepReason = WakeSleepReason.OTHER,
     val lastSleepReason: WakeSleepReason = WakeSleepReason.OTHER,
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
index 78f4b4b..072d322 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/dagger/PanelsModule.kt
@@ -30,14 +30,10 @@
 import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
 import com.android.systemui.qs.panels.shared.model.PaginatedGridLayoutType
 import com.android.systemui.qs.panels.shared.model.PanelsLog
-import com.android.systemui.qs.panels.shared.model.PartitionedGridLayoutType
-import com.android.systemui.qs.panels.shared.model.StretchedGridLayoutType
 import com.android.systemui.qs.panels.ui.compose.GridLayout
 import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
 import com.android.systemui.qs.panels.ui.compose.PaginatableGridLayout
 import com.android.systemui.qs.panels.ui.compose.PaginatedGridLayout
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.compose.StretchedGridLayout
 import com.android.systemui.qs.panels.ui.viewmodel.FixedColumnsSizeViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.FixedColumnsSizeViewModelImpl
 import com.android.systemui.qs.panels.ui.viewmodel.IconLabelVisibilityViewModel
@@ -102,22 +98,6 @@
 
         @Provides
         @IntoSet
-        fun provideStretchedGridLayout(
-            gridLayout: StretchedGridLayout
-        ): Pair<GridLayoutType, GridLayout> {
-            return Pair(StretchedGridLayoutType, gridLayout)
-        }
-
-        @Provides
-        @IntoSet
-        fun providePartitionedGridLayout(
-            gridLayout: PartitionedGridLayout
-        ): Pair<GridLayoutType, GridLayout> {
-            return Pair(PartitionedGridLayoutType, gridLayout)
-        }
-
-        @Provides
-        @IntoSet
         fun providePaginatedGridLayout(
             gridLayout: PaginatedGridLayout
         ): Pair<GridLayoutType, GridLayout> {
@@ -148,22 +128,6 @@
 
         @Provides
         @IntoSet
-        fun provideStretchedGridConsistencyInteractor(
-            consistencyInteractor: NoopGridConsistencyInteractor
-        ): Pair<GridLayoutType, GridTypeConsistencyInteractor> {
-            return Pair(StretchedGridLayoutType, consistencyInteractor)
-        }
-
-        @Provides
-        @IntoSet
-        fun providePartitionedGridConsistencyInteractor(
-            consistencyInteractor: NoopGridConsistencyInteractor
-        ): Pair<GridLayoutType, GridTypeConsistencyInteractor> {
-            return Pair(PartitionedGridLayoutType, consistencyInteractor)
-        }
-
-        @Provides
-        @IntoSet
         fun providePaginatedGridConsistencyInteractor(
             @PaginatedBaseLayoutType consistencyInteractor: GridTypeConsistencyInteractor,
         ): Pair<GridLayoutType, GridTypeConsistencyInteractor> {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutType.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutType.kt
index b1942fe..323f39b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutType.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/GridLayoutType.kt
@@ -26,14 +26,5 @@
 /** Grid type representing a scrollable vertical grid. */
 data object InfiniteGridLayoutType : GridLayoutType
 
-/**
- * Grid type representing a scrollable vertical grid where tiles will stretch to fill in empty
- * spaces.
- */
-data object StretchedGridLayoutType : GridLayoutType
-
-/** Grid type grouping large tiles on top and icon tiles at the bottom. */
-data object PartitionedGridLayoutType : GridLayoutType
-
 /** Grid type for a paginated list of tiles. It will delegate to some other layout type. */
 data object PaginatedGridLayoutType : GridLayoutType
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/PartitionedGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/PartitionedGridLayout.kt
deleted file mode 100644
index 6c84edd..0000000
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/PartitionedGridLayout.kt
+++ /dev/null
@@ -1,404 +0,0 @@
-/*
- * Copyright (C) 2024 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.qs.panels.ui.compose
-
-import androidx.compose.foundation.layout.Arrangement
-import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.Row
-import androidx.compose.foundation.layout.Spacer
-import androidx.compose.foundation.layout.fillMaxSize
-import androidx.compose.foundation.layout.fillMaxWidth
-import androidx.compose.foundation.layout.height
-import androidx.compose.foundation.layout.padding
-import androidx.compose.foundation.lazy.grid.GridCells
-import androidx.compose.foundation.lazy.grid.GridItemSpan
-import androidx.compose.foundation.lazy.grid.LazyGridScope
-import androidx.compose.foundation.rememberScrollState
-import androidx.compose.foundation.shape.RoundedCornerShape
-import androidx.compose.foundation.verticalScroll
-import androidx.compose.material3.MaterialTheme
-import androidx.compose.material3.Switch
-import androidx.compose.material3.Text
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.DisposableEffect
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.rememberUpdatedState
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.drawWithContent
-import androidx.compose.ui.graphics.Color
-import androidx.compose.ui.graphics.Path
-import androidx.compose.ui.graphics.PathEffect
-import androidx.compose.ui.graphics.Shape
-import androidx.compose.ui.graphics.addOutline
-import androidx.compose.ui.graphics.drawscope.Stroke
-import androidx.compose.ui.res.dimensionResource
-import androidx.compose.ui.text.font.FontWeight
-import androidx.compose.ui.unit.Dp
-import androidx.compose.ui.unit.dp
-import androidx.lifecycle.compose.collectAsStateWithLifecycle
-import com.android.compose.modifiers.background
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.qs.panels.shared.model.SizedTile
-import com.android.systemui.qs.panels.ui.viewmodel.EditTileViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.PartitionedGridViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
-import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor
-import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.android.systemui.res.R
-import javax.inject.Inject
-
-@SysUISingleton
-class PartitionedGridLayout @Inject constructor(private val viewModel: PartitionedGridViewModel) :
-    PaginatableGridLayout {
-    @Composable
-    override fun TileGrid(
-        tiles: List<TileViewModel>,
-        modifier: Modifier,
-        editModeStart: () -> Unit,
-    ) {
-        DisposableEffect(tiles) {
-            val token = Any()
-            tiles.forEach { it.startListening(token) }
-            onDispose { tiles.forEach { it.stopListening(token) } }
-        }
-        val columns by viewModel.columns.collectAsStateWithLifecycle()
-        val showLabels by viewModel.showLabels.collectAsStateWithLifecycle()
-        val largeTileHeight = tileHeight()
-        val iconTileHeight = tileHeight(showLabels)
-        val (smallTiles, largeTiles) = tiles.partition { viewModel.isIconTile(it.spec) }
-
-        TileLazyGrid(modifier = modifier, columns = GridCells.Fixed(columns)) {
-            // Large tiles
-            items(largeTiles.size, span = { GridItemSpan(2) }) { index ->
-                Tile(
-                    tile = largeTiles[index],
-                    iconOnly = false,
-                    modifier = Modifier.height(largeTileHeight)
-                )
-            }
-            fillUpRow(nTiles = largeTiles.size, columns = columns / 2)
-
-            // Small tiles
-            items(smallTiles.size) { index ->
-                Tile(
-                    tile = smallTiles[index],
-                    iconOnly = true,
-                    showLabels = showLabels,
-                    modifier = Modifier.height(iconTileHeight)
-                )
-            }
-        }
-    }
-
-    @Composable
-    override fun EditTileGrid(
-        tiles: List<EditTileViewModel>,
-        modifier: Modifier,
-        onAddTile: (TileSpec, Int) -> Unit,
-        onRemoveTile: (TileSpec) -> Unit,
-    ) {
-        val columns by viewModel.columns.collectAsStateWithLifecycle()
-        val showLabels by viewModel.showLabels.collectAsStateWithLifecycle()
-
-        val listState = rememberEditListState(tiles)
-        val dragAndDropState = rememberDragAndDropState(listState)
-
-        val (currentTiles, otherTiles) = listState.tiles.partition { it.isCurrent }
-        val addTileToEnd: (TileSpec) -> Unit by rememberUpdatedState {
-            onAddTile(it, CurrentTilesInteractor.POSITION_AT_END)
-        }
-        val onDoubleTap: (TileSpec) -> Unit by rememberUpdatedState { tileSpec ->
-            viewModel.resize(tileSpec, !viewModel.isIconTile(tileSpec))
-        }
-        val largeTileHeight = tileHeight()
-        val iconTileHeight = tileHeight(showLabels)
-        val tilePadding = dimensionResource(R.dimen.qs_tile_margin_vertical)
-
-        Column(
-            verticalArrangement = Arrangement.spacedBy(tilePadding),
-            modifier = modifier.fillMaxSize().verticalScroll(rememberScrollState())
-        ) {
-            Row(
-                modifier =
-                    Modifier.background(
-                            color = MaterialTheme.colorScheme.surfaceVariant,
-                            alpha = { 1f },
-                            shape = RoundedCornerShape(dimensionResource(R.dimen.qs_corner_radius))
-                        )
-                        .padding(tilePadding)
-            ) {
-                Column(Modifier.padding(start = tilePadding)) {
-                    Text(
-                        text = "Show text labels",
-                        color = MaterialTheme.colorScheme.onBackground,
-                        fontWeight = FontWeight.Bold
-                    )
-                    Text(
-                        text = "Display names under each tile",
-                        color = MaterialTheme.colorScheme.onBackground
-                    )
-                }
-                Spacer(modifier = Modifier.weight(1f))
-                Switch(checked = showLabels, onCheckedChange = { viewModel.setShowLabels(it) })
-            }
-
-            CurrentTiles(
-                tiles = currentTiles,
-                largeTileHeight = largeTileHeight,
-                iconTileHeight = iconTileHeight,
-                tilePadding = tilePadding,
-                onAdd = onAddTile,
-                onRemove = onRemoveTile,
-                onDoubleTap = onDoubleTap,
-                isIconOnly = viewModel::isIconTile,
-                columns = columns,
-                showLabels = showLabels,
-                dragAndDropState = dragAndDropState,
-            )
-            AvailableTiles(
-                tiles = otherTiles.filter { !dragAndDropState.isMoving(it.tileSpec) },
-                largeTileHeight = largeTileHeight,
-                iconTileHeight = iconTileHeight,
-                tilePadding = tilePadding,
-                addTileToEnd = addTileToEnd,
-                onRemove = onRemoveTile,
-                onDoubleTap = onDoubleTap,
-                isIconOnly = viewModel::isIconTile,
-                showLabels = showLabels,
-                columns = columns,
-                dragAndDropState = dragAndDropState,
-            )
-        }
-    }
-
-    override fun splitIntoPages(
-        tiles: List<TileViewModel>,
-        rows: Int,
-        columns: Int,
-    ): List<List<TileViewModel>> {
-        val (smallTiles, largeTiles) = tiles.partition { viewModel.isIconTile(it.spec) }
-
-        val sizedLargeTiles = largeTiles.map { SizedTile(it, 2) }
-        val sizedSmallTiles = smallTiles.map { SizedTile(it, 1) }
-        val largeTilesRows = PaginatableGridLayout.splitInRows(sizedLargeTiles, columns)
-        val smallTilesRows = PaginatableGridLayout.splitInRows(sizedSmallTiles, columns)
-        return (largeTilesRows + smallTilesRows).chunked(rows).map { it.flatten().map { it.tile } }
-    }
-
-    @Composable
-    private fun CurrentTiles(
-        tiles: List<EditTileViewModel>,
-        largeTileHeight: Dp,
-        iconTileHeight: Dp,
-        tilePadding: Dp,
-        onAdd: (TileSpec, Int) -> Unit,
-        onRemove: (TileSpec) -> Unit,
-        onDoubleTap: (TileSpec) -> Unit,
-        isIconOnly: (TileSpec) -> Boolean,
-        showLabels: Boolean,
-        columns: Int,
-        dragAndDropState: DragAndDropState,
-    ) {
-        val (smallTiles, largeTiles) = tiles.partition { isIconOnly(it.tileSpec) }
-
-        val largeGridHeight = gridHeight(largeTiles.size, largeTileHeight, columns / 2, tilePadding)
-        val smallGridHeight = gridHeight(smallTiles.size, iconTileHeight, columns, tilePadding)
-
-        CurrentTilesContainer {
-            TileLazyGrid(
-                columns = GridCells.Fixed(columns),
-                modifier =
-                    Modifier.height(largeGridHeight)
-                        .dragAndDropTileList(dragAndDropState, { !isIconOnly(it) }, onAdd)
-            ) {
-                editTiles(
-                    tiles = largeTiles,
-                    clickAction = ClickAction.REMOVE,
-                    onClick = onRemove,
-                    onDoubleTap = onDoubleTap,
-                    isIconOnly = { false },
-                    dragAndDropState = dragAndDropState,
-                    acceptDrops = { !isIconOnly(it) },
-                    onDrop = onAdd,
-                    indicatePosition = true,
-                )
-            }
-        }
-
-        CurrentTilesContainer {
-            TileLazyGrid(
-                columns = GridCells.Fixed(columns),
-                modifier =
-                    Modifier.height(smallGridHeight)
-                        .dragAndDropTileList(dragAndDropState, { isIconOnly(it) }, onAdd)
-            ) {
-                editTiles(
-                    tiles = smallTiles,
-                    clickAction = ClickAction.REMOVE,
-                    onClick = onRemove,
-                    onDoubleTap = onDoubleTap,
-                    isIconOnly = { true },
-                    showLabels = showLabels,
-                    dragAndDropState = dragAndDropState,
-                    acceptDrops = { isIconOnly(it) },
-                    onDrop = onAdd,
-                    indicatePosition = true,
-                )
-            }
-        }
-    }
-
-    @Composable
-    private fun AvailableTiles(
-        tiles: List<EditTileViewModel>,
-        largeTileHeight: Dp,
-        iconTileHeight: Dp,
-        tilePadding: Dp,
-        addTileToEnd: (TileSpec) -> Unit,
-        onRemove: (TileSpec) -> Unit,
-        onDoubleTap: (TileSpec) -> Unit,
-        isIconOnly: (TileSpec) -> Boolean,
-        showLabels: Boolean,
-        columns: Int,
-        dragAndDropState: DragAndDropState,
-    ) {
-        val (tilesStock, tilesCustom) = tiles.partition { it.appName == null }
-        val (smallTiles, largeTiles) = tilesStock.partition { isIconOnly(it.tileSpec) }
-
-        val largeGridHeight = gridHeight(largeTiles.size, largeTileHeight, columns / 2, tilePadding)
-        val smallGridHeight = gridHeight(smallTiles.size, iconTileHeight, columns, tilePadding)
-        val largeGridHeightCustom =
-            gridHeight(tilesCustom.size, iconTileHeight, columns, tilePadding)
-
-        // Add up the height of all three grids and add padding in between
-        val gridHeight =
-            largeGridHeight + smallGridHeight + largeGridHeightCustom + (tilePadding * 2)
-
-        val onDrop: (TileSpec, Int) -> Unit by rememberUpdatedState { tileSpec, _ ->
-            onRemove(tileSpec)
-        }
-
-        AvailableTilesContainer {
-            TileLazyGrid(
-                columns = GridCells.Fixed(columns),
-                modifier =
-                    Modifier.height(gridHeight)
-                        .dragAndDropTileList(dragAndDropState, { true }, onDrop)
-            ) {
-                // Large tiles
-                editTiles(
-                    largeTiles,
-                    ClickAction.ADD,
-                    addTileToEnd,
-                    isIconOnly,
-                    dragAndDropState,
-                    onDoubleTap = onDoubleTap,
-                    acceptDrops = { true },
-                    onDrop = onDrop,
-                )
-                fillUpRow(nTiles = largeTiles.size, columns = columns / 2)
-
-                // Small tiles
-                editTiles(
-                    smallTiles,
-                    ClickAction.ADD,
-                    addTileToEnd,
-                    isIconOnly,
-                    dragAndDropState,
-                    onDoubleTap = onDoubleTap,
-                    showLabels = showLabels,
-                    acceptDrops = { true },
-                    onDrop = onDrop,
-                )
-                fillUpRow(nTiles = smallTiles.size, columns = columns)
-
-                // Custom tiles, all icons
-                editTiles(
-                    tilesCustom,
-                    ClickAction.ADD,
-                    addTileToEnd,
-                    isIconOnly,
-                    dragAndDropState,
-                    onDoubleTap = onDoubleTap,
-                    showLabels = showLabels,
-                    acceptDrops = { true },
-                    onDrop = onDrop,
-                )
-            }
-        }
-    }
-
-    @Composable
-    private fun CurrentTilesContainer(content: @Composable () -> Unit) {
-        Box(
-            Modifier.fillMaxWidth()
-                .dashedBorder(
-                    color = MaterialTheme.colorScheme.onBackground.copy(alpha = .5f),
-                    shape = Dimensions.ContainerShape,
-                )
-                .padding(dimensionResource(R.dimen.qs_tile_margin_vertical))
-        ) {
-            content()
-        }
-    }
-
-    @Composable
-    private fun AvailableTilesContainer(content: @Composable () -> Unit) {
-        Box(
-            Modifier.fillMaxWidth()
-                .background(
-                    color = MaterialTheme.colorScheme.background,
-                    alpha = { 1f },
-                    shape = Dimensions.ContainerShape,
-                )
-                .padding(dimensionResource(R.dimen.qs_tile_margin_vertical))
-        ) {
-            content()
-        }
-    }
-
-    /** Fill up the rest of the row if it's not complete. */
-    private fun LazyGridScope.fillUpRow(nTiles: Int, columns: Int) {
-        if (nTiles % columns != 0) {
-            item(span = { GridItemSpan(maxCurrentLineSpan) }) { Spacer(Modifier) }
-        }
-    }
-
-    private fun Modifier.dashedBorder(
-        color: Color,
-        shape: Shape,
-    ): Modifier {
-        return this.drawWithContent {
-            val outline = shape.createOutline(size, layoutDirection, this)
-            val path = Path()
-            path.addOutline(outline)
-            val stroke =
-                Stroke(
-                    width = 1.dp.toPx(),
-                    pathEffect = PathEffect.dashPathEffect(floatArrayOf(10f, 10f))
-                )
-            this.drawContent()
-            drawPath(path = path, style = stroke, color = color)
-        }
-    }
-
-    private object Dimensions {
-        // Corner radius is half the height of a tile + padding
-        val ContainerShape = RoundedCornerShape(48.dp)
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/StretchedGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/StretchedGridLayout.kt
deleted file mode 100644
index 3e48245..0000000
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/StretchedGridLayout.kt
+++ /dev/null
@@ -1,139 +0,0 @@
-/*
- * Copyright (C) 2024 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.qs.panels.ui.compose
-
-import androidx.compose.foundation.layout.height
-import androidx.compose.foundation.lazy.grid.GridCells
-import androidx.compose.foundation.lazy.grid.GridItemSpan
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.DisposableEffect
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.remember
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.res.dimensionResource
-import androidx.lifecycle.compose.collectAsStateWithLifecycle
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.qs.panels.shared.model.SizedTile
-import com.android.systemui.qs.panels.shared.model.TileRow
-import com.android.systemui.qs.panels.ui.viewmodel.EditTileViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.FixedColumnsSizeViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.IconTilesViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
-import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.android.systemui.res.R
-import javax.inject.Inject
-
-@SysUISingleton
-class StretchedGridLayout
-@Inject
-constructor(
-    private val iconTilesViewModel: IconTilesViewModel,
-    private val gridSizeViewModel: FixedColumnsSizeViewModel,
-) : GridLayout {
-
-    @Composable
-    override fun TileGrid(
-        tiles: List<TileViewModel>,
-        modifier: Modifier,
-        editModeStart: () -> Unit,
-    ) {
-        DisposableEffect(tiles) {
-            val token = Any()
-            tiles.forEach { it.startListening(token) }
-            onDispose { tiles.forEach { it.stopListening(token) } }
-        }
-
-        // Tile widths [normal|stretched]
-        // Icon [3 | 4]
-        // Large [6 | 8]
-        val columns = 12
-        val stretchedTiles =
-            remember(tiles) {
-                val sizedTiles =
-                    tiles.map {
-                        SizedTile(
-                            it,
-                            if (iconTilesViewModel.isIconTile(it.spec)) {
-                                3
-                            } else {
-                                6
-                            }
-                        )
-                    }
-                splitInRows(sizedTiles, columns)
-            }
-
-        TileLazyGrid(columns = GridCells.Fixed(columns), modifier = modifier) {
-            items(stretchedTiles.size, span = { GridItemSpan(stretchedTiles[it].width) }) { index ->
-                Tile(
-                    tile = stretchedTiles[index].tile,
-                    iconOnly = iconTilesViewModel.isIconTile(stretchedTiles[index].tile.spec),
-                    modifier = Modifier.height(dimensionResource(id = R.dimen.qs_tile_height))
-                )
-            }
-        }
-    }
-
-    @Composable
-    override fun EditTileGrid(
-        tiles: List<EditTileViewModel>,
-        modifier: Modifier,
-        onAddTile: (TileSpec, Int) -> Unit,
-        onRemoveTile: (TileSpec) -> Unit
-    ) {
-        val columns by gridSizeViewModel.columns.collectAsStateWithLifecycle()
-
-        DefaultEditTileGrid(
-            tiles = tiles,
-            isIconOnly = iconTilesViewModel::isIconTile,
-            columns = columns,
-            modifier = modifier,
-            onAddTile = onAddTile,
-            onRemoveTile = onRemoveTile,
-            onResize = iconTilesViewModel::resize,
-        )
-    }
-
-    private fun splitInRows(
-        tiles: List<SizedTile<TileViewModel>>,
-        columns: Int
-    ): List<SizedTile<TileViewModel>> {
-        val row = TileRow<TileViewModel>(columns)
-
-        return buildList {
-            for (tile in tiles) {
-                if (row.maybeAddTile(tile)) {
-                    if (row.isFull()) {
-                        // Row is full, no need to stretch tiles
-                        addAll(row.tiles)
-                        row.clear()
-                    }
-                } else {
-                    if (row.isFull()) {
-                        addAll(row.tiles)
-                    } else {
-                        // Stretching tiles when row isn't full
-                        addAll(row.tiles.map { it.copy(width = it.width + (it.width / 3)) })
-                    }
-                    row.clear()
-                    row.maybeAddTile(tile)
-                }
-            }
-            addAll(row.tiles)
-        }
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewImpl.kt b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewImpl.kt
index dbfe818..abc0453 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/QSTileViewImpl.kt
@@ -586,6 +586,15 @@
                     )
                 )
             )
+        } else {
+            if (isLongClickable) {
+                info.addAction(
+                    AccessibilityNodeInfo.AccessibilityAction(
+                        AccessibilityNodeInfo.AccessibilityAction.ACTION_LONG_CLICK.id,
+                        resources.getString(R.string.accessibility_long_click_tile)
+                    )
+                )
+            }
         }
         if (!TextUtils.isEmpty(accessibilityClass)) {
             info.className =
@@ -597,14 +606,6 @@
             if (Switch::class.java.name == accessibilityClass) {
                 info.isChecked = tileState
                 info.isCheckable = true
-                if (isLongClickable) {
-                    info.addAction(
-                        AccessibilityNodeInfo.AccessibilityAction(
-                            AccessibilityNodeInfo.AccessibilityAction.ACTION_LONG_CLICK.id,
-                            resources.getString(R.string.accessibility_long_click_tile)
-                        )
-                    )
-                }
             }
         }
         if (position != INVALID) {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/SubtitleArrayMapping.kt b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/SubtitleArrayMapping.kt
index 53594bb..f702da4 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tileimpl/SubtitleArrayMapping.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tileimpl/SubtitleArrayMapping.kt
@@ -27,7 +27,6 @@
         subtitleIdsMap["cell"] = R.array.tile_states_cell
         subtitleIdsMap["battery"] = R.array.tile_states_battery
         subtitleIdsMap["dnd"] = R.array.tile_states_dnd
-        subtitleIdsMap["modes"] = R.array.tile_states_modes
         subtitleIdsMap["flashlight"] = R.array.tile_states_flashlight
         subtitleIdsMap["rotation"] = R.array.tile_states_rotation
         subtitleIdsMap["bt"] = R.array.tile_states_bt
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/DndTile.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/DndTile.java
index bdf935e..b927134 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/DndTile.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/DndTile.java
@@ -49,6 +49,7 @@
 import com.android.systemui.animation.Expandable;
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.flags.RefactorFlagUtils;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.plugins.FalsingManager;
 import com.android.systemui.plugins.qs.QSTile.BooleanState;
@@ -105,6 +106,11 @@
     ) {
         super(host, uiEventLogger, backgroundLooper, mainHandler, falsingManager, metricsLogger,
                 statusBarStateController, activityStarter, qsLogger);
+
+        // If the flag is on, this shouldn't run at all since the modes tile replaces the DND tile.
+        RefactorFlagUtils.INSTANCE.assertInLegacyMode(android.app.Flags.modesUi(),
+                android.app.Flags.FLAG_MODES_UI);
+
         mController = zenModeController;
         mSharedPreferences = sharedPreferences;
         mController.observe(getLifecycle(), mZenCallback);
@@ -253,18 +259,20 @@
             case Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS:
                 state.contentDescription =
                         mContext.getString(R.string.accessibility_quick_settings_dnd) + ", "
-                        + state.secondaryLabel;
+                                + state.secondaryLabel;
                 break;
             case Global.ZEN_MODE_NO_INTERRUPTIONS:
                 state.contentDescription =
                         mContext.getString(R.string.accessibility_quick_settings_dnd) + ", " +
-                        mContext.getString(R.string.accessibility_quick_settings_dnd_none_on)
+                                mContext.getString(
+                                        R.string.accessibility_quick_settings_dnd_none_on)
                                 + ", " + state.secondaryLabel;
                 break;
             case ZEN_MODE_ALARMS:
                 state.contentDescription =
                         mContext.getString(R.string.accessibility_quick_settings_dnd) + ", " +
-                        mContext.getString(R.string.accessibility_quick_settings_dnd_alarms_on)
+                                mContext.getString(
+                                        R.string.accessibility_quick_settings_dnd_alarms_on)
                                 + ", " + state.secondaryLabel;
                 break;
             default:
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
index a300031..2a33a16 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
@@ -23,13 +23,15 @@
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.coroutineScope
 import androidx.lifecycle.repeatOnLifecycle
+import com.android.internal.R.attr.contentDescription
 import com.android.internal.logging.MetricsLogger
 import com.android.systemui.animation.Expandable
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.flags.RefactorFlagUtils.isUnexpectedlyInLegacyMode
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.plugins.FalsingManager
-import com.android.systemui.plugins.qs.QSTile.BooleanState
+import com.android.systemui.plugins.qs.QSTile
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.qs.QSHost
 import com.android.systemui.qs.QsEventLogger
@@ -63,7 +65,7 @@
     private val tileMapper: ModesTileMapper,
     private val userActionInteractor: ModesTileUserActionInteractor,
 ) :
-    QSTileImpl<BooleanState>(
+    QSTileImpl<QSTile.State>(
         host,
         uiEventLogger,
         backgroundLooper,
@@ -79,6 +81,8 @@
     private val config = qsTileConfigProvider.getConfig(TILE_SPEC)
 
     init {
+        /* Check if */ isUnexpectedlyInLegacyMode(Flags.modesUi(), Flags.FLAG_MODES_UI)
+
         lifecycle.coroutineScope.launch {
             lifecycle.repeatOnLifecycle(Lifecycle.State.RESUMED) {
                 dataInteractor.tileData().collect { refreshState(it) }
@@ -90,7 +94,7 @@
 
     override fun getTileLabel(): CharSequence = tileState.label
 
-    override fun newTileState() = BooleanState()
+    override fun newTileState() = QSTile.State()
 
     override fun handleClick(expandable: Expandable?) = runBlocking {
         userActionInteractor.handleClick(expandable)
@@ -98,22 +102,22 @@
 
     override fun getLongClickIntent(): Intent = userActionInteractor.longClickIntent
 
-    override fun handleUpdateState(booleanState: BooleanState?, arg: Any?) {
+    override fun handleUpdateState(state: QSTile.State?, arg: Any?) {
         if (arg is ModesTileModel) {
             tileState = tileMapper.map(config, arg)
 
-            booleanState?.apply {
-                state = tileState.activationState.legacyState
+            state?.apply {
+                this.state = tileState.activationState.legacyState
                 icon = ResourceIcon.get(tileState.iconRes ?: R.drawable.qs_dnd_icon_off)
                 label = tileLabel
                 secondaryLabel = tileState.secondaryLabel
                 contentDescription = tileState.contentDescription
-                forceExpandIcon = true
+                expandedAccessibilityClassName = tileState.expandedAccessibilityClassName
             }
         }
     }
 
     companion object {
-        const val TILE_SPEC = "modes"
+        const val TILE_SPEC = "dnd"
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogController.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogController.java
index 158eb6e..b2873c5 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogController.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogController.java
@@ -736,7 +736,8 @@
         // Set network description for the carrier network when connecting to the carrier network
         // under the airplane mode ON.
         if (activeNetworkIsCellular() || isCarrierNetworkActive()) {
-            summary = context.getString(R.string.preference_summary_default_combination,
+            summary = context.getString(
+                    com.android.settingslib.R.string.preference_summary_default_combination,
                     context.getString(
                             isForDds // if nonDds is active, explains Dds status as poor connection
                                     ? (isOnNonDds ? R.string.mobile_data_poor_connection
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
index 31e91aa..92efa40 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
@@ -19,20 +19,27 @@
 import android.app.Flags
 import android.os.UserHandle
 import com.android.settingslib.notification.data.repository.ZenModeRepository
+import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.qs.tiles.base.interactor.DataUpdateTrigger
 import com.android.systemui.qs.tiles.base.interactor.QSTileDataInteractor
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 
-class ModesTileDataInteractor @Inject constructor(val zenModeRepository: ZenModeRepository) :
-    QSTileDataInteractor<ModesTileModel> {
-    private val zenModeActive =
+class ModesTileDataInteractor
+@Inject
+constructor(
+    val zenModeRepository: ZenModeRepository,
+    @Background val bgDispatcher: CoroutineDispatcher,
+) : QSTileDataInteractor<ModesTileModel> {
+    private val activeModes =
         zenModeRepository.modes
-            .map { modes -> modes.any { mode -> mode.isActive } }
+            .map { modes -> modes.filter { mode -> mode.isActive }.map { it.name } }
             .distinctUntilChanged()
 
     override fun tileData(
@@ -45,7 +52,10 @@
      *
      * TODO(b/299909989): Remove after the transition.
      */
-    fun tileData() = zenModeActive.map { ModesTileModel(isActivated = it) }
+    fun tileData() =
+        activeModes
+            .map { ModesTileModel(isActivated = it.isNotEmpty(), activeModes = it) }
+            .flowOn(bgDispatcher)
 
     override fun availability(user: UserHandle): Flow<Boolean> = flowOf(Flags.modesUi())
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractor.kt
index 4c6563d..083bf05 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractor.kt
@@ -16,14 +16,10 @@
 
 package com.android.systemui.qs.tiles.impl.modes.domain.interactor
 
-//noinspection CleanArchitectureDependencyViolation: dialog needs to be opened on click
 import android.content.Intent
 import android.provider.Settings
-import com.android.internal.jank.InteractionJankMonitor
-import com.android.systemui.animation.DialogCuj
-import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.animation.Expandable
-import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.qs.tiles.base.actions.QSTileIntentUserInputHandler
 import com.android.systemui.qs.tiles.base.interactor.QSTileInput
 import com.android.systemui.qs.tiles.base.interactor.QSTileUserActionInteractor
@@ -31,15 +27,13 @@
 import com.android.systemui.qs.tiles.viewmodel.QSTileUserAction
 import com.android.systemui.statusbar.policy.ui.dialog.ModesDialogDelegate
 import javax.inject.Inject
-import kotlin.coroutines.CoroutineContext
-import kotlinx.coroutines.withContext
 
+@SysUISingleton
 class ModesTileUserActionInteractor
 @Inject
 constructor(
-    @Main private val coroutineContext: CoroutineContext,
-    private val qsTileIntentUserActionHandler: QSTileIntentUserInputHandler,
-    private val dialogTransitionAnimator: DialogTransitionAnimator,
+    private val qsTileIntentUserInputHandler: QSTileIntentUserInputHandler,
+    // TODO(b/353896370): The domain layer should not have to depend on the UI layer.
     private val dialogDelegate: ModesDialogDelegate,
 ) : QSTileUserActionInteractor<ModesTileModel> {
     val longClickIntent = Intent(Settings.ACTION_ZEN_MODE_SETTINGS)
@@ -51,29 +45,14 @@
                     handleClick(action.expandable)
                 }
                 is QSTileUserAction.LongClick -> {
-                    qsTileIntentUserActionHandler.handle(action.expandable, longClickIntent)
+                    qsTileIntentUserInputHandler.handle(action.expandable, longClickIntent)
                 }
             }
         }
     }
 
     suspend fun handleClick(expandable: Expandable?) {
-        // Show a dialog with the list of modes to configure. Dialogs shown by the
-        // DialogTransitionAnimator must be created and shown on the main thread, so we post it to
-        // the UI handler.
-        withContext(coroutineContext) {
-            val dialog = dialogDelegate.createDialog()
-
-            expandable
-                ?.dialogTransitionController(
-                    DialogCuj(InteractionJankMonitor.CUJ_SHADE_DIALOG_OPEN, INTERACTION_JANK_TAG)
-                )
-                ?.let { controller -> dialogTransitionAnimator.show(dialog, controller) }
-                ?: dialog.show()
-        }
-    }
-
-    companion object {
-        private const val INTERACTION_JANK_TAG = "configure_priority_modes"
+        // Show a dialog with the list of modes to configure.
+        dialogDelegate.showDialog(expandable)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
index e44413a..cc509ea 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
@@ -15,4 +15,4 @@
  */
 
 package com.android.systemui.qs.tiles.impl.modes.domain.model
-data class ModesTileModel(val isActivated: Boolean)
+data class ModesTileModel(val isActivated: Boolean, val activeModes: List<String>)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
index 7048ada..7afdb75 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
@@ -17,6 +17,8 @@
 package com.android.systemui.qs.tiles.impl.modes.ui
 
 import android.content.res.Resources
+import android.icu.text.MessageFormat
+import android.widget.Button
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.qs.tiles.base.interactor.QSTileDataToStateMapper
@@ -24,6 +26,7 @@
 import com.android.systemui.qs.tiles.viewmodel.QSTileConfig
 import com.android.systemui.qs.tiles.viewmodel.QSTileState
 import com.android.systemui.res.R
+import java.util.Locale
 import javax.inject.Inject
 
 class ModesTileMapper
@@ -46,19 +49,32 @@
                     contentDescription = null,
                 )
             this.icon = { icon }
-            if (data.isActivated) {
-                activationState = QSTileState.ActivationState.ACTIVE
-                secondaryLabel = "Some modes enabled idk" // TODO(b/346519570)
-            } else {
-                activationState = QSTileState.ActivationState.INACTIVE
-                secondaryLabel = "Off" // TODO(b/346519570)
-            }
-            contentDescription = label
+            activationState =
+                if (data.isActivated) {
+                    QSTileState.ActivationState.ACTIVE
+                } else {
+                    QSTileState.ActivationState.INACTIVE
+                }
+            secondaryLabel = getModesStatus(data, resources)
+            contentDescription = "$label. $secondaryLabel"
             supportedActions =
                 setOf(
                     QSTileState.UserAction.CLICK,
                     QSTileState.UserAction.LONG_CLICK,
                 )
             sideViewIcon = QSTileState.SideViewIcon.Chevron
+            expandedAccessibilityClass = Button::class
         }
+
+    private fun getModesStatus(data: ModesTileModel, resources: Resources): String {
+        val msgFormat =
+            MessageFormat(resources.getString(R.string.zen_mode_active_modes), Locale.getDefault())
+        val count = data.activeModes.count()
+        val args: MutableMap<String, Any> = HashMap()
+        args["count"] = count
+        if (count >= 1) {
+            args["mode"] = data.activeModes[0]
+        }
+        return msgFormat.format(args)
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt
index ba0a8d6..2cdcc24 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/viewmodel/QSTileViewModelAdapter.kt
@@ -19,7 +19,6 @@
 import android.content.Context
 import android.os.UserHandle
 import android.util.Log
-import androidx.annotation.GuardedBy
 import com.android.internal.logging.InstanceId
 import com.android.systemui.Dumpable
 import com.android.systemui.animation.Expandable
@@ -34,6 +33,7 @@
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import java.io.PrintWriter
+import java.util.concurrent.CopyOnWriteArraySet
 import java.util.function.Supplier
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
@@ -57,10 +57,8 @@
     private val context
         get() = qsHost.context
 
-    @GuardedBy("callbacks")
-    private val callbacks: MutableCollection<QSTile.Callback> = mutableSetOf()
-    @GuardedBy("listeningClients")
-    private val listeningClients: MutableCollection<Any> = mutableSetOf()
+    private val callbacks = CopyOnWriteArraySet<QSTile.Callback>()
+    private val listeningClients = CopyOnWriteArraySet<Any>()
 
     // Cancels the jobs when the adapter is no longer alive
     private var tileAdapterJob: Job? = null
@@ -113,19 +111,17 @@
 
     override fun addCallback(callback: QSTile.Callback?) {
         callback ?: return
-        synchronized(callbacks) {
-            callbacks.add(callback)
-            state?.let(callback::onStateChanged)
-        }
+        callbacks.add(callback)
+        state?.let(callback::onStateChanged)
     }
 
     override fun removeCallback(callback: QSTile.Callback?) {
         callback ?: return
-        synchronized(callbacks) { callbacks.remove(callback) }
+        callbacks.remove(callback)
     }
 
     override fun removeCallbacks() {
-        synchronized(callbacks) { callbacks.clear() }
+        callbacks.clear()
     }
 
     override fun click(expandable: Expandable?) {
@@ -163,32 +159,28 @@
 
     override fun setListening(client: Any?, listening: Boolean) {
         client ?: return
-        synchronized(listeningClients) {
-            if (listening) {
-                listeningClients.add(client)
-                if (listeningClients.size == 1) {
-                    stateJob =
-                        qsTileViewModel.state
-                            .filterNotNull()
-                            .map { mapState(context, it, qsTileViewModel.config) }
-                            .onEach { legacyState ->
-                                synchronized(callbacks) {
-                                    callbacks.forEach { it.onStateChanged(legacyState) }
-                                }
-                            }
-                            .launchIn(applicationScope)
-                }
-            } else {
-                listeningClients.remove(client)
-                if (listeningClients.isEmpty()) {
-                    stateJob?.cancel()
-                }
+        if (listening) {
+            listeningClients.add(client)
+            if (listeningClients.size == 1) {
+                stateJob =
+                    qsTileViewModel.state
+                        .filterNotNull()
+                        .map { mapState(context, it, qsTileViewModel.config) }
+                        .onEach { legacyState ->
+                            callbacks.forEach { it.onStateChanged(legacyState) }
+                        }
+                        .launchIn(applicationScope)
+            }
+        } else {
+            listeningClients.remove(client)
+            if (listeningClients.isEmpty()) {
+                stateJob?.cancel()
             }
         }
     }
 
     override fun isListening(): Boolean =
-        synchronized(listeningClients) { listeningClients.isNotEmpty() }
+        listeningClients.isNotEmpty()
 
     override fun setDetailListening(show: Boolean) {
         // do nothing like QSTileImpl
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
index 3fca84e..5b50133 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
@@ -172,12 +172,28 @@
 
     private fun resetShadeSessions() {
         applicationScope.launch {
-            sceneBackInteractor.backStack
-                // We are in a session if either Shade or QuickSettings is on the back stack
-                .map { backStack ->
-                    backStack.asIterable().any { it == Scenes.Shade || it == Scenes.QuickSettings }
+            combine(
+                    sceneBackInteractor.backStack
+                        // We are in a session if either Shade or QuickSettings is on the back stack
+                        .map { backStack ->
+                            backStack.asIterable().any {
+                                it == Scenes.Shade || it == Scenes.QuickSettings
+                            }
+                        }
+                        .distinctUntilChanged(),
+                    sceneInteractor.transitionState
+                        .mapNotNull { state ->
+                            // We are also in a session if either Shade or QuickSettings is the
+                            // current scene
+                            when (state) {
+                                is ObservableTransitionState.Idle -> state.currentScene
+                                is ObservableTransitionState.Transition -> state.fromScene
+                            }.let { it == Scenes.Shade || it == Scenes.QuickSettings }
+                        }
+                        .distinctUntilChanged()
+                ) { inBackStack, isCurrentScene ->
+                    inBackStack || isCurrentScene
                 }
-                .distinctUntilChanged()
                 // Once a session has ended, clear the session storage.
                 .filter { inSession -> !inSession }
                 .collect { shadeSessionStorage.clear() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt b/packages/SystemUI/src/com/android/systemui/screenshot/InteractiveScreenshotHandler.kt
similarity index 62%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
copy to packages/SystemUI/src/com/android/systemui/screenshot/InteractiveScreenshotHandler.kt
index 37c9552..26405f0 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/InteractiveScreenshotHandler.kt
@@ -14,11 +14,20 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.domain.interactor
+package com.android.systemui.screenshot
 
-import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
+import android.view.Display
 
-val Kosmos.partitionedGridLayout by
-    Kosmos.Fixture { PartitionedGridLayout(partitionedGridViewModel) }
+interface InteractiveScreenshotHandler : ScreenshotHandler {
+    fun isPendingSharedTransition(): Boolean
+
+    fun requestDismissal(event: ScreenshotEvent)
+
+    fun removeWindow()
+
+    fun onDestroy()
+
+    interface Factory {
+        fun create(display: Display): InteractiveScreenshotHandler
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
new file mode 100644
index 0000000..a2583e6
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
@@ -0,0 +1,848 @@
+/*
+ * Copyright (C) 2024 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.screenshot;
+
+import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
+import static android.view.WindowManager.LayoutParams.TYPE_SCREENSHOT;
+
+import static com.android.systemui.Flags.screenshotPrivateProfileAccessibilityAnnouncementFix;
+import static com.android.systemui.Flags.screenshotSaveImageExporter;
+import static com.android.systemui.screenshot.LogConfig.DEBUG_ANIM;
+import static com.android.systemui.screenshot.LogConfig.DEBUG_CALLBACK;
+import static com.android.systemui.screenshot.LogConfig.DEBUG_INPUT;
+import static com.android.systemui.screenshot.LogConfig.DEBUG_UI;
+import static com.android.systemui.screenshot.LogConfig.DEBUG_WINDOW;
+import static com.android.systemui.screenshot.LogConfig.logTag;
+import static com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER;
+import static com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_INTERACTION_TIMEOUT;
+
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
+import android.annotation.MainThread;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.app.ICompatCameraControlCallback;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.content.pm.ActivityInfo;
+import android.content.res.Configuration;
+import android.graphics.Bitmap;
+import android.graphics.Insets;
+import android.graphics.Rect;
+import android.net.Uri;
+import android.os.Process;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.provider.Settings;
+import android.util.DisplayMetrics;
+import android.util.Log;
+import android.view.Display;
+import android.view.ScrollCaptureResponse;
+import android.view.View;
+import android.view.ViewGroup;
+import android.view.ViewRootImpl;
+import android.view.ViewTreeObserver;
+import android.view.WindowInsets;
+import android.view.WindowManager;
+import android.widget.Toast;
+import android.window.WindowContext;
+
+import com.android.internal.logging.UiEventLogger;
+import com.android.internal.policy.PhoneWindow;
+import com.android.settingslib.applications.InterestingConfigChanges;
+import com.android.systemui.broadcast.BroadcastDispatcher;
+import com.android.systemui.broadcast.BroadcastSender;
+import com.android.systemui.clipboardoverlay.ClipboardOverlayController;
+import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.res.R;
+import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback;
+import com.android.systemui.screenshot.scroll.ScrollCaptureExecutor;
+import com.android.systemui.util.Assert;
+
+import com.google.common.util.concurrent.ListenableFuture;
+
+import dagger.assisted.Assisted;
+import dagger.assisted.AssistedFactory;
+import dagger.assisted.AssistedInject;
+
+import kotlin.Unit;
+
+import java.util.UUID;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.function.Consumer;
+
+import javax.inject.Provider;
+
+/**
+ * Controls the state and flow for screenshots.
+ */
+public class LegacyScreenshotController implements InteractiveScreenshotHandler {
+    private static final String TAG = logTag(LegacyScreenshotController.class);
+
+    // From WizardManagerHelper.java
+    private static final String SETTINGS_SECURE_USER_SETUP_COMPLETE = "user_setup_complete";
+
+    static final int SCREENSHOT_CORNER_DEFAULT_TIMEOUT_MILLIS = 6000;
+
+    private final WindowContext mContext;
+    private final FeatureFlags mFlags;
+    private final ScreenshotShelfViewProxy mViewProxy;
+    private final ScreenshotNotificationsController mNotificationsController;
+    private final ScreenshotSmartActions mScreenshotSmartActions;
+    private final UiEventLogger mUiEventLogger;
+    private final ImageExporter mImageExporter;
+    private final ImageCapture mImageCapture;
+    private final Executor mMainExecutor;
+    private final ExecutorService mBgExecutor;
+    private final BroadcastSender mBroadcastSender;
+    private final BroadcastDispatcher mBroadcastDispatcher;
+    private final ScreenshotActionsController mActionsController;
+
+    private final WindowManager mWindowManager;
+    private final WindowManager.LayoutParams mWindowLayoutParams;
+    @Nullable
+    private final ScreenshotSoundController mScreenshotSoundController;
+    private final PhoneWindow mWindow;
+    private final Display mDisplay;
+    private final ScrollCaptureExecutor mScrollCaptureExecutor;
+    private final ScreenshotNotificationSmartActionsProvider
+            mScreenshotNotificationSmartActionsProvider;
+    private final TimeoutHandler mScreenshotHandler;
+    private final UserManager mUserManager;
+    private final AssistContentRequester mAssistContentRequester;
+    private final ActionExecutor mActionExecutor;
+
+
+    private final MessageContainerController mMessageContainerController;
+    private final AnnouncementResolver mAnnouncementResolver;
+    private Bitmap mScreenBitmap;
+    private SaveImageInBackgroundTask mSaveInBgTask;
+    private boolean mScreenshotTakenInPortrait;
+    private boolean mAttachRequested;
+    private boolean mDetachRequested;
+    private Animator mScreenshotAnimation;
+    private RequestCallback mCurrentRequestCallback;
+    private String mPackageName = "";
+    private final BroadcastReceiver mCopyBroadcastReceiver;
+
+    /** Tracks config changes that require re-creating UI */
+    private final InterestingConfigChanges mConfigChanges = new InterestingConfigChanges(
+            ActivityInfo.CONFIG_ORIENTATION
+                    | ActivityInfo.CONFIG_LAYOUT_DIRECTION
+                    | ActivityInfo.CONFIG_LOCALE
+                    | ActivityInfo.CONFIG_UI_MODE
+                    | ActivityInfo.CONFIG_SCREEN_LAYOUT
+                    | ActivityInfo.CONFIG_ASSETS_PATHS);
+
+
+    @AssistedInject
+    LegacyScreenshotController(
+            Context context,
+            WindowManager windowManager,
+            FeatureFlags flags,
+            ScreenshotShelfViewProxy.Factory viewProxyFactory,
+            ScreenshotSmartActions screenshotSmartActions,
+            ScreenshotNotificationsController.Factory screenshotNotificationsControllerFactory,
+            UiEventLogger uiEventLogger,
+            ImageExporter imageExporter,
+            ImageCapture imageCapture,
+            @Main Executor mainExecutor,
+            ScrollCaptureExecutor scrollCaptureExecutor,
+            TimeoutHandler timeoutHandler,
+            BroadcastSender broadcastSender,
+            BroadcastDispatcher broadcastDispatcher,
+            ScreenshotNotificationSmartActionsProvider screenshotNotificationSmartActionsProvider,
+            ScreenshotActionsController.Factory screenshotActionsControllerFactory,
+            ActionExecutor.Factory actionExecutorFactory,
+            UserManager userManager,
+            AssistContentRequester assistContentRequester,
+            MessageContainerController messageContainerController,
+            Provider<ScreenshotSoundController> screenshotSoundController,
+            AnnouncementResolver announcementResolver,
+            @Assisted Display display
+    ) {
+        mScreenshotSmartActions = screenshotSmartActions;
+        mNotificationsController = screenshotNotificationsControllerFactory.create(
+                display.getDisplayId());
+        mUiEventLogger = uiEventLogger;
+        mImageExporter = imageExporter;
+        mImageCapture = imageCapture;
+        mMainExecutor = mainExecutor;
+        mScrollCaptureExecutor = scrollCaptureExecutor;
+        mScreenshotNotificationSmartActionsProvider = screenshotNotificationSmartActionsProvider;
+        mBgExecutor = Executors.newSingleThreadExecutor();
+        mBroadcastSender = broadcastSender;
+        mBroadcastDispatcher = broadcastDispatcher;
+
+        mScreenshotHandler = timeoutHandler;
+        mScreenshotHandler.setDefaultTimeoutMillis(SCREENSHOT_CORNER_DEFAULT_TIMEOUT_MILLIS);
+
+        mDisplay = display;
+        mWindowManager = windowManager;
+        final Context displayContext = context.createDisplayContext(display);
+        mContext = (WindowContext) displayContext.createWindowContext(TYPE_SCREENSHOT, null);
+        mFlags = flags;
+        mUserManager = userManager;
+        mMessageContainerController = messageContainerController;
+        mAssistContentRequester = assistContentRequester;
+        mAnnouncementResolver = announcementResolver;
+
+        mViewProxy = viewProxyFactory.getProxy(mContext, mDisplay.getDisplayId());
+
+        mScreenshotHandler.setOnTimeoutRunnable(() -> {
+            if (DEBUG_UI) {
+                Log.d(TAG, "Corner timeout hit");
+            }
+            mViewProxy.requestDismissal(SCREENSHOT_INTERACTION_TIMEOUT);
+        });
+
+        // Setup the window that we are going to use
+        mWindowLayoutParams = FloatingWindowUtil.getFloatingWindowParams();
+        mWindowLayoutParams.setTitle("ScreenshotAnimation");
+
+        mWindow = FloatingWindowUtil.getFloatingWindow(mContext);
+        mWindow.setWindowManager(mWindowManager, null, null);
+
+        mConfigChanges.applyNewConfig(context.getResources());
+        reloadAssets();
+
+        mActionExecutor = actionExecutorFactory.create(mWindow, mViewProxy,
+                () -> {
+                    finishDismiss();
+                    return Unit.INSTANCE;
+                });
+        mActionsController = screenshotActionsControllerFactory.getController(mActionExecutor);
+
+
+        // Sound is only reproduced from the controller of the default display.
+        if (mDisplay.getDisplayId() == Display.DEFAULT_DISPLAY) {
+            mScreenshotSoundController = screenshotSoundController.get();
+        } else {
+            mScreenshotSoundController = null;
+        }
+
+        mCopyBroadcastReceiver = new BroadcastReceiver() {
+            @Override
+            public void onReceive(Context context, Intent intent) {
+                if (ClipboardOverlayController.COPY_OVERLAY_ACTION.equals(intent.getAction())) {
+                    mViewProxy.requestDismissal(SCREENSHOT_DISMISSED_OTHER);
+                }
+            }
+        };
+        mBroadcastDispatcher.registerReceiver(mCopyBroadcastReceiver, new IntentFilter(
+                        ClipboardOverlayController.COPY_OVERLAY_ACTION), null, null,
+                Context.RECEIVER_NOT_EXPORTED, ClipboardOverlayController.SELF_PERMISSION);
+    }
+
+    @Override
+    public void handleScreenshot(ScreenshotData screenshot, Consumer<Uri> finisher,
+            RequestCallback requestCallback) {
+        Assert.isMainThread();
+
+        mCurrentRequestCallback = requestCallback;
+        if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_FULLSCREEN
+                && screenshot.getBitmap() == null) {
+            Rect bounds = getFullScreenRect();
+            screenshot.setBitmap(mImageCapture.captureDisplay(mDisplay.getDisplayId(), bounds));
+            screenshot.setScreenBounds(bounds);
+        }
+
+        if (screenshot.getBitmap() == null) {
+            Log.e(TAG, "handleScreenshot: Screenshot bitmap was null");
+            mNotificationsController.notifyScreenshotError(
+                    R.string.screenshot_failed_to_capture_text);
+            if (mCurrentRequestCallback != null) {
+                mCurrentRequestCallback.reportError();
+            }
+            return;
+        }
+
+        mScreenBitmap = screenshot.getBitmap();
+        String oldPackageName = mPackageName;
+        mPackageName = screenshot.getPackageNameString();
+
+        if (!isUserSetupComplete(Process.myUserHandle())) {
+            Log.w(TAG, "User setup not complete, displaying toast only");
+            // User setup isn't complete, so we don't want to show any UI beyond a toast, as editing
+            // and sharing shouldn't be exposed to the user.
+            saveScreenshotAndToast(screenshot, finisher);
+            return;
+        }
+
+        mBroadcastSender.sendBroadcast(new Intent(ClipboardOverlayController.SCREENSHOT_ACTION),
+                ClipboardOverlayController.SELF_PERMISSION);
+
+        mScreenshotTakenInPortrait =
+                mContext.getResources().getConfiguration().orientation == ORIENTATION_PORTRAIT;
+
+        // Optimizations
+        mScreenBitmap.setHasAlpha(false);
+        mScreenBitmap.prepareToDraw();
+
+        prepareViewForNewScreenshot(screenshot, oldPackageName);
+
+        final UUID requestId;
+        requestId = mActionsController.setCurrentScreenshot(screenshot);
+        saveScreenshotInBackground(screenshot, requestId, finisher, result -> {
+            if (result.uri != null) {
+                ScreenshotSavedResult savedScreenshot = new ScreenshotSavedResult(
+                        result.uri, screenshot.getUserOrDefault(), result.timestamp);
+                mActionsController.setCompletedScreenshot(requestId, savedScreenshot);
+            }
+        });
+
+        if (screenshot.getTaskId() >= 0) {
+            mAssistContentRequester.requestAssistContent(
+                    screenshot.getTaskId(),
+                    assistContent ->
+                            mActionsController.onAssistContent(requestId, assistContent));
+        } else {
+            mActionsController.onAssistContent(requestId, null);
+        }
+
+        // The window is focusable by default
+        setWindowFocusable(true);
+        mViewProxy.requestFocus();
+
+        enqueueScrollCaptureRequest(requestId, screenshot.getUserHandle());
+
+        attachWindow();
+
+        boolean showFlash;
+        if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE) {
+            if (screenshot.getScreenBounds() != null
+                    && aspectRatiosMatch(screenshot.getBitmap(), screenshot.getInsets(),
+                    screenshot.getScreenBounds())) {
+                showFlash = false;
+            } else {
+                showFlash = true;
+                screenshot.setInsets(Insets.NONE);
+                screenshot.setScreenBounds(new Rect(0, 0, screenshot.getBitmap().getWidth(),
+                        screenshot.getBitmap().getHeight()));
+            }
+        } else {
+            showFlash = true;
+        }
+
+        mViewProxy.prepareEntranceAnimation(
+                () -> startAnimation(screenshot.getScreenBounds(), showFlash,
+                        () -> mMessageContainerController.onScreenshotTaken(screenshot)));
+
+        mViewProxy.setScreenshot(screenshot);
+
+        // ignore system bar insets for the purpose of window layout
+        mWindow.getDecorView().setOnApplyWindowInsetsListener(
+                (v, insets) -> WindowInsets.CONSUMED);
+    }
+
+    void prepareViewForNewScreenshot(@NonNull ScreenshotData screenshot, String oldPackageName) {
+        withWindowAttached(() -> {
+            if (screenshotPrivateProfileAccessibilityAnnouncementFix()) {
+                mAnnouncementResolver.getScreenshotAnnouncement(
+                        screenshot.getUserHandle().getIdentifier(),
+                        mViewProxy::announceForAccessibility);
+            } else {
+                if (mUserManager.isManagedProfile(screenshot.getUserHandle().getIdentifier())) {
+                    mViewProxy.announceForAccessibility(mContext.getResources().getString(
+                            R.string.screenshot_saving_work_profile_title));
+                } else {
+                    mViewProxy.announceForAccessibility(
+                            mContext.getResources().getString(R.string.screenshot_saving_title));
+                }
+            }
+        });
+
+        mViewProxy.reset();
+
+        if (mViewProxy.isAttachedToWindow()) {
+            // if we didn't already dismiss for another reason
+            if (!mViewProxy.isDismissing()) {
+                mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_REENTERED, 0,
+                        oldPackageName);
+            }
+            if (DEBUG_WINDOW) {
+                Log.d(TAG, "saveScreenshot: screenshotView is already attached, resetting. "
+                        + "(dismissing=" + mViewProxy.isDismissing() + ")");
+            }
+        }
+
+        mViewProxy.setPackageName(mPackageName);
+    }
+
+    /**
+     * Requests the view to dismiss the current screenshot (may be ignored, if screenshot is already
+     * being dismissed)
+     */
+    @Override
+    public void requestDismissal(ScreenshotEvent event) {
+        mViewProxy.requestDismissal(event);
+    }
+
+    @Override
+    public boolean isPendingSharedTransition() {
+        return mActionExecutor.isPendingSharedTransition();
+    }
+
+    // Any cleanup needed when the service is being destroyed.
+    @Override
+    public void onDestroy() {
+        if (mSaveInBgTask != null) {
+            // just log success/failure for the pre-existing screenshot
+            mSaveInBgTask.setActionsReadyListener(this::logSuccessOnActionsReady);
+        }
+        removeWindow();
+        releaseMediaPlayer();
+        releaseContext();
+        mBgExecutor.shutdown();
+    }
+
+    /**
+     * Release the constructed window context.
+     */
+    private void releaseContext() {
+        mBroadcastDispatcher.unregisterReceiver(mCopyBroadcastReceiver);
+        mContext.release();
+    }
+
+    private void releaseMediaPlayer() {
+        if (mScreenshotSoundController == null) return;
+        mScreenshotSoundController.releaseScreenshotSoundAsync();
+    }
+
+    /**
+     * Update resources on configuration change. Reinflate for theme/color changes.
+     */
+    private void reloadAssets() {
+        if (DEBUG_UI) {
+            Log.d(TAG, "reloadAssets()");
+        }
+
+        mMessageContainerController.setView(mViewProxy.getView());
+        mViewProxy.setCallbacks(new ScreenshotShelfViewProxy.ScreenshotViewCallback() {
+            @Override
+            public void onUserInteraction() {
+                if (DEBUG_INPUT) {
+                    Log.d(TAG, "onUserInteraction");
+                }
+                mScreenshotHandler.resetTimeout();
+            }
+
+            @Override
+            public void onDismiss() {
+                finishDismiss();
+            }
+
+            @Override
+            public void onTouchOutside() {
+                // TODO(159460485): Remove this when focus is handled properly in the system
+                setWindowFocusable(false);
+            }
+        });
+
+        if (DEBUG_WINDOW) {
+            Log.d(TAG, "setContentView: " + mViewProxy.getView());
+        }
+        mWindow.setContentView(mViewProxy.getView());
+    }
+
+    private void enqueueScrollCaptureRequest(UUID requestId, UserHandle owner) {
+        // Wait until this window is attached to request because it is
+        // the reference used to locate the target window (below).
+        withWindowAttached(() -> {
+            requestScrollCapture(requestId, owner);
+            mWindow.peekDecorView().getViewRootImpl().setActivityConfigCallback(
+                    new ViewRootImpl.ActivityConfigCallback() {
+                        @Override
+                        public void onConfigurationChanged(Configuration overrideConfig,
+                                int newDisplayId) {
+                            if (mConfigChanges.applyNewConfig(mContext.getResources())) {
+                                // Hide the scroll chip until we know it's available in this
+                                // orientation
+                                mActionsController.onScrollChipInvalidated();
+                                // Delay scroll capture eval a bit to allow the underlying activity
+                                // to set up in the new orientation.
+                                mScreenshotHandler.postDelayed(
+                                        () -> requestScrollCapture(requestId, owner), 150);
+                                mViewProxy.updateInsets(
+                                        mWindowManager.getCurrentWindowMetrics().getWindowInsets());
+                                // Screenshot animation calculations won't be valid anymore,
+                                // so just end
+                                if (mScreenshotAnimation != null
+                                        && mScreenshotAnimation.isRunning()) {
+                                    mScreenshotAnimation.end();
+                                }
+                            }
+                        }
+
+                        @Override
+                        public void requestCompatCameraControl(boolean showControl,
+                                boolean transformationApplied,
+                                ICompatCameraControlCallback callback) {
+                            Log.w(TAG, "Unexpected requestCompatCameraControl callback");
+                        }
+                    });
+        });
+    }
+
+    private void requestScrollCapture(UUID requestId, UserHandle owner) {
+        mScrollCaptureExecutor.requestScrollCapture(
+                mDisplay.getDisplayId(),
+                mWindow.getDecorView().getWindowToken(),
+                (response) -> {
+                    mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_IMPRESSION,
+                            0, response.getPackageName());
+                    mActionsController.onScrollChipReady(requestId,
+                            () -> onScrollButtonClicked(owner, response));
+                    return Unit.INSTANCE;
+                }
+        );
+    }
+
+    private void onScrollButtonClicked(UserHandle owner, ScrollCaptureResponse response) {
+        if (DEBUG_INPUT) {
+            Log.d(TAG, "scroll chip tapped");
+        }
+        mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_REQUESTED, 0,
+                response.getPackageName());
+        Bitmap newScreenshot = mImageCapture.captureDisplay(mDisplay.getDisplayId(),
+                getFullScreenRect());
+        if (newScreenshot == null) {
+            Log.e(TAG, "Failed to capture current screenshot for scroll transition!");
+            return;
+        }
+        // delay starting scroll capture to make sure scrim is up before the app moves
+        mViewProxy.prepareScrollingTransition(response, mScreenBitmap, newScreenshot,
+                mScreenshotTakenInPortrait, () -> executeBatchScrollCapture(response, owner));
+    }
+
+    private void executeBatchScrollCapture(ScrollCaptureResponse response, UserHandle owner) {
+        mScrollCaptureExecutor.executeBatchScrollCapture(response,
+                () -> {
+                    final Intent intent = ActionIntentCreator.INSTANCE.createLongScreenshotIntent(
+                            owner, mContext);
+                    mContext.startActivity(intent);
+                },
+                mViewProxy::restoreNonScrollingUi,
+                mViewProxy::startLongScreenshotTransition);
+    }
+
+    private void withWindowAttached(Runnable action) {
+        View decorView = mWindow.getDecorView();
+        if (decorView.isAttachedToWindow()) {
+            action.run();
+        } else {
+            decorView.getViewTreeObserver().addOnWindowAttachListener(
+                    new ViewTreeObserver.OnWindowAttachListener() {
+                        @Override
+                        public void onWindowAttached() {
+                            mAttachRequested = false;
+                            decorView.getViewTreeObserver().removeOnWindowAttachListener(this);
+                            action.run();
+                        }
+
+                        @Override
+                        public void onWindowDetached() {
+                        }
+                    });
+
+        }
+    }
+
+    @MainThread
+    private void attachWindow() {
+        View decorView = mWindow.getDecorView();
+        if (decorView.isAttachedToWindow() || mAttachRequested) {
+            return;
+        }
+        if (DEBUG_WINDOW) {
+            Log.d(TAG, "attachWindow");
+        }
+        mAttachRequested = true;
+        mWindowManager.addView(decorView, mWindowLayoutParams);
+        decorView.requestApplyInsets();
+
+        ViewGroup layout = decorView.requireViewById(android.R.id.content);
+        layout.setClipChildren(false);
+        layout.setClipToPadding(false);
+    }
+
+    @Override
+    public void removeWindow() {
+        final View decorView = mWindow.peekDecorView();
+        if (decorView != null && decorView.isAttachedToWindow()) {
+            if (DEBUG_WINDOW) {
+                Log.d(TAG, "Removing screenshot window");
+            }
+            mWindowManager.removeViewImmediate(decorView);
+            mDetachRequested = false;
+        }
+        if (mAttachRequested && !mDetachRequested) {
+            mDetachRequested = true;
+            withWindowAttached(this::removeWindow);
+        }
+
+        mViewProxy.stopInputListening();
+    }
+
+    private void playCameraSoundIfNeeded() {
+        if (mScreenshotSoundController == null) return;
+        // the controller is not-null only on the default display controller
+        mScreenshotSoundController.playScreenshotSoundAsync();
+    }
+
+    /**
+     * Save the bitmap but don't show the normal screenshot UI.. just a toast (or notification on
+     * failure).
+     */
+    private void saveScreenshotAndToast(ScreenshotData screenshot, Consumer<Uri> finisher) {
+        // Play the shutter sound to notify that we've taken a screenshot
+        playCameraSoundIfNeeded();
+
+        if (screenshotSaveImageExporter()) {
+            saveScreenshotInBackground(screenshot, UUID.randomUUID(), finisher, result -> {
+                if (result.uri != null) {
+                    mScreenshotHandler.post(() -> Toast.makeText(mContext,
+                            R.string.screenshot_saved_title, Toast.LENGTH_SHORT).show());
+                }
+            });
+        } else {
+            saveScreenshotInWorkerThread(
+                    screenshot.getUserHandle(),
+                    /* onComplete */ finisher,
+                    /* actionsReadyListener */ imageData -> {
+                        if (DEBUG_CALLBACK) {
+                            Log.d(TAG,
+                                    "returning URI to finisher (Consumer<URI>): " + imageData.uri);
+                        }
+                        finisher.accept(imageData.uri);
+                        if (imageData.uri == null) {
+                            mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_NOT_SAVED, 0,
+                                    mPackageName);
+                            mNotificationsController.notifyScreenshotError(
+                                    R.string.screenshot_failed_to_save_text);
+                        } else {
+                            mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_SAVED, 0, mPackageName);
+                            mScreenshotHandler.post(() -> Toast.makeText(mContext,
+                                    R.string.screenshot_saved_title, Toast.LENGTH_SHORT).show());
+                        }
+                    },
+                    null);
+        }
+    }
+
+    /**
+     * Starts the animation after taking the screenshot
+     */
+    private void startAnimation(Rect screenRect, boolean showFlash, Runnable onAnimationComplete) {
+        if (mScreenshotAnimation != null && mScreenshotAnimation.isRunning()) {
+            mScreenshotAnimation.cancel();
+        }
+
+        mScreenshotAnimation =
+                mViewProxy.createScreenshotDropInAnimation(screenRect, showFlash);
+        if (onAnimationComplete != null) {
+            mScreenshotAnimation.addListener(new AnimatorListenerAdapter() {
+                @Override
+                public void onAnimationEnd(Animator animation) {
+                    super.onAnimationEnd(animation);
+                    onAnimationComplete.run();
+                }
+            });
+        }
+
+        // Play the shutter sound to notify that we've taken a screenshot
+        playCameraSoundIfNeeded();
+
+        if (DEBUG_ANIM) {
+            Log.d(TAG, "starting post-screenshot animation");
+        }
+        mScreenshotAnimation.start();
+    }
+
+    /** Reset screenshot view and then call onCompleteRunnable */
+    private void finishDismiss() {
+        Log.d(TAG, "finishDismiss");
+        mActionsController.endScreenshotSession();
+        mScrollCaptureExecutor.close();
+        if (mCurrentRequestCallback != null) {
+            mCurrentRequestCallback.onFinish();
+            mCurrentRequestCallback = null;
+        }
+        mViewProxy.reset();
+        removeWindow();
+        mScreenshotHandler.cancelTimeout();
+    }
+
+    private void saveScreenshotInBackground(ScreenshotData screenshot, UUID requestId,
+            Consumer<Uri> finisher, Consumer<ImageExporter.Result> onResult) {
+        ListenableFuture<ImageExporter.Result> future = mImageExporter.export(mBgExecutor,
+                requestId, screenshot.getBitmap(), screenshot.getUserOrDefault(),
+                mDisplay.getDisplayId());
+        future.addListener(() -> {
+            try {
+                ImageExporter.Result result = future.get();
+                Log.d(TAG, "Saved screenshot: " + result);
+                logScreenshotResultStatus(result.uri, screenshot.getUserHandle());
+                onResult.accept(result);
+                if (DEBUG_CALLBACK) {
+                    Log.d(TAG, "finished background processing, Calling (Consumer<Uri>) "
+                            + "finisher.accept(\"" + result.uri + "\"");
+                }
+                finisher.accept(result.uri);
+            } catch (Exception e) {
+                Log.d(TAG, "Failed to store screenshot", e);
+                if (DEBUG_CALLBACK) {
+                    Log.d(TAG, "Calling (Consumer<Uri>) finisher.accept(null)");
+                }
+                finisher.accept(null);
+            }
+        }, mMainExecutor);
+    }
+
+    /**
+     * Creates a new worker thread and saves the screenshot to the media store.
+     */
+    private void saveScreenshotInWorkerThread(
+            UserHandle owner,
+            @NonNull Consumer<Uri> finisher,
+            @Nullable SaveImageInBackgroundTask.ActionsReadyListener actionsReadyListener,
+            @Nullable SaveImageInBackgroundTask.QuickShareActionReadyListener
+                    quickShareActionsReadyListener) {
+        SaveImageInBackgroundTask.SaveImageInBackgroundData
+                data = new SaveImageInBackgroundTask.SaveImageInBackgroundData();
+        data.image = mScreenBitmap;
+        data.finisher = finisher;
+        data.mActionsReadyListener = actionsReadyListener;
+        data.mQuickShareActionsReadyListener = quickShareActionsReadyListener;
+        data.owner = owner;
+        data.displayId = mDisplay.getDisplayId();
+
+        if (mSaveInBgTask != null) {
+            // just log success/failure for the pre-existing screenshot
+            mSaveInBgTask.setActionsReadyListener(this::logSuccessOnActionsReady);
+        }
+
+        mSaveInBgTask = new SaveImageInBackgroundTask(mContext, mFlags, mImageExporter,
+                mScreenshotSmartActions, data,
+                mScreenshotNotificationSmartActionsProvider);
+        mSaveInBgTask.execute();
+    }
+
+    /**
+     * Logs success/failure of the screenshot saving task, and shows an error if it failed.
+     */
+    private void logScreenshotResultStatus(Uri uri, UserHandle owner) {
+        if (uri == null) {
+            mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_NOT_SAVED, 0, mPackageName);
+            mNotificationsController.notifyScreenshotError(
+                    R.string.screenshot_failed_to_save_text);
+        } else {
+            mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_SAVED, 0, mPackageName);
+            if (mUserManager.isManagedProfile(owner.getIdentifier())) {
+                mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_SAVED_TO_WORK_PROFILE, 0,
+                        mPackageName);
+            }
+        }
+    }
+
+    /**
+     * Logs success/failure of the screenshot saving task, and shows an error if it failed.
+     */
+    private void logSuccessOnActionsReady(SaveImageInBackgroundTask.SavedImageData imageData) {
+        logScreenshotResultStatus(imageData.uri, imageData.owner);
+    }
+
+    private boolean isUserSetupComplete(UserHandle owner) {
+        return Settings.Secure.getInt(mContext.createContextAsUser(owner, 0)
+                .getContentResolver(), SETTINGS_SECURE_USER_SETUP_COMPLETE, 0) == 1;
+    }
+
+    /**
+     * Updates the window focusability.  If the window is already showing, then it updates the
+     * window immediately, otherwise the layout params will be applied when the window is next
+     * shown.
+     */
+    private void setWindowFocusable(boolean focusable) {
+        if (DEBUG_WINDOW) {
+            Log.d(TAG, "setWindowFocusable: " + focusable);
+        }
+        int flags = mWindowLayoutParams.flags;
+        if (focusable) {
+            mWindowLayoutParams.flags &= ~WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
+        } else {
+            mWindowLayoutParams.flags |= WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
+        }
+        if (mWindowLayoutParams.flags == flags) {
+            if (DEBUG_WINDOW) {
+                Log.d(TAG, "setWindowFocusable: skipping, already " + focusable);
+            }
+            return;
+        }
+        final View decorView = mWindow.peekDecorView();
+        if (decorView != null && decorView.isAttachedToWindow()) {
+            mWindowManager.updateViewLayout(decorView, mWindowLayoutParams);
+        }
+    }
+
+    private Rect getFullScreenRect() {
+        DisplayMetrics displayMetrics = new DisplayMetrics();
+        mDisplay.getRealMetrics(displayMetrics);
+        return new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels);
+    }
+
+    /** Does the aspect ratio of the bitmap with insets removed match the bounds. */
+    private static boolean aspectRatiosMatch(Bitmap bitmap, Insets bitmapInsets,
+            Rect screenBounds) {
+        int insettedWidth = bitmap.getWidth() - bitmapInsets.left - bitmapInsets.right;
+        int insettedHeight = bitmap.getHeight() - bitmapInsets.top - bitmapInsets.bottom;
+
+        if (insettedHeight == 0 || insettedWidth == 0 || bitmap.getWidth() == 0
+                || bitmap.getHeight() == 0) {
+            if (DEBUG_UI) {
+                Log.e(TAG, "Provided bitmap and insets create degenerate region: "
+                        + bitmap.getWidth() + "x" + bitmap.getHeight() + " " + bitmapInsets);
+            }
+            return false;
+        }
+
+        float insettedBitmapAspect = ((float) insettedWidth) / insettedHeight;
+        float boundsAspect = ((float) screenBounds.width()) / screenBounds.height();
+
+        boolean matchWithinTolerance = Math.abs(insettedBitmapAspect - boundsAspect) < 0.1f;
+        if (DEBUG_UI) {
+            Log.d(TAG, "aspectRatiosMatch: don't match bitmap: " + insettedBitmapAspect
+                    + ", bounds: " + boundsAspect);
+        }
+        return matchWithinTolerance;
+    }
+
+    /** Injectable factory to create screenshot controller instances for a specific display. */
+    @AssistedFactory
+    public interface Factory extends InteractiveScreenshotHandler.Factory {
+        /**
+         * Creates an instance of the controller for that specific display.
+         *
+         * @param display                 display to capture
+         */
+        LegacyScreenshotController create(Display display);
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/SaveImageInBackgroundTask.java b/packages/SystemUI/src/com/android/systemui/screenshot/SaveImageInBackgroundTask.java
index 54ae225..9bc3bd8 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/SaveImageInBackgroundTask.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/SaveImageInBackgroundTask.java
@@ -286,7 +286,7 @@
                     ScreenshotNotificationSmartActionsProvider.ACTION_TYPE,
                     ScreenshotNotificationSmartActionsProvider.DEFAULT_ACTION_TYPE);
             Intent intent = new Intent(context, SmartActionsReceiver.class)
-                    .putExtra(ScreenshotController.EXTRA_ACTION_INTENT, action.actionIntent)
+                    .putExtra(SmartActionsReceiver.EXTRA_ACTION_INTENT, action.actionIntent)
                     .addFlags(Intent.FLAG_RECEIVER_FOREGROUND);
             addIntentExtras(mScreenshotId, intent, actionType, true /* smartActionsEnabled */);
             PendingIntent broadcastIntent = PendingIntent.getBroadcast(context,
@@ -302,9 +302,9 @@
     private static void addIntentExtras(String screenshotId, Intent intent, String actionType,
             boolean smartActionsEnabled) {
         intent
-                .putExtra(ScreenshotController.EXTRA_ACTION_TYPE, actionType)
-                .putExtra(ScreenshotController.EXTRA_ID, screenshotId)
-                .putExtra(ScreenshotController.EXTRA_SMART_ACTIONS_ENABLED, smartActionsEnabled);
+                .putExtra(SmartActionsReceiver.EXTRA_ACTION_TYPE, actionType)
+                .putExtra(SmartActionsReceiver.EXTRA_ID, screenshotId)
+                .putExtra(SmartActionsReceiver.EXTRA_SMART_ACTIONS_ENABLED, smartActionsEnabled);
     }
 
     /**
@@ -327,8 +327,8 @@
         }
 
         Intent wrappedIntent = new Intent(mContext, SmartActionsReceiver.class)
-                .putExtra(ScreenshotController.EXTRA_ACTION_INTENT, quickShare.actionIntent)
-                .putExtra(ScreenshotController.EXTRA_ACTION_INTENT_FILLIN,
+                .putExtra(SmartActionsReceiver.EXTRA_ACTION_INTENT, quickShare.actionIntent)
+                .putExtra(SmartActionsReceiver.EXTRA_ACTION_INTENT_FILLIN,
                         createFillInIntent(uri, imageTime))
                 .addFlags(Intent.FLAG_RECEIVER_FOREGROUND);
         Bundle extras = quickShare.getExtras();
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index 0a4635e..653e49f 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -95,28 +95,9 @@
 /**
  * Controls the state and flow for screenshots.
  */
-public class ScreenshotController implements ScreenshotHandler {
+public class ScreenshotController implements InteractiveScreenshotHandler {
     private static final String TAG = logTag(ScreenshotController.class);
 
-    public interface TransitionDestination {
-        /**
-         * Allows the long screenshot activity to call back with a destination location (the bounds
-         * on screen of the destination for the transitioning view) and a Runnable to be run once
-         * the transition animation is complete.
-         */
-        void setTransitionDestination(Rect transitionDestination, Runnable onTransitionEnd);
-    }
-
-    // These strings are used for communicating the action invoked to
-    // ScreenshotNotificationSmartActionsProvider.
-    public static final String EXTRA_ACTION_TYPE = "android:screenshot_action_type";
-    public static final String EXTRA_ID = "android:screenshot_id";
-    public static final String EXTRA_SMART_ACTIONS_ENABLED = "android:smart_actions_enabled";
-    public static final String EXTRA_ACTION_INTENT = "android:screenshot_action_intent";
-    public static final String EXTRA_ACTION_INTENT_FILLIN =
-            "android:screenshot_action_intent_fillin";
-
-
     // From WizardManagerHelper.java
     private static final String SETTINGS_SECURE_USER_SETUP_COMPLETE = "user_setup_complete";
 
@@ -378,9 +359,7 @@
             if (screenshotPrivateProfileAccessibilityAnnouncementFix()) {
                 mAnnouncementResolver.getScreenshotAnnouncement(
                         screenshot.getUserHandle().getIdentifier(),
-                        announcement -> {
-                            mViewProxy.announceForAccessibility(announcement);
-                        });
+                        mViewProxy::announceForAccessibility);
             } else {
                 if (mUserManager.isManagedProfile(screenshot.getUserHandle().getIdentifier())) {
                     mViewProxy.announceForAccessibility(mContext.getResources().getString(
@@ -413,16 +392,19 @@
      * Requests the view to dismiss the current screenshot (may be ignored, if screenshot is already
      * being dismissed)
      */
-    void requestDismissal(ScreenshotEvent event) {
+    @Override
+    public void requestDismissal(ScreenshotEvent event) {
         mViewProxy.requestDismissal(event);
     }
 
-    boolean isPendingSharedTransition() {
+    @Override
+    public boolean isPendingSharedTransition() {
         return mActionExecutor.isPendingSharedTransition();
     }
 
     // Any cleanup needed when the service is being destroyed.
-    void onDestroy() {
+    @Override
+    public void onDestroy() {
         if (mSaveInBgTask != null) {
             // just log success/failure for the pre-existing screenshot
             mSaveInBgTask.setActionsReadyListener(this::logSuccessOnActionsReady);
@@ -603,7 +585,8 @@
         layout.setClipToPadding(false);
     }
 
-    void removeWindow() {
+    @Override
+    public void removeWindow() {
         final View decorView = mWindow.peekDecorView();
         if (decorView != null && decorView.isAttachedToWindow()) {
             if (DEBUG_WINDOW) {
@@ -854,12 +837,12 @@
 
     /** Injectable factory to create screenshot controller instances for a specific display. */
     @AssistedFactory
-    public interface Factory {
+    public interface Factory extends InteractiveScreenshotHandler.Factory {
         /**
          * Creates an instance of the controller for that specific display.
          *
          * @param display                 display to capture
          */
-        ScreenshotController create(Display display);
+        LegacyScreenshotController create(Display display);
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/SmartActionsReceiver.java b/packages/SystemUI/src/com/android/systemui/screenshot/SmartActionsReceiver.java
index f8b22a6..f902693 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/SmartActionsReceiver.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/SmartActionsReceiver.java
@@ -17,10 +17,6 @@
 package com.android.systemui.screenshot;
 
 import static com.android.systemui.screenshot.LogConfig.DEBUG_ACTIONS;
-import static com.android.systemui.screenshot.ScreenshotController.EXTRA_ACTION_INTENT;
-import static com.android.systemui.screenshot.ScreenshotController.EXTRA_ACTION_INTENT_FILLIN;
-import static com.android.systemui.screenshot.ScreenshotController.EXTRA_ACTION_TYPE;
-import static com.android.systemui.screenshot.ScreenshotController.EXTRA_ID;
 
 import android.app.ActivityOptions;
 import android.app.PendingIntent;
@@ -36,6 +32,15 @@
  */
 public class SmartActionsReceiver extends BroadcastReceiver {
     private static final String TAG = "SmartActionsReceiver";
+    // These strings are used for communicating the action invoked to
+    // ScreenshotNotificationSmartActionsProvider.
+    public static final String EXTRA_ACTION_TYPE = "android:screenshot_action_type";
+    public static final String EXTRA_ID = "android:screenshot_id";
+    public static final String EXTRA_SMART_ACTIONS_ENABLED = "android:smart_actions_enabled";
+    public static final String EXTRA_ACTION_INTENT = "android:screenshot_action_intent";
+    public static final String EXTRA_ACTION_INTENT_FILLIN =
+            "android:screenshot_action_intent_fillin";
+
 
     private final ScreenshotSmartActions mScreenshotSmartActions;
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
index 07f6e85..50ea3bb 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
@@ -74,7 +74,7 @@
 class TakeScreenshotExecutorImpl
 @Inject
 constructor(
-    private val screenshotControllerFactory: ScreenshotController.Factory,
+    private val interactiveScreenshotHandlerFactory: InteractiveScreenshotHandler.Factory,
     displayRepository: DisplayRepository,
     @Application private val mainScope: CoroutineScope,
     private val screenshotRequestProcessor: ScreenshotRequestProcessor,
@@ -83,7 +83,7 @@
     private val headlessScreenshotHandler: HeadlessScreenshotHandler,
 ) : TakeScreenshotExecutor {
     private val displays = displayRepository.displays
-    private var screenshotController: ScreenshotController? = null
+    private var screenshotController: InteractiveScreenshotHandler? = null
     private val notificationControllers = mutableMapOf<Int, ScreenshotNotificationsController>()
 
     /**
@@ -183,7 +183,7 @@
 
     /** Propagates the close system dialog signal to the ScreenshotController. */
     override fun onCloseSystemDialogsReceived() {
-        if (screenshotController?.isPendingSharedTransition == false) {
+        if (screenshotController?.isPendingSharedTransition() == false) {
             screenshotController?.requestDismissal(SCREENSHOT_DISMISSED_OTHER)
         }
     }
@@ -218,8 +218,9 @@
         }
     }
 
-    private fun getScreenshotController(display: Display): ScreenshotController {
-        val controller = screenshotController ?: screenshotControllerFactory.create(display)
+    private fun getScreenshotController(display: Display): InteractiveScreenshotHandler {
+        val controller =
+            screenshotController ?: interactiveScreenshotHandlerFactory.create(display)
         screenshotController = controller
         return controller
     }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
index 8feefa4..9db1f24 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
@@ -46,13 +46,17 @@
 import android.os.ResultReceiver;
 import android.util.Log;
 import android.view.View;
+import android.view.ViewGroup;
+import android.widget.ArrayAdapter;
 import android.widget.Button;
 import android.widget.CheckBox;
 import android.widget.ImageView;
+import android.widget.ListPopupWindow;
 import android.widget.TextView;
 
 import androidx.activity.ComponentActivity;
 import androidx.annotation.Nullable;
+import androidx.appcompat.content.res.AppCompatResources;
 import androidx.core.graphics.Insets;
 import androidx.core.view.ViewCompat;
 import androidx.core.view.WindowInsetsCompat;
@@ -67,6 +71,7 @@
 import com.android.systemui.screenshot.scroll.CropView;
 import com.android.systemui.settings.UserTracker;
 
+import java.util.List;
 import java.util.Set;
 
 import javax.inject.Inject;
@@ -92,6 +97,7 @@
 
     private static final String TAG = AppClipsActivity.class.getSimpleName();
     private static final ApplicationInfoFlags APPLICATION_INFO_FLAGS = ApplicationInfoFlags.of(0);
+    private static final int DRAWABLE_END = 2;
 
     private final AppClipsViewModel.Factory mViewModelFactory;
     private final PackageManager mPackageManager;
@@ -192,6 +198,7 @@
         mViewModel.getResultLiveData().observe(this, this::setResultThenFinish);
         mViewModel.getErrorLiveData().observe(this, this::setErrorThenFinish);
         mViewModel.getBacklinksLiveData().observe(this, this::setBacklinksData);
+        mViewModel.mSelectedBacklinksLiveData.observe(this, this::updateBacklinksTextView);
 
         if (savedInstanceState == null) {
             int displayId = getDisplayId();
@@ -305,8 +312,8 @@
 
         if (mBacklinksIncludeDataCheckBox.getVisibility() == View.VISIBLE
                 && mBacklinksIncludeDataCheckBox.isChecked()
-                && mViewModel.getBacklinksLiveData().getValue() != null) {
-            ClipData backlinksData = mViewModel.getBacklinksLiveData().getValue().getClipData();
+                && mViewModel.mSelectedBacklinksLiveData.getValue() != null) {
+            ClipData backlinksData = mViewModel.mSelectedBacklinksLiveData.getValue().getClipData();
             data.putParcelable(EXTRA_CLIP_DATA, backlinksData);
 
             DebugLogger.INSTANCE.logcatMessage(this,
@@ -330,18 +337,80 @@
         finish();
     }
 
-    private void setBacklinksData(InternalBacklinksData backlinksData) {
+    private void setBacklinksData(List<InternalBacklinksData> backlinksData) {
         mBacklinksIncludeDataCheckBox.setVisibility(View.VISIBLE);
         mBacklinksDataTextView.setVisibility(
                 mBacklinksIncludeDataCheckBox.isChecked() ? View.VISIBLE : View.GONE);
 
-        mBacklinksDataTextView.setText(backlinksData.getClipData().getDescription().getLabel());
+        // Set up the dropdown when multiple backlinks are available.
+        if (backlinksData.size() > 1) {
+            setUpListPopupWindow(backlinksData, mBacklinksDataTextView);
+        }
+    }
 
+    private void setUpListPopupWindow(List<InternalBacklinksData> backlinksData, View anchor) {
+        ListPopupWindow listPopupWindow = new ListPopupWindow(this);
+        listPopupWindow.setAnchorView(anchor);
+        listPopupWindow.setOverlapAnchor(true);
+        listPopupWindow.setBackgroundDrawable(
+                AppCompatResources.getDrawable(this, R.drawable.backlinks_rounded_rectangle));
+        listPopupWindow.setOnItemClickListener((parent, view, position, id) -> {
+            mViewModel.mSelectedBacklinksLiveData.setValue(backlinksData.get(position));
+            listPopupWindow.dismiss();
+        });
+
+        ArrayAdapter<InternalBacklinksData> adapter = new ArrayAdapter<>(this,
+                R.layout.app_clips_backlinks_drop_down_entry) {
+            @Override
+            public View getView(int position, @Nullable View convertView, ViewGroup parent) {
+                TextView itemView = (TextView) super.getView(position, convertView, parent);
+                InternalBacklinksData data = backlinksData.get(position);
+                itemView.setText(data.getClipData().getDescription().getLabel());
+
+                Drawable icon = data.getAppIcon();
+                icon.setBounds(createBacklinksTextViewDrawableBounds());
+                itemView.setCompoundDrawablesRelative(/* start= */ icon, /* top= */ null,
+                        /* end= */ null, /* bottom= */ null);
+
+                return itemView;
+            }
+        };
+        adapter.addAll(backlinksData);
+        listPopupWindow.setAdapter(adapter);
+
+        mBacklinksDataTextView.setOnClickListener(unused -> listPopupWindow.show());
+    }
+
+    /**
+     * Updates the {@link #mBacklinksDataTextView} with the currently selected
+     * {@link InternalBacklinksData}. The {@link AppClipsViewModel#getBacklinksLiveData()} is
+     * expected to be already set when this method is called.
+     */
+    private void updateBacklinksTextView(InternalBacklinksData backlinksData) {
+        mBacklinksDataTextView.setText(backlinksData.getClipData().getDescription().getLabel());
         Drawable appIcon = backlinksData.getAppIcon();
-        int size = getResources().getDimensionPixelSize(R.dimen.appclips_backlinks_icon_size);
-        appIcon.setBounds(/* left= */ 0, /* top= */ 0, /* right= */ size, /* bottom= */ size);
+        Rect compoundDrawableBounds = createBacklinksTextViewDrawableBounds();
+        appIcon.setBounds(compoundDrawableBounds);
+
+        // Try to reuse the dropdown down arrow icon if available, will be null if never set.
+        Drawable dropDownIcon = mBacklinksDataTextView.getCompoundDrawablesRelative()[DRAWABLE_END];
+        if (mViewModel.getBacklinksLiveData().getValue().size() > 1 && dropDownIcon == null) {
+            // Set up the dropdown down arrow drawable only if it is required.
+            dropDownIcon = AppCompatResources.getDrawable(this, R.drawable.arrow_pointing_down);
+            dropDownIcon.setBounds(compoundDrawableBounds);
+            dropDownIcon.setTint(Utils.getColorAttr(this,
+                    android.R.attr.textColorSecondary).getDefaultColor());
+        }
+
         mBacklinksDataTextView.setCompoundDrawablesRelative(/* start= */ appIcon, /* top= */
-                null, /* end= */ null, /* bottom= */ null);
+                null, /* end= */ dropDownIcon, /* bottom= */ null);
+    }
+
+    private Rect createBacklinksTextViewDrawableBounds() {
+        int size = getResources().getDimensionPixelSize(R.dimen.appclips_backlinks_icon_size);
+        Rect bounds = new Rect();
+        bounds.set(/* left= */ 0, /* top= */ 0, /* right= */ size, /* bottom= */ size);
+        return bounds;
     }
 
     private void setError(int errorCode) {
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java
index bd9e295..3530b3f 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsViewModel.java
@@ -21,12 +21,11 @@
 import static android.content.Intent.CAPTURE_CONTENT_FOR_NOTE_FAILED;
 import static android.content.Intent.CATEGORY_LAUNCHER;
 
-import static com.google.common.util.concurrent.Futures.withTimeout;
-
 import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;
 
 import android.app.ActivityTaskManager.RootTaskInfo;
 import android.app.IActivityTaskManager;
+import android.app.TaskInfo;
 import android.app.WindowConfiguration;
 import android.app.assist.AssistContent;
 import android.content.ClipData;
@@ -41,7 +40,6 @@
 import android.graphics.RenderNode;
 import android.graphics.drawable.Drawable;
 import android.net.Uri;
-import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.Log;
 import android.view.Display;
@@ -94,7 +92,8 @@
     private final MutableLiveData<Bitmap> mScreenshotLiveData;
     private final MutableLiveData<Uri> mResultLiveData;
     private final MutableLiveData<Integer> mErrorLiveData;
-    private final MutableLiveData<InternalBacklinksData> mBacklinksLiveData;
+    private final MutableLiveData<List<InternalBacklinksData>> mBacklinksLiveData;
+    final MutableLiveData<InternalBacklinksData> mSelectedBacklinksLiveData;
 
     private AppClipsViewModel(AppClipsCrossProcessHelper appClipsCrossProcessHelper,
             ImageExporter imageExporter, IActivityTaskManager atmService,
@@ -112,6 +111,7 @@
         mResultLiveData = new MutableLiveData<>();
         mErrorLiveData = new MutableLiveData<>();
         mBacklinksLiveData = new MutableLiveData<>();
+        mSelectedBacklinksLiveData = new MutableLiveData<>();
     }
 
     /**
@@ -135,10 +135,11 @@
     /**
      * Triggers the Backlinks flow which:
      * <ul>
-     *     <li>Evaluates the task to query.
-     *     <li>Requests {@link AssistContent} from that task.
-     *     <li>Transforms the {@link AssistContent} into {@link ClipData} for Backlinks.
-     *     <li>The {@link ClipData} is reported to activity via {@link #getBacklinksLiveData()}.
+     *     <li>Evaluates the tasks to query.
+     *     <li>Requests {@link AssistContent} from all valid tasks.
+     *     <li>Transforms {@link AssistContent} into {@link InternalBacklinksData} for Backlinks.
+     *     <li>The {@link InternalBacklinksData}s are reported to activity via
+     *     {@link #getBacklinksLiveData()}.
      * </ul>
      *
      * @param taskIdsToIgnore id of the tasks to ignore when querying for {@link AssistContent}
@@ -146,24 +147,24 @@
      */
     void triggerBacklinks(Set<Integer> taskIdsToIgnore, int displayId) {
         DebugLogger.INSTANCE.logcatMessage(this, () -> "Backlinks triggered");
-        mBgExecutor.execute(() -> {
-            ListenableFuture<InternalBacklinksData> backlinksData = getBacklinksData(
-                    taskIdsToIgnore, displayId);
-            Futures.addCallback(backlinksData, new FutureCallback<>() {
-                @Override
-                public void onSuccess(@Nullable InternalBacklinksData result) {
-                    if (result != null) {
-                        mBacklinksLiveData.setValue(result);
-                    }
+        ListenableFuture<List<InternalBacklinksData>> backlinksData = getAllAvailableBacklinks(
+                taskIdsToIgnore, displayId);
+        Futures.addCallback(backlinksData, new FutureCallback<>() {
+            @Override
+            public void onSuccess(@Nullable List<InternalBacklinksData> result) {
+                if (result != null && !result.isEmpty()) {
+                    // Set the list of backlinks before setting the selected backlink as this is
+                    // required when updating the backlink data text view.
+                    mBacklinksLiveData.setValue(result);
+                    mSelectedBacklinksLiveData.setValue(result.get(0));
                 }
+            }
 
-                @Override
-                public void onFailure(Throwable t) {
-                    Log.e(TAG, "Error querying for Backlinks data", t);
-                }
-            }, mMainExecutor);
-
-        });
+            @Override
+            public void onFailure(Throwable t) {
+                Log.e(TAG, "Error querying for Backlinks data", t);
+            }
+        }, mMainExecutor);
     }
 
     /** Returns a {@link LiveData} that holds the captured screenshot. */
@@ -184,8 +185,11 @@
         return mErrorLiveData;
     }
 
-    /** Returns a {@link LiveData} that holds Backlinks data in {@link InternalBacklinksData}. */
-    LiveData<InternalBacklinksData> getBacklinksLiveData() {
+    /**
+     * Returns a {@link LiveData} that holds all the available Backlinks data and the currently
+     * selected index for displaying the Backlinks in the UI.
+     */
+    LiveData<List<InternalBacklinksData>> getBacklinksLiveData() {
         return mBacklinksLiveData;
     }
 
@@ -230,26 +234,58 @@
         return HardwareRenderer.createHardwareBitmap(output, bounds.width(), bounds.height());
     }
 
-    private ListenableFuture<InternalBacklinksData> getBacklinksData(Set<Integer> taskIdsToIgnore,
-            int displayId) {
-        return getAllRootTaskInfosOnDisplay(displayId)
-                .stream()
-                .filter(taskInfo -> shouldIncludeTask(taskInfo, taskIdsToIgnore))
-                .findFirst()
-                .map(this::getBacklinksDataForTaskId)
-                .orElse(Futures.immediateFuture(null));
+    private ListenableFuture<List<InternalBacklinksData>> getAllAvailableBacklinks(
+            Set<Integer> taskIdsToIgnore, int displayId) {
+        ListenableFuture<List<TaskInfo>> allTasksOnDisplayFuture = getAllTasksOnDisplay(displayId);
+
+        ListenableFuture<List<ListenableFuture<InternalBacklinksData>>> backlinksNestedListFuture =
+                Futures.transform(allTasksOnDisplayFuture, allTasksOnDisplay ->
+                        allTasksOnDisplay
+                                .stream()
+                                .filter(taskInfo -> shouldIncludeTask(taskInfo, taskIdsToIgnore))
+                                .map(this::getBacklinksDataForTaskInfo)
+                                .toList(),
+                        mBgExecutor);
+
+        return Futures.transformAsync(backlinksNestedListFuture, Futures::allAsList, mBgExecutor);
     }
 
-    private List<RootTaskInfo> getAllRootTaskInfosOnDisplay(int displayId) {
-        try {
-            return mAtmService.getAllRootTaskInfosOnDisplay(displayId);
-        } catch (RemoteException e) {
-            Log.e(TAG, String.format("Error while querying for tasks on display %d", displayId), e);
-            return Collections.emptyList();
-        }
+    /**
+     * Returns all tasks on a given display after querying {@link IActivityTaskManager} from the
+     * {@link #mBgExecutor}.
+     */
+    private ListenableFuture<List<TaskInfo>> getAllTasksOnDisplay(int displayId) {
+        SettableFuture<List<TaskInfo>> recentTasksFuture = SettableFuture.create();
+        mBgExecutor.execute(() -> {
+            try {
+                // Directly call into ActivityTaskManagerService instead of going through WMShell
+                // because WMShell is only available in the main SysUI process and App Clips runs
+                // in its own separate process as it deals with bitmaps.
+                List<TaskInfo> allTasksOnDisplay = mAtmService.getTasks(
+                                /* maxNum= */ Integer.MAX_VALUE,
+                                // PIP tasks are not visible in recents. So _not_ filtering for
+                                // tasks that are only visible in recents.
+                                /* filterOnlyVisibleRecents= */ false,
+                                /* keepIntentExtra= */ false,
+                                displayId)
+                        .stream()
+                        .map(runningTaskInfo -> (TaskInfo) runningTaskInfo)
+                        .toList();
+                recentTasksFuture.set(allTasksOnDisplay);
+            } catch (Exception e) {
+                Log.e(TAG, String.format("Error getting all tasks on displayId %d", displayId), e);
+                recentTasksFuture.set(Collections.emptyList());
+            }
+        });
+
+        return withTimeout(recentTasksFuture);
     }
 
-    private boolean shouldIncludeTask(RootTaskInfo taskInfo, Set<Integer> taskIdsToIgnore) {
+    /**
+     * Returns whether the app represented by the provided {@link TaskInfo} should be included for
+     * querying for {@link AssistContent}.
+     */
+    private boolean shouldIncludeTask(TaskInfo taskInfo, Set<Integer> taskIdsToIgnore) {
         DebugLogger.INSTANCE.logcatMessage(this,
                 () -> String.format("shouldIncludeTask taskId %d; topActivity %s", taskInfo.taskId,
                         taskInfo.topActivity));
@@ -262,11 +298,14 @@
                 && taskInfo.numActivities > 0
                 && taskInfo.topActivity != null
                 && taskInfo.topActivityInfo != null
-                && taskInfo.childTaskIds.length > 0
                 && taskInfo.getActivityType() == WindowConfiguration.ACTIVITY_TYPE_STANDARD
                 && canAppStartThroughLauncher(taskInfo.topActivity.getPackageName());
     }
 
+    /**
+     * Returns whether the app represented by the provided {@code packageName} can be launched
+     * through the all apps tray by a user.
+     */
     private boolean canAppStartThroughLauncher(String packageName) {
         // Use Intent.resolveActivity API to check if the intent resolves as that is what Android
         // uses internally when apps use Context.startActivity.
@@ -274,8 +313,12 @@
                 != null;
     }
 
-    private ListenableFuture<InternalBacklinksData> getBacklinksDataForTaskId(
-            RootTaskInfo taskInfo) {
+    /**
+     * Returns an {@link InternalBacklinksData} that represents the Backlink data internally, which
+     * is captured by querying the system using {@link TaskInfo#taskId}.
+     */
+    private ListenableFuture<InternalBacklinksData> getBacklinksDataForTaskInfo(
+            TaskInfo taskInfo) {
         DebugLogger.INSTANCE.logcatMessage(this,
                 () -> String.format("getBacklinksDataForTaskId for taskId %d; topActivity %s",
                         taskInfo.taskId, taskInfo.topActivity));
@@ -284,7 +327,13 @@
         int taskId = taskInfo.taskId;
         mAssistContentRequester.requestAssistContent(taskId, assistContent ->
                 backlinksData.set(getBacklinksDataFromAssistContent(taskInfo, assistContent)));
-        return withTimeout(backlinksData, 5L, TimeUnit.SECONDS, newSingleThreadScheduledExecutor());
+        return withTimeout(backlinksData);
+    }
+
+    /** Returns the same {@link ListenableFuture} but with a 5 {@link TimeUnit#SECONDS} timeout. */
+    private static <V> ListenableFuture<V> withTimeout(ListenableFuture<V> future) {
+        return Futures.withTimeout(future, 5L, TimeUnit.SECONDS,
+                newSingleThreadScheduledExecutor());
     }
 
     /**
@@ -306,7 +355,7 @@
      * @param content the {@link AssistContent} to map into Backlinks {@link ClipData}.
      * @return {@link InternalBacklinksData} that represents the Backlinks data along with app icon.
      */
-    private InternalBacklinksData getBacklinksDataFromAssistContent(RootTaskInfo taskInfo,
+    private InternalBacklinksData getBacklinksDataFromAssistContent(TaskInfo taskInfo,
             @Nullable AssistContent content) {
         DebugLogger.INSTANCE.logcatMessage(this,
                 () -> String.format("getBacklinksDataFromAssistContent taskId %d; topActivity %s",
@@ -365,7 +414,7 @@
         return resolvedComponent.getPackageName().equals(requiredPackageName);
     }
 
-    private String getAppNameOfTask(RootTaskInfo taskInfo) {
+    private String getAppNameOfTask(TaskInfo taskInfo) {
         return taskInfo.topActivityInfo.loadLabel(mPackageManager).toString();
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
index 682f848..254dde4 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/dagger/ScreenshotModule.java
@@ -16,12 +16,17 @@
 
 package com.android.systemui.screenshot.dagger;
 
+import static com.android.systemui.Flags.screenshotUiControllerRefactor;
+
 import android.app.Service;
 import android.view.accessibility.AccessibilityManager;
 
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.screenshot.ImageCapture;
 import com.android.systemui.screenshot.ImageCaptureImpl;
+import com.android.systemui.screenshot.InteractiveScreenshotHandler;
+import com.android.systemui.screenshot.LegacyScreenshotController;
+import com.android.systemui.screenshot.ScreenshotController;
 import com.android.systemui.screenshot.ScreenshotPolicy;
 import com.android.systemui.screenshot.ScreenshotPolicyImpl;
 import com.android.systemui.screenshot.ScreenshotSoundController;
@@ -90,4 +95,15 @@
             AccessibilityManager accessibilityManager) {
         return new ScreenshotViewModel(accessibilityManager);
     }
+
+    @Provides
+    static InteractiveScreenshotHandler.Factory providesScreenshotController(
+            LegacyScreenshotController.Factory legacyScreenshotController,
+            ScreenshotController.Factory screenshotController) {
+        if (screenshotUiControllerRefactor()) {
+            return screenshotController;
+        } else {
+            return legacyScreenshotController;
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotData.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotData.java
index ebac5bf..08c1fca 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotData.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotData.java
@@ -16,8 +16,9 @@
 
 package com.android.systemui.screenshot.scroll;
 
+import android.graphics.Rect;
+
 import com.android.systemui.dagger.SysUISingleton;
-import com.android.systemui.screenshot.ScreenshotController;
 
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -30,9 +31,18 @@
 @SysUISingleton
 public class LongScreenshotData {
     private final AtomicReference<ScrollCaptureController.LongScreenshot> mLongScreenshot;
-    private final AtomicReference<ScreenshotController.TransitionDestination>
+    private final AtomicReference<TransitionDestination>
             mTransitionDestinationCallback;
 
+    public interface TransitionDestination {
+        /**
+         * Allows the long screenshot activity to call back with a destination location (the bounds
+         * on screen of the destination for the transitioning view) and a Runnable to be run once
+         * the transition animation is complete.
+         */
+        void setTransitionDestination(Rect transitionDestination, Runnable onTransitionEnd);
+    }
+
     @Inject
     public LongScreenshotData() {
         mLongScreenshot = new AtomicReference<>();
@@ -63,15 +73,14 @@
     /**
      * Set the holder's TransitionDestination callback.
      */
-    public void setTransitionDestinationCallback(
-            ScreenshotController.TransitionDestination destination) {
+    public void setTransitionDestinationCallback(TransitionDestination destination) {
         mTransitionDestinationCallback.set(destination);
     }
 
     /**
      * Return the current TransitionDestination callback and clear.
      */
-    public ScreenshotController.TransitionDestination takeTransitionDestinationCallback() {
+    public TransitionDestination takeTransitionDestinationCallback() {
         return mTransitionDestinationCallback.getAndSet(null);
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/settings/SystemSettingsRepositoryModule.kt b/packages/SystemUI/src/com/android/systemui/settings/SystemSettingsRepositoryModule.kt
new file mode 100644
index 0000000..02ce74a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/settings/SystemSettingsRepositoryModule.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 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.settings
+
+import android.content.ContentResolver
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.shared.settings.data.repository.SystemSettingsRepository
+import com.android.systemui.shared.settings.data.repository.SystemSettingsRepositoryImpl
+import dagger.Module
+import dagger.Provides
+import kotlinx.coroutines.CoroutineDispatcher
+
+@Module
+object SystemSettingsRepositoryModule {
+    @JvmStatic
+    @Provides
+    @SysUISingleton
+    fun provideSystemSettingsRepository(
+        contentResolver: ContentResolver,
+        @Background backgroundDispatcher: CoroutineDispatcher,
+    ): SystemSettingsRepository =
+        SystemSettingsRepositoryImpl(contentResolver, backgroundDispatcher)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt b/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt
index 25d1cd1..05c50fe 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt
@@ -133,13 +133,6 @@
     private var touchMonitor: TouchMonitor? = null
 
     /**
-     * The width of the area in which a right edge swipe can open the hub, in pixels. Read from
-     * resources when [initView] is called.
-     */
-    // TODO(b/320786721): support RTL layouts
-    private var rightEdgeSwipeRegionWidth: Int = 0
-
-    /**
      * True if we are currently tracking a touch intercepted by the hub, either because the hub is
      * open or being opened.
      */
@@ -265,11 +258,6 @@
 
         communalContainerView = containerView
 
-        rightEdgeSwipeRegionWidth =
-            containerView.resources.getDimensionPixelSize(
-                R.dimen.communal_right_edge_swipe_region_width
-            )
-
         val topEdgeSwipeRegionWidth =
             containerView.resources.getDimensionPixelSize(
                 R.dimen.communal_top_edge_swipe_region_height
@@ -286,7 +274,7 @@
             // Run when the touch handling lifecycle is RESUMED, meaning the hub is visible and not
             // occluded.
             lifecycleRegistry.repeatOnLifecycle(Lifecycle.State.RESUMED) {
-                // Avoid adding exclusion to right/left edges to allow back gestures.
+                // Avoid adding exclusion to end/start edges to allow back gestures.
                 val insets =
                     if (glanceableHubBackGesture()) {
                         containerView.rootWindowInsets.getInsets(WindowInsets.Type.systemGestures())
@@ -294,17 +282,22 @@
                         Insets.NONE
                     }
 
+                val ltr = containerView.layoutDirection == View.LAYOUT_DIRECTION_LTR
+
+                val backGestureInset =
+                    Rect(
+                        if (ltr) 0 else insets.left,
+                        0,
+                        if (ltr) insets.right else containerView.right,
+                        containerView.bottom,
+                    )
+
                 containerView.systemGestureExclusionRects =
                     if (Flags.hubmodeFullscreenVerticalSwipe()) {
                         listOf(
                             // Disable back gestures on the left side of the screen, to avoid
                             // conflicting with scene transitions.
-                            Rect(
-                                0,
-                                0,
-                                insets.right,
-                                containerView.bottom,
-                            )
+                            backGestureInset
                         )
                     } else {
                         listOf(
@@ -318,12 +311,7 @@
                             ),
                             // Disable back gestures on the left side of the screen, to avoid
                             // conflicting with scene transitions.
-                            Rect(
-                                0,
-                                0,
-                                insets.right,
-                                containerView.bottom,
-                            )
+                            backGestureInset
                         )
                     }
             }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
index bc5cf2a..b60c193 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
@@ -41,10 +41,10 @@
 import android.view.View;
 import android.view.ViewGroup;
 import android.view.WindowInsets;
-import android.view.WindowManager;
 import android.view.WindowManager.LayoutParams;
 import android.view.WindowManagerGlobal;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.systemui.Dumpable;
 import com.android.systemui.biometrics.AuthController;
@@ -101,7 +101,7 @@
 
     private final Context mContext;
     private final WindowRootViewComponent.Factory mWindowRootViewComponentFactory;
-    private final WindowManager mWindowManager;
+    private final ViewCaptureAwareWindowManager mWindowManager;
     private final IActivityManager mActivityManager;
     private final DozeParameters mDozeParameters;
     private final KeyguardStateController mKeyguardStateController;
@@ -145,7 +145,7 @@
     public NotificationShadeWindowControllerImpl(
             Context context,
             WindowRootViewComponent.Factory windowRootViewComponentFactory,
-            WindowManager windowManager,
+            ViewCaptureAwareWindowManager viewCaptureAwareWindowManager,
             IActivityManager activityManager,
             DozeParameters dozeParameters,
             StatusBarStateController statusBarStateController,
@@ -165,7 +165,7 @@
             Lazy<CommunalInteractor> communalInteractor) {
         mContext = context;
         mWindowRootViewComponentFactory = windowRootViewComponentFactory;
-        mWindowManager = windowManager;
+        mWindowManager = viewCaptureAwareWindowManager;
         mActivityManager = activityManager;
         mDozeParameters = dozeParameters;
         mKeyguardStateController = keyguardStateController;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt
index aff57bd..e50d64b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt
@@ -50,10 +50,10 @@
 import javax.inject.Inject
 
 /**
- * Coordinates heads up notification (HUN) interactions with the notification pipeline based on
- * the HUN state reported by the [HeadsUpManager]. In this class we only consider one
- * notification, in particular the [HeadsUpManager.getTopEntry], to be HeadsUpping at a
- * time even though other notifications may be queued to heads up next.
+ * Coordinates heads up notification (HUN) interactions with the notification pipeline based on the
+ * HUN state reported by the [HeadsUpManager]. In this class we only consider one notification, in
+ * particular the [HeadsUpManager.getTopEntry], to be HeadsUpping at a time even though other
+ * notifications may be queued to heads up next.
  *
  * The current HUN, but not HUNs that are queued to heads up, will be:
  * - Lifetime extended until it's no longer heads upping.
@@ -64,7 +64,9 @@
  * Note: The inflation callback in [PreparationCoordinator] handles showing HUNs.
  */
 @CoordinatorScope
-class HeadsUpCoordinator @Inject constructor(
+class HeadsUpCoordinator
+@Inject
+constructor(
     private val mLogger: HeadsUpCoordinatorLogger,
     private val mSystemClock: SystemClock,
     private val mHeadsUpManager: HeadsUpManager,
@@ -104,8 +106,8 @@
     }
 
     /**
-     * Once the pipeline starts running, we can look through posted entries and quickly process
-     * any that don't have groups, and thus will never gave a group heads up edge case.
+     * Once the pipeline starts running, we can look through posted entries and quickly process any
+     * that don't have groups, and thus will never gave a group heads up edge case.
      */
     fun onBeforeTransformGroups(list: List<ListEntry>) {
         mNow = mSystemClock.currentTimeMillis()
@@ -128,120 +130,137 @@
      * we know that stability and [NotifPromoter]s have been applied, so we can use the location of
      * notifications in this list to determine what kind of group heads up behavior should happen.
      */
-    fun onBeforeFinalizeFilter(list: List<ListEntry>) = mHeadsUpManager.modifyHuns { hunMutator ->
-        // Nothing to do if there are no other adds/updates
-        if (mPostedEntries.isEmpty()) {
-            return@modifyHuns
-        }
-        // Calculate a bunch of information about the logical group and the locations of group
-        // entries in the nearly-finalized shade list.  These may be used in the per-group loop.
-        val postedEntriesByGroup = mPostedEntries.values.groupBy { it.entry.sbn.groupKey }
-        val logicalMembersByGroup = mNotifPipeline.allNotifs.asSequence()
-            .filter { postedEntriesByGroup.contains(it.sbn.groupKey) }
-            .groupBy { it.sbn.groupKey }
-        val groupLocationsByKey: Map<String, GroupLocation> by lazy { getGroupLocationsByKey(list) }
-        mLogger.logEvaluatingGroups(postedEntriesByGroup.size)
-        // For each group, determine which notification(s) for a group should heads up.
-        postedEntriesByGroup.forEach { (groupKey, postedEntries) ->
-            // get and classify the logical members
-            val logicalMembers = logicalMembersByGroup[groupKey] ?: emptyList()
-            val logicalSummary = logicalMembers.find { it.sbn.notification.isGroupSummary }
+    fun onBeforeFinalizeFilter(list: List<ListEntry>) =
+        mHeadsUpManager.modifyHuns { hunMutator ->
+            // Nothing to do if there are no other adds/updates
+            if (mPostedEntries.isEmpty()) {
+                return@modifyHuns
+            }
+            // Calculate a bunch of information about the logical group and the locations of group
+            // entries in the nearly-finalized shade list.  These may be used in the per-group loop.
+            val postedEntriesByGroup = mPostedEntries.values.groupBy { it.entry.sbn.groupKey }
+            val logicalMembersByGroup =
+                mNotifPipeline.allNotifs
+                    .asSequence()
+                    .filter { postedEntriesByGroup.contains(it.sbn.groupKey) }
+                    .groupBy { it.sbn.groupKey }
+            val groupLocationsByKey: Map<String, GroupLocation> by lazy {
+                getGroupLocationsByKey(list)
+            }
+            mLogger.logEvaluatingGroups(postedEntriesByGroup.size)
+            // For each group, determine which notification(s) for a group should heads up.
+            postedEntriesByGroup.forEach { (groupKey, postedEntries) ->
+                // get and classify the logical members
+                val logicalMembers = logicalMembersByGroup[groupKey] ?: emptyList()
+                val logicalSummary = logicalMembers.find { it.sbn.notification.isGroupSummary }
 
-            // Report the start of this group's evaluation
-            mLogger.logEvaluatingGroup(groupKey, postedEntries.size, logicalMembers.size)
+                // Report the start of this group's evaluation
+                mLogger.logEvaluatingGroup(groupKey, postedEntries.size, logicalMembers.size)
 
-            // If there is no logical summary, then there is no heads up to transfer
-            if (logicalSummary == null) {
-                postedEntries.forEach {
-                    handlePostedEntry(it, hunMutator, scenario = "logical-summary-missing")
+                // If there is no logical summary, then there is no heads up to transfer
+                if (logicalSummary == null) {
+                    postedEntries.forEach {
+                        handlePostedEntry(it, hunMutator, scenario = "logical-summary-missing")
+                    }
+                    return@forEach
                 }
-                return@forEach
-            }
 
-            // If summary isn't wanted to be heads up, then there is no heads up to transfer
-            if (!isGoingToShowHunStrict(logicalSummary)) {
-                postedEntries.forEach {
-                    handlePostedEntry(it, hunMutator, scenario = "logical-summary-not-heads-up")
+                // If summary isn't wanted to be heads up, then there is no heads up to transfer
+                if (!isGoingToShowHunStrict(logicalSummary)) {
+                    postedEntries.forEach {
+                        handlePostedEntry(it, hunMutator, scenario = "logical-summary-not-heads-up")
+                    }
+                    return@forEach
                 }
-                return@forEach
-            }
 
-            // The group is heads up! Overall goals:
-            //  - Maybe transfer its heads up to a child
-            //  - Also let any/all newly heads up children still heads up
-            var childToReceiveParentHeadsUp: NotificationEntry?
-            var targetType = "undefined"
+                // The group is heads up! Overall goals:
+                //  - Maybe transfer its heads up to a child
+                //  - Also let any/all newly heads up children still heads up
+                var childToReceiveParentHeadsUp: NotificationEntry?
+                var targetType = "undefined"
 
-            // If the parent is heads up, always look at the posted notification with the newest
-            // 'when', and if it is isolated with GROUP_ALERT_SUMMARY, then it should receive the
-            // parent's heads up.
-            childToReceiveParentHeadsUp =
-                findHeadsUpOverride(postedEntries, groupLocationsByKey::getLocation)
-            if (childToReceiveParentHeadsUp != null) {
-                targetType = "headsUpOverride"
-            }
-
-            // If the summary is Detached and we have not picked a receiver of the heads up, then we
-            // need to look for the best child to heads up in place of the summary.
-            val isSummaryAttached = groupLocationsByKey.contains(logicalSummary.key)
-            if (!isSummaryAttached && childToReceiveParentHeadsUp == null) {
+                // If the parent is heads up, always look at the posted notification with the newest
+                // 'when', and if it is isolated with GROUP_ALERT_SUMMARY, then it should receive
+                // the
+                // parent's heads up.
                 childToReceiveParentHeadsUp =
-                    findBestTransferChild(logicalMembers, groupLocationsByKey::getLocation)
+                    findHeadsUpOverride(postedEntries, groupLocationsByKey::getLocation)
                 if (childToReceiveParentHeadsUp != null) {
-                    targetType = "bestChild"
+                    targetType = "headsUpOverride"
                 }
-            }
 
-            // If there is no child to receive the parent heads up, then just handle the posted
-            // entries and return.
-            if (childToReceiveParentHeadsUp == null) {
-                postedEntries.forEach {
-                    handlePostedEntry(it, hunMutator, scenario = "no-transfer-target")
+                // If the summary is Detached and we have not picked a receiver of the heads up,
+                // then we
+                // need to look for the best child to heads up in place of the summary.
+                val isSummaryAttached = groupLocationsByKey.contains(logicalSummary.key)
+                if (!isSummaryAttached && childToReceiveParentHeadsUp == null) {
+                    childToReceiveParentHeadsUp =
+                        findBestTransferChild(logicalMembers, groupLocationsByKey::getLocation)
+                    if (childToReceiveParentHeadsUp != null) {
+                        targetType = "bestChild"
+                    }
                 }
-                return@forEach
-            }
 
-            // At this point we just need to initiate the transfer
-            val summaryUpdate = mPostedEntries[logicalSummary.key]
+                // If there is no child to receive the parent heads up, then just handle the posted
+                // entries and return.
+                if (childToReceiveParentHeadsUp == null) {
+                    postedEntries.forEach {
+                        handlePostedEntry(it, hunMutator, scenario = "no-transfer-target")
+                    }
+                    return@forEach
+                }
 
-            // Because we now know for certain that some child is going to heads up for this summary
-            // (as we have found a child to transfer the heads up to), mark the group as having
-            // interrupted. This will allow us to know in the future that the "should heads up"
-            // state of this group has already been handled, just not via the summary entry itself.
-            logicalSummary.setInterruption()
-            mLogger.logSummaryMarkedInterrupted(logicalSummary.key, childToReceiveParentHeadsUp.key)
+                // At this point we just need to initiate the transfer
+                val summaryUpdate = mPostedEntries[logicalSummary.key]
 
-            // If the summary was not attached, then remove the heads up from the detached summary.
-            // Otherwise we can simply ignore its posted update.
-            if (!isSummaryAttached) {
-                val summaryUpdateForRemoval = summaryUpdate?.also {
-                    it.shouldHeadsUpEver = false
-                } ?: PostedEntry(
-                        logicalSummary,
-                        wasAdded = false,
-                        wasUpdated = false,
-                        shouldHeadsUpEver = false,
-                        shouldHeadsUpAgain = false,
-                        isHeadsUpEntry = mHeadsUpManager.isHeadsUpEntry(logicalSummary.key),
-                        isBinding = isEntryBinding(logicalSummary),
+                // Because we now know for certain that some child is going to heads up for this
+                // summary
+                // (as we have found a child to transfer the heads up to), mark the group as having
+                // interrupted. This will allow us to know in the future that the "should heads up"
+                // state of this group has already been handled, just not via the summary entry
+                // itself.
+                logicalSummary.setInterruption()
+                mLogger.logSummaryMarkedInterrupted(
+                    logicalSummary.key,
+                    childToReceiveParentHeadsUp.key
                 )
-                // If we transfer the heads up notification and the summary isn't even attached,
-                // that means we should ensure the summary is no longer a heads up notification,
-                // so we remove it here.
-                handlePostedEntry(
+
+                // If the summary was not attached, then remove the heads up from the detached
+                // summary.
+                // Otherwise we can simply ignore its posted update.
+                if (!isSummaryAttached) {
+                    val summaryUpdateForRemoval =
+                        summaryUpdate?.also { it.shouldHeadsUpEver = false }
+                            ?: PostedEntry(
+                                logicalSummary,
+                                wasAdded = false,
+                                wasUpdated = false,
+                                shouldHeadsUpEver = false,
+                                shouldHeadsUpAgain = false,
+                                isHeadsUpEntry = mHeadsUpManager.isHeadsUpEntry(logicalSummary.key),
+                                isBinding = isEntryBinding(logicalSummary),
+                            )
+                    // If we transfer the heads up notification and the summary isn't even attached,
+                    // that means we should ensure the summary is no longer a heads up notification,
+                    // so we remove it here.
+                    handlePostedEntry(
                         summaryUpdateForRemoval,
                         hunMutator,
-                        scenario = "detached-summary-remove-heads-up")
-            } else if (summaryUpdate != null) {
-                mLogger.logPostedEntryWillNotEvaluate(
+                        scenario = "detached-summary-remove-heads-up"
+                    )
+                } else if (summaryUpdate != null) {
+                    mLogger.logPostedEntryWillNotEvaluate(
                         summaryUpdate,
-                        reason = "attached-summary-transferred")
-            }
+                        reason = "attached-summary-transferred"
+                    )
+                }
 
-            // Handle all posted entries -- if the child receiving the parent's heads up is in the
-            // list, then set its flags to ensure it heads up.
-            var didHeadsUpChildToReceiveParentHeadsUp = false
-            postedEntries.asSequence()
+                // Handle all posted entries -- if the child receiving the parent's heads up is in
+                // the
+                // list, then set its flags to ensure it heads up.
+                var didHeadsUpChildToReceiveParentHeadsUp = false
+                postedEntries
+                    .asSequence()
                     .filter { it.key != logicalSummary.key }
                     .forEach { postedEntry ->
                         if (childToReceiveParentHeadsUp.key == postedEntry.key) {
@@ -249,44 +268,49 @@
                             postedEntry.shouldHeadsUpEver = true
                             postedEntry.shouldHeadsUpAgain = true
                             handlePostedEntry(
-                                    postedEntry,
-                                    hunMutator,
-                                    scenario = "child-heads-up-transfer-target-$targetType")
+                                postedEntry,
+                                hunMutator,
+                                scenario = "child-heads-up-transfer-target-$targetType"
+                            )
                             didHeadsUpChildToReceiveParentHeadsUp = true
                         } else {
                             handlePostedEntry(
-                                    postedEntry,
-                                    hunMutator,
-                                    scenario = "child-heads-up-non-target")
+                                postedEntry,
+                                hunMutator,
+                                scenario = "child-heads-up-non-target"
+                            )
                         }
                     }
 
-            // If the child receiving the heads up notification was not updated on this tick
-            // (which can happen in a standard heads up transfer scenario), then construct an update
-            // so that we can apply it.
-            if (!didHeadsUpChildToReceiveParentHeadsUp) {
-                val posted = PostedEntry(
-                        childToReceiveParentHeadsUp,
-                        wasAdded = false,
-                        wasUpdated = false,
-                        shouldHeadsUpEver = true,
-                        shouldHeadsUpAgain = true,
-                        isHeadsUpEntry =
+                // If the child receiving the heads up notification was not updated on this tick
+                // (which can happen in a standard heads up transfer scenario), then construct an
+                // update
+                // so that we can apply it.
+                if (!didHeadsUpChildToReceiveParentHeadsUp) {
+                    val posted =
+                        PostedEntry(
+                            childToReceiveParentHeadsUp,
+                            wasAdded = false,
+                            wasUpdated = false,
+                            shouldHeadsUpEver = true,
+                            shouldHeadsUpAgain = true,
+                            isHeadsUpEntry =
                                 mHeadsUpManager.isHeadsUpEntry(childToReceiveParentHeadsUp.key),
-                        isBinding = isEntryBinding(childToReceiveParentHeadsUp),
-                )
-                handlePostedEntry(
+                            isBinding = isEntryBinding(childToReceiveParentHeadsUp),
+                        )
+                    handlePostedEntry(
                         posted,
                         hunMutator,
-                        scenario = "non-posted-child-heads-up-transfer-target-$targetType")
+                        scenario = "non-posted-child-heads-up-transfer-target-$targetType"
+                    )
+                }
             }
-        }
-        // After this method runs, all posted entries should have been handled (or skipped).
-        mPostedEntries.clear()
+            // After this method runs, all posted entries should have been handled (or skipped).
+            mPostedEntries.clear()
 
-        // Also take this opportunity to clean up any stale entry update times
-        cleanUpEntryTimes()
-    }
+            // Also take this opportunity to clean up any stale entry update times
+            cleanUpEntryTimes()
+        }
 
     /**
      * Find the posted child with the newest when, and return it if it is isolated and has
@@ -295,34 +319,38 @@
     private fun findHeadsUpOverride(
         postedEntries: List<PostedEntry>,
         locationLookupByKey: (String) -> GroupLocation,
-    ): NotificationEntry? = postedEntries.asSequence()
-        .filter { posted -> !posted.entry.sbn.notification.isGroupSummary }
-        .sortedBy { posted ->
-            -posted.entry.sbn.notification.getWhen()
-        }
-        .firstOrNull()
-        ?.let { posted ->
-            posted.entry.takeIf { entry ->
-                locationLookupByKey(entry.key) == GroupLocation.Isolated &&
+    ): NotificationEntry? =
+        postedEntries
+            .asSequence()
+            .filter { posted -> !posted.entry.sbn.notification.isGroupSummary }
+            .sortedBy { posted -> -posted.entry.sbn.notification.getWhen() }
+            .firstOrNull()
+            ?.let { posted ->
+                posted.entry.takeIf { entry ->
+                    locationLookupByKey(entry.key) == GroupLocation.Isolated &&
                         entry.sbn.notification.groupAlertBehavior == GROUP_ALERT_SUMMARY
+                }
             }
-        }
 
     /**
-     * Of children which are attached, look for the child to receive the notification:
-     * First prefer children which were updated, then looking for the ones with the newest 'when'
+     * Of children which are attached, look for the child to receive the notification: First prefer
+     * children which were updated, then looking for the ones with the newest 'when'
      */
     private fun findBestTransferChild(
         logicalMembers: List<NotificationEntry>,
         locationLookupByKey: (String) -> GroupLocation,
-    ): NotificationEntry? = logicalMembers.asSequence()
-        .filter { !it.sbn.notification.isGroupSummary }
-        .filter { locationLookupByKey(it.key) != GroupLocation.Detached }
-        .sortedWith(compareBy(
-            { !mPostedEntries.contains(it.key) },
-            { -it.sbn.notification.getWhen() },
-        ))
-        .firstOrNull()
+    ): NotificationEntry? =
+        logicalMembers
+            .asSequence()
+            .filter { !it.sbn.notification.isGroupSummary }
+            .filter { locationLookupByKey(it.key) != GroupLocation.Detached }
+            .sortedWith(
+                compareBy(
+                    { !mPostedEntries.contains(it.key) },
+                    { -it.sbn.notification.getWhen() },
+                )
+            )
+            .firstOrNull()
 
     private fun getGroupLocationsByKey(list: List<ListEntry>): Map<String, GroupLocation> =
         mutableMapOf<String, GroupLocation>().also { map ->
@@ -387,197 +415,217 @@
         mHeadsUpViewBinder.bindHeadsUpView(posted.entry, this::onHeadsUpViewBound)
     }
 
-    private val mNotifCollectionListener = object : NotifCollectionListener {
-        /**
-         * Notification was just added and if it should heads up, bind the view and then show it.
-         */
-        override fun onEntryAdded(entry: NotificationEntry) {
-            // First check whether this notification should launch a full screen intent, and
-            // launch it if needed.
-            val fsiDecision =
-                mVisualInterruptionDecisionProvider.makeUnloggedFullScreenIntentDecision(entry)
-            mVisualInterruptionDecisionProvider.logFullScreenIntentDecision(fsiDecision)
-            if (fsiDecision.shouldInterrupt) {
-                mLaunchFullScreenIntentProvider.launchFullScreenIntent(entry)
-            } else if (fsiDecision.wouldInterruptWithoutDnd) {
-                // If DND was the only reason this entry was suppressed, note it for potential
-                // reconsideration on later ranking updates.
-                addForFSIReconsideration(entry, mSystemClock.currentTimeMillis())
-            }
-
-            // makeAndLogHeadsUpDecision includes check for whether this notification should be
-            // filtered
-            val shouldHeadsUpEver =
-                mVisualInterruptionDecisionProvider.makeAndLogHeadsUpDecision(entry).shouldInterrupt
-            mPostedEntries[entry.key] = PostedEntry(
-                entry,
-                wasAdded = true,
-                wasUpdated = false,
-                shouldHeadsUpEver = shouldHeadsUpEver,
-                shouldHeadsUpAgain = true,
-                isHeadsUpEntry = false,
-                isBinding = false,
-            )
-
-            // Record the last updated time for this key
-            setUpdateTime(entry, mSystemClock.currentTimeMillis())
-        }
-
-        /**
-         * Notification could've updated to be heads up or not heads up. Even if it did update to
-         * heads up, if the notification specified that it only wants to heads up once, don't heads
-         * up again.
-         */
-        override fun onEntryUpdated(entry: NotificationEntry) {
-            val shouldHeadsUpEver =
-                mVisualInterruptionDecisionProvider.makeAndLogHeadsUpDecision(entry).shouldInterrupt
-            val shouldHeadsUpAgain = shouldHunAgain(entry)
-            val isHeadsUpEntry = mHeadsUpManager.isHeadsUpEntry(entry.key)
-            val isBinding = isEntryBinding(entry)
-            val posted = mPostedEntries.compute(entry.key) { _, value ->
-                value?.also { update ->
-                    update.wasUpdated = true
-                    update.shouldHeadsUpEver = shouldHeadsUpEver
-                    update.shouldHeadsUpAgain = update.shouldHeadsUpAgain || shouldHeadsUpAgain
-                    update.isHeadsUpEntry = isHeadsUpEntry
-                    update.isBinding = isBinding
-                } ?: PostedEntry(
-                    entry,
-                    wasAdded = false,
-                    wasUpdated = true,
-                    shouldHeadsUpEver = shouldHeadsUpEver,
-                    shouldHeadsUpAgain = shouldHeadsUpAgain,
-                    isHeadsUpEntry = isHeadsUpEntry,
-                    isBinding = isBinding,
-                )
-            }
-            // Handle cancelling heads up here, rather than in the OnBeforeFinalizeFilter, so that
-            // work can be done before the ShadeListBuilder is run. This prevents re-entrant
-            // behavior between this Coordinator, HeadsUpManager, and VisualStabilityManager.
-            if (posted?.shouldHeadsUpEver == false) {
-                if (posted.isHeadsUpEntry) {
-                    // We don't want this to be interrupting anymore, let's remove it
-                    mHeadsUpManager.removeNotification(posted.key, false /*removeImmediately*/)
-                } else if (posted.isBinding) {
-                    // Don't let the bind finish
-                    cancelHeadsUpBind(posted.entry)
+    private val mNotifCollectionListener =
+        object : NotifCollectionListener {
+            /**
+             * Notification was just added and if it should heads up, bind the view and then show
+             * it.
+             */
+            override fun onEntryAdded(entry: NotificationEntry) {
+                // First check whether this notification should launch a full screen intent, and
+                // launch it if needed.
+                val fsiDecision =
+                    mVisualInterruptionDecisionProvider.makeUnloggedFullScreenIntentDecision(entry)
+                mVisualInterruptionDecisionProvider.logFullScreenIntentDecision(fsiDecision)
+                if (fsiDecision.shouldInterrupt) {
+                    mLaunchFullScreenIntentProvider.launchFullScreenIntent(entry)
+                } else if (fsiDecision.wouldInterruptWithoutDnd) {
+                    // If DND was the only reason this entry was suppressed, note it for potential
+                    // reconsideration on later ranking updates.
+                    addForFSIReconsideration(entry, mSystemClock.currentTimeMillis())
                 }
+
+                // makeAndLogHeadsUpDecision includes check for whether this notification should be
+                // filtered
+                val shouldHeadsUpEver =
+                    mVisualInterruptionDecisionProvider
+                        .makeAndLogHeadsUpDecision(entry)
+                        .shouldInterrupt
+                mPostedEntries[entry.key] =
+                    PostedEntry(
+                        entry,
+                        wasAdded = true,
+                        wasUpdated = false,
+                        shouldHeadsUpEver = shouldHeadsUpEver,
+                        shouldHeadsUpAgain = true,
+                        isHeadsUpEntry = false,
+                        isBinding = false,
+                    )
+
+                // Record the last updated time for this key
+                setUpdateTime(entry, mSystemClock.currentTimeMillis())
             }
 
-            // Update last updated time for this entry
-            setUpdateTime(entry, mSystemClock.currentTimeMillis())
-        }
-
-        /**
-         * Stop showing as heads up once removed from the notification collection
-         */
-        override fun onEntryRemoved(entry: NotificationEntry, reason: Int) {
-            mPostedEntries.remove(entry.key)
-            mEntriesUpdateTimes.remove(entry.key)
-            cancelHeadsUpBind(entry)
-            val entryKey = entry.key
-            if (mHeadsUpManager.isHeadsUpEntry(entryKey)) {
-                // TODO: This should probably know the RemoteInputCoordinator's conditions,
-                //  or otherwise reference that coordinator's state, rather than replicate its logic
-                val removeImmediatelyForRemoteInput = (mRemoteInputManager.isSpinning(entryKey) &&
-                        !NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY)
-                mHeadsUpManager.removeNotification(entry.key, removeImmediatelyForRemoteInput)
-            }
-        }
-
-        override fun onEntryCleanUp(entry: NotificationEntry) {
-            mHeadsUpViewBinder.abortBindCallback(entry)
-        }
-
-        /**
-         * Identify notifications whose heads-up state changes when the notification rankings are
-         * updated, and have those changed notifications heads up if necessary.
-         *
-         * This method will occur after any operations in onEntryAdded or onEntryUpdated, so any
-         * handling of ranking changes needs to take into account that we may have just made a
-         * PostedEntry for some of these notifications.
-         */
-        override fun onRankingApplied() {
-            // Because a ranking update may cause some notifications that are no longer (or were
-            // never) in mPostedEntries to need to heads up, we need to check every notification
-            // known to the pipeline.
-            for (entry in mNotifPipeline.allNotifs) {
-                // Only consider entries that are recent enough, since we want to apply a fairly
-                // strict threshold for when an entry should be updated via only ranking and not an
-                // app-provided notification update.
-                if (!isNewEnoughForRankingUpdate(entry)) continue
-
-                // The only entries we consider heads up for here are entries that have never
-                // interrupted and that now say they should heads up or FSI; if they've heads uped in
-                // the past, we don't want to incorrectly heads up a second time if there wasn't an
-                // explicit notification update.
-                if (entry.hasInterrupted()) continue
-
-                // Before potentially allowing heads-up, check for any candidates for a FSI launch.
-                // Any entry that is a candidate meets two criteria:
-                //   - was suppressed from FSI launch only by a DND suppression
-                //   - is within the recency window for reconsideration
-                // If any of these entries are no longer suppressed, launch the FSI now.
-                if (isCandidateForFSIReconsideration(entry)) {
-                    val decision =
-                        mVisualInterruptionDecisionProvider.makeUnloggedFullScreenIntentDecision(
-                            entry
-                        )
-                    if (decision.shouldInterrupt) {
-                        // Log both the launch of the full screen and also that this was via a
-                        // ranking update, and finally revoke candidacy for FSI reconsideration
-                        mLogger.logEntryUpdatedToFullScreen(entry.key, decision.logReason)
-                        mVisualInterruptionDecisionProvider.logFullScreenIntentDecision(decision)
-                        mLaunchFullScreenIntentProvider.launchFullScreenIntent(entry)
-                        mFSIUpdateCandidates.remove(entry.key)
-
-                        // if we launch the FSI then this is no longer a candidate for HUN
-                        continue
-                    } else if (decision.wouldInterruptWithoutDnd) {
-                        // decision has not changed; no need to log
-                    } else {
-                        // some other condition is now blocking FSI; log that and revoke candidacy
-                        // for FSI reconsideration
-                        mLogger.logEntryDisqualifiedFromFullScreen(entry.key, decision.logReason)
-                        mVisualInterruptionDecisionProvider.logFullScreenIntentDecision(decision)
-                        mFSIUpdateCandidates.remove(entry.key)
+            /**
+             * Notification could've updated to be heads up or not heads up. Even if it did update
+             * to heads up, if the notification specified that it only wants to heads up once, don't
+             * heads up again.
+             */
+            override fun onEntryUpdated(entry: NotificationEntry) {
+                val shouldHeadsUpEver =
+                    mVisualInterruptionDecisionProvider
+                        .makeAndLogHeadsUpDecision(entry)
+                        .shouldInterrupt
+                val shouldHeadsUpAgain = shouldHunAgain(entry)
+                val isHeadsUpEntry = mHeadsUpManager.isHeadsUpEntry(entry.key)
+                val isBinding = isEntryBinding(entry)
+                val posted =
+                    mPostedEntries.compute(entry.key) { _, value ->
+                        value?.also { update ->
+                            update.wasUpdated = true
+                            update.shouldHeadsUpEver = shouldHeadsUpEver
+                            update.shouldHeadsUpAgain =
+                                update.shouldHeadsUpAgain || shouldHeadsUpAgain
+                            update.isHeadsUpEntry = isHeadsUpEntry
+                            update.isBinding = isBinding
+                        }
+                            ?: PostedEntry(
+                                entry,
+                                wasAdded = false,
+                                wasUpdated = true,
+                                shouldHeadsUpEver = shouldHeadsUpEver,
+                                shouldHeadsUpAgain = shouldHeadsUpAgain,
+                                isHeadsUpEntry = isHeadsUpEntry,
+                                isBinding = isBinding,
+                            )
+                    }
+                // Handle cancelling heads up here, rather than in the OnBeforeFinalizeFilter, so
+                // that
+                // work can be done before the ShadeListBuilder is run. This prevents re-entrant
+                // behavior between this Coordinator, HeadsUpManager, and VisualStabilityManager.
+                if (posted?.shouldHeadsUpEver == false) {
+                    if (posted.isHeadsUpEntry) {
+                        // We don't want this to be interrupting anymore, let's remove it
+                        mHeadsUpManager.removeNotification(posted.key, false /*removeImmediately*/)
+                    } else if (posted.isBinding) {
+                        // Don't let the bind finish
+                        cancelHeadsUpBind(posted.entry)
                     }
                 }
 
-                // The cases where we should consider this notification to be updated:
-                // - if this entry is not present in PostedEntries, and is now in a shouldHeadsUp
-                //   state
-                // - if it is present in PostedEntries and the previous state of shouldHeadsUp
-                //   differs from the updated one
-                val decision =
-                    mVisualInterruptionDecisionProvider.makeUnloggedHeadsUpDecision(entry)
-                val shouldHeadsUpEver = decision.shouldInterrupt
-                val postedShouldHeadsUpEver = mPostedEntries[entry.key]?.shouldHeadsUpEver ?: false
-                val shouldUpdateEntry = postedShouldHeadsUpEver != shouldHeadsUpEver
+                // Update last updated time for this entry
+                setUpdateTime(entry, mSystemClock.currentTimeMillis())
+            }
 
-                if (shouldUpdateEntry) {
-                    mLogger.logEntryUpdatedByRanking(
-                        entry.key,
-                        shouldHeadsUpEver,
-                        decision.logReason
-                    )
-                    onEntryUpdated(entry)
+            /** Stop showing as heads up once removed from the notification collection */
+            override fun onEntryRemoved(entry: NotificationEntry, reason: Int) {
+                mPostedEntries.remove(entry.key)
+                mEntriesUpdateTimes.remove(entry.key)
+                cancelHeadsUpBind(entry)
+                val entryKey = entry.key
+                if (mHeadsUpManager.isHeadsUpEntry(entryKey)) {
+                    // TODO: This should probably know the RemoteInputCoordinator's conditions,
+                    //  or otherwise reference that coordinator's state, rather than replicate its
+                    // logic
+                    val removeImmediatelyForRemoteInput =
+                        (mRemoteInputManager.isSpinning(entryKey) &&
+                            !NotificationRemoteInputManager.FORCE_REMOTE_INPUT_HISTORY)
+                    mHeadsUpManager.removeNotification(entry.key, removeImmediatelyForRemoteInput)
+                }
+            }
+
+            override fun onEntryCleanUp(entry: NotificationEntry) {
+                mHeadsUpViewBinder.abortBindCallback(entry)
+            }
+
+            /**
+             * Identify notifications whose heads-up state changes when the notification rankings
+             * are updated, and have those changed notifications heads up if necessary.
+             *
+             * This method will occur after any operations in onEntryAdded or onEntryUpdated, so any
+             * handling of ranking changes needs to take into account that we may have just made a
+             * PostedEntry for some of these notifications.
+             */
+            override fun onRankingApplied() {
+                // Because a ranking update may cause some notifications that are no longer (or were
+                // never) in mPostedEntries to need to heads up, we need to check every notification
+                // known to the pipeline.
+                for (entry in mNotifPipeline.allNotifs) {
+                    // Only consider entries that are recent enough, since we want to apply a fairly
+                    // strict threshold for when an entry should be updated via only ranking and not
+                    // an
+                    // app-provided notification update.
+                    if (!isNewEnoughForRankingUpdate(entry)) continue
+
+                    // The only entries we consider heads up for here are entries that have never
+                    // interrupted and that now say they should heads up or FSI; if they've heads
+                    // uped in
+                    // the past, we don't want to incorrectly heads up a second time if there wasn't
+                    // an
+                    // explicit notification update.
+                    if (entry.hasInterrupted()) continue
+
+                    // Before potentially allowing heads-up, check for any candidates for a FSI
+                    // launch.
+                    // Any entry that is a candidate meets two criteria:
+                    //   - was suppressed from FSI launch only by a DND suppression
+                    //   - is within the recency window for reconsideration
+                    // If any of these entries are no longer suppressed, launch the FSI now.
+                    if (isCandidateForFSIReconsideration(entry)) {
+                        val decision =
+                            mVisualInterruptionDecisionProvider
+                                .makeUnloggedFullScreenIntentDecision(entry)
+                        if (decision.shouldInterrupt) {
+                            // Log both the launch of the full screen and also that this was via a
+                            // ranking update, and finally revoke candidacy for FSI reconsideration
+                            mLogger.logEntryUpdatedToFullScreen(entry.key, decision.logReason)
+                            mVisualInterruptionDecisionProvider.logFullScreenIntentDecision(
+                                decision
+                            )
+                            mLaunchFullScreenIntentProvider.launchFullScreenIntent(entry)
+                            mFSIUpdateCandidates.remove(entry.key)
+
+                            // if we launch the FSI then this is no longer a candidate for HUN
+                            continue
+                        } else if (decision.wouldInterruptWithoutDnd) {
+                            // decision has not changed; no need to log
+                        } else {
+                            // some other condition is now blocking FSI; log that and revoke
+                            // candidacy
+                            // for FSI reconsideration
+                            mLogger.logEntryDisqualifiedFromFullScreen(
+                                entry.key,
+                                decision.logReason
+                            )
+                            mVisualInterruptionDecisionProvider.logFullScreenIntentDecision(
+                                decision
+                            )
+                            mFSIUpdateCandidates.remove(entry.key)
+                        }
+                    }
+
+                    // The cases where we should consider this notification to be updated:
+                    // - if this entry is not present in PostedEntries, and is now in a
+                    // shouldHeadsUp
+                    //   state
+                    // - if it is present in PostedEntries and the previous state of shouldHeadsUp
+                    //   differs from the updated one
+                    val decision =
+                        mVisualInterruptionDecisionProvider.makeUnloggedHeadsUpDecision(entry)
+                    val shouldHeadsUpEver = decision.shouldInterrupt
+                    val postedShouldHeadsUpEver =
+                        mPostedEntries[entry.key]?.shouldHeadsUpEver ?: false
+                    val shouldUpdateEntry = postedShouldHeadsUpEver != shouldHeadsUpEver
+
+                    if (shouldUpdateEntry) {
+                        mLogger.logEntryUpdatedByRanking(
+                            entry.key,
+                            shouldHeadsUpEver,
+                            decision.logReason
+                        )
+                        onEntryUpdated(entry)
+                    }
                 }
             }
         }
-    }
 
-    /**
-     * Checks whether an update for a notification warrants an heads up for the user.
-     */
+    /** Checks whether an update for a notification warrants an heads up for the user. */
     private fun shouldHunAgain(entry: NotificationEntry): Boolean {
         return (!entry.hasInterrupted() ||
-                (entry.sbn.notification.flags and Notification.FLAG_ONLY_ALERT_ONCE) == 0)
+            (entry.sbn.notification.flags and Notification.FLAG_ONLY_ALERT_ONCE) == 0)
     }
 
-    /**
-     * Sets the updated time for the given entry to the specified time.
-     */
+    /** Sets the updated time for the given entry to the specified time. */
     @VisibleForTesting
     fun setUpdateTime(entry: NotificationEntry, time: Long) {
         mEntriesUpdateTimes[entry.key] = time
@@ -593,10 +641,10 @@
     }
 
     /**
-     * Checks whether the entry is new enough to be updated via ranking update.
-     * We want to avoid updating an entry too long after it was originally posted/updated when we're
-     * only reacting to a ranking change, as relevant ranking updates are expected to come in
-     * fairly soon after the posting of a notification.
+     * Checks whether the entry is new enough to be updated via ranking update. We want to avoid
+     * updating an entry too long after it was originally posted/updated when we're only reacting to
+     * a ranking change, as relevant ranking updates are expected to come in fairly soon after the
+     * posting of a notification.
      */
     private fun isNewEnoughForRankingUpdate(entry: NotificationEntry): Boolean {
         // If we don't have an update time for this key, default to "too old"
@@ -648,72 +696,92 @@
      * @see HeadsUpManager.setUserActionMayIndirectlyRemove
      * @see HeadsUpManager.canRemoveImmediately
      */
-    private val mActionPressListener = Consumer<NotificationEntry> { entry ->
-        mHeadsUpManager.setUserActionMayIndirectlyRemove(entry)
-        mExecutor.execute { endNotifLifetimeExtensionIfExtended(entry) }
-    }
-
-    private val mLifetimeExtender = object : NotifLifetimeExtender {
-        override fun getName() = TAG
-
-        override fun setCallback(callback: OnEndLifetimeExtensionCallback) {
-            mEndLifetimeExtension = callback
+    private val mActionPressListener =
+        Consumer<NotificationEntry> { entry ->
+            mHeadsUpManager.setUserActionMayIndirectlyRemove(entry)
+            mExecutor.execute { endNotifLifetimeExtensionIfExtended(entry) }
         }
 
-        override fun maybeExtendLifetime(entry: NotificationEntry, reason: Int): Boolean {
-            if (mHeadsUpManager.canRemoveImmediately(entry.key)) {
-                return false
+    private val mLifetimeExtender =
+        object : NotifLifetimeExtender {
+            override fun getName() = TAG
+
+            override fun setCallback(callback: OnEndLifetimeExtensionCallback) {
+                mEndLifetimeExtension = callback
             }
-            if (isSticky(entry)) {
-                val removeAfterMillis = mHeadsUpManager.getEarliestRemovalTime(entry.key)
-                mNotifsExtendingLifetime[entry] = mExecutor.executeDelayed({
-                    mHeadsUpManager.removeNotification(entry.key, /* releaseImmediately */ true)
-                }, removeAfterMillis)
-            } else {
-                mExecutor.execute {
-                    mHeadsUpManager.removeNotification(entry.key, /* releaseImmediately */ false)
+
+            override fun maybeExtendLifetime(entry: NotificationEntry, reason: Int): Boolean {
+                if (mHeadsUpManager.canRemoveImmediately(entry.key)) {
+                    return false
                 }
-                mNotifsExtendingLifetime[entry] = null
+                if (isSticky(entry)) {
+                    val removeAfterMillis = mHeadsUpManager.getEarliestRemovalTime(entry.key)
+                    mNotifsExtendingLifetime[entry] =
+                        mExecutor.executeDelayed(
+                            {
+                                mHeadsUpManager.removeNotification(
+                                    entry.key, /* releaseImmediately */
+                                    true
+                                )
+                            },
+                            removeAfterMillis
+                        )
+                } else {
+                    mExecutor.execute {
+                        mHeadsUpManager.removeNotification(
+                            entry.key, /* releaseImmediately */
+                            false
+                        )
+                    }
+                    mNotifsExtendingLifetime[entry] = null
+                }
+                return true
             }
-            return true
-        }
 
-        override fun cancelLifetimeExtension(entry: NotificationEntry) {
-            mNotifsExtendingLifetime.remove(entry)?.run()
-        }
-    }
-
-    private val mNotifPromoter = object : NotifPromoter(TAG) {
-        override fun shouldPromoteToTopLevel(entry: NotificationEntry): Boolean =
-            isGoingToShowHunNoRetract(entry)
-    }
-
-    val sectioner = object : NotifSectioner("HeadsUp", BUCKET_HEADS_UP) {
-        override fun isInSection(entry: ListEntry): Boolean =
-            // TODO: This check won't notice if a child of the group is going to HUN...
-            isGoingToShowHunNoRetract(entry)
-
-        override fun getComparator(): NotifComparator {
-            return object : NotifComparator("HeadsUp") {
-                override fun compare(o1: ListEntry, o2: ListEntry): Int =
-                    mHeadsUpManager.compare(o1.representativeEntry, o2.representativeEntry)
+            override fun cancelLifetimeExtension(entry: NotificationEntry) {
+                mNotifsExtendingLifetime.remove(entry)?.run()
             }
         }
 
-        override fun getHeaderNodeController(): NodeController? =
-            // TODO: remove SHOW_ALL_SECTIONS, this redundant method, and mIncomingHeaderController
-            if (RankingCoordinator.SHOW_ALL_SECTIONS) mIncomingHeaderController else null
-    }
+    private val mNotifPromoter =
+        object : NotifPromoter(TAG) {
+            override fun shouldPromoteToTopLevel(entry: NotificationEntry): Boolean =
+                isGoingToShowHunNoRetract(entry)
+        }
 
-    private val mOnHeadsUpChangedListener = object : OnHeadsUpChangedListener {
-        override fun onHeadsUpStateChanged(entry: NotificationEntry, isHeadsUp: Boolean) {
-            if (!isHeadsUp) {
-                mNotifPromoter.invalidateList("headsUpEnded: ${entry.logKey}")
-                mHeadsUpViewBinder.unbindHeadsUpView(entry)
-                endNotifLifetimeExtensionIfExtended(entry)
+    val sectioner =
+        object : NotifSectioner("HeadsUp", BUCKET_HEADS_UP) {
+            override fun isInSection(entry: ListEntry): Boolean =
+                // TODO: This check won't notice if a child of the group is going to HUN...
+                isGoingToShowHunNoRetract(entry)
+
+            override fun getComparator(): NotifComparator {
+                return object : NotifComparator("HeadsUp") {
+                    override fun compare(o1: ListEntry, o2: ListEntry): Int =
+                        mHeadsUpManager.compare(o1.representativeEntry, o2.representativeEntry)
+                }
+            }
+
+            override fun getHeaderNodeController(): NodeController? =
+                // TODO: remove SHOW_ALL_SECTIONS, this redundant method, and
+                // mIncomingHeaderController
+                if (RankingCoordinator.SHOW_ALL_SECTIONS) mIncomingHeaderController else null
+        }
+
+    private val mOnHeadsUpChangedListener =
+        object : OnHeadsUpChangedListener {
+            override fun onHeadsUpStateChanged(entry: NotificationEntry, isHeadsUp: Boolean) {
+                if (!isHeadsUp) {
+                    mNotifPromoter.invalidateList("headsUpEnded: ${entry.logKey}")
+                    mHeadsUpViewBinder.unbindHeadsUpView(entry)
+                    endNotifLifetimeExtensionIfExtended(entry)
+                }
+            }
+
+            override fun onHeadsUpAnimatingAwayEnded(entry: NotificationEntry) {
+                mNotifPromoter.invalidateList("headsUpAnimatingAwayEnded: ${entry.logKey}")
             }
         }
-    }
 
     private fun isSticky(entry: NotificationEntry) = mHeadsUpManager.isSticky(entry.key)
 
@@ -726,8 +794,9 @@
      * Whether the notification is already heads up or binding so that it can imminently heads up
      */
     private fun isAttemptingToShowHun(entry: ListEntry) =
-        mHeadsUpManager.isHeadsUpEntry(entry.key) || isEntryBinding(entry)
-                || isHeadsUpAnimatingAway(entry)
+        mHeadsUpManager.isHeadsUpEntry(entry.key) ||
+            isEntryBinding(entry) ||
+            isHeadsUpAnimatingAway(entry)
 
     private fun isHeadsUpAnimatingAway(entry: ListEntry): Boolean {
         if (!GroupHunAnimationFix.isEnabled) return false
@@ -735,19 +804,19 @@
     }
 
     /**
-     * Whether the notification is already heads up/binding per [isAttemptingToShowHun] OR if it
-     * has been updated so that it should heads up this update.  This method is permissive because
-     * it returns `true` even if the update would (in isolation of its group) cause the heads up to
-     * be retracted.  This is important for not retracting transferred group heads ups.
+     * Whether the notification is already heads up/binding per [isAttemptingToShowHun] OR if it has
+     * been updated so that it should heads up this update. This method is permissive because it
+     * returns `true` even if the update would (in isolation of its group) cause the heads up to be
+     * retracted. This is important for not retracting transferred group heads ups.
      */
     private fun isGoingToShowHunNoRetract(entry: ListEntry) =
         mPostedEntries[entry.key]?.calculateShouldBeHeadsUpNoRetract ?: isAttemptingToShowHun(entry)
 
     /**
      * If the notification has been updated, then whether it should HUN in isolation, otherwise
-     * defers to the already heads up/binding state of [isAttemptingToShowHun].  This method is
-     * strict because any update which would revoke the heads up supersedes the current
-     * heads up/binding state.
+     * defers to the already heads up/binding state of [isAttemptingToShowHun]. This method is
+     * strict because any update which would revoke the heads up supersedes the current heads
+     * up/binding state.
      */
     private fun isGoingToShowHunStrict(entry: ListEntry) =
         mPostedEntries[entry.key]?.calculateShouldBeHeadsUpStrict ?: isAttemptingToShowHun(entry)
@@ -779,14 +848,21 @@
         val key = entry.key
         val isHeadsUpAlready: Boolean
             get() = isHeadsUpEntry || isBinding
+
         val calculateShouldBeHeadsUpStrict: Boolean
             get() = shouldHeadsUpEver && (wasAdded || shouldHeadsUpAgain || isHeadsUpAlready)
+
         val calculateShouldBeHeadsUpNoRetract: Boolean
             get() = isHeadsUpAlready || (shouldHeadsUpEver && (wasAdded || shouldHeadsUpAgain))
     }
 }
 
-private enum class GroupLocation { Detached, Isolated, Summary, Child }
+private enum class GroupLocation {
+    Detached,
+    Isolated,
+    Summary,
+    Child
+}
 
 private fun Map<String, GroupLocation>.getLocation(key: String): GroupLocation =
     getOrDefault(key, GroupLocation.Detached)
@@ -804,6 +880,7 @@
 /** Mutates the HeadsUp state of notifications. */
 private interface HunMutator {
     fun updateNotification(key: String, shouldHeadsUpAgain: Boolean)
+
     fun removeNotification(key: String, releaseImmediately: Boolean)
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/NotificationSettingsRepositoryModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/NotificationSettingsRepositoryModule.kt
index a7970c7..af21e75 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/NotificationSettingsRepositoryModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/data/NotificationSettingsRepositoryModule.kt
@@ -19,14 +19,16 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.settings.SecureSettingsRepositoryModule
+import com.android.systemui.settings.SystemSettingsRepositoryModule
 import com.android.systemui.shared.notifications.data.repository.NotificationSettingsRepository
 import com.android.systemui.shared.settings.data.repository.SecureSettingsRepository
+import com.android.systemui.shared.settings.data.repository.SystemSettingsRepository
 import dagger.Module
 import dagger.Provides
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 
-@Module(includes = [SecureSettingsRepositoryModule::class])
+@Module(includes = [SecureSettingsRepositoryModule::class, SystemSettingsRepositoryModule::class])
 object NotificationSettingsRepositoryModule {
     @Provides
     @SysUISingleton
@@ -34,10 +36,12 @@
         @Background backgroundScope: CoroutineScope,
         @Background backgroundDispatcher: CoroutineDispatcher,
         secureSettingsRepository: SecureSettingsRepository,
+        systemSettingsRepository: SystemSettingsRepository,
     ): NotificationSettingsRepository =
         NotificationSettingsRepository(
             backgroundScope,
             backgroundDispatcher,
-            secureSettingsRepository
+            secureSettingsRepository,
+            systemSettingsRepository
         )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/SeenNotificationsInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/SeenNotificationsInteractor.kt
index 90a05ef..2956432 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/SeenNotificationsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/SeenNotificationsInteractor.kt
@@ -109,7 +109,7 @@
             .map {
                 secureSettings.getIntForUser(
                     name = Settings.Secure.LOCK_SCREEN_SHOW_ONLY_UNSEEN_NOTIFICATIONS,
-                    def = 0,
+                    default = 0,
                     userHandle = UserHandle.USER_CURRENT,
                 ) == 1
             }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt
index 9d09595..2a8db56 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/CommonVisualInterruptionSuppressors.kt
@@ -46,6 +46,7 @@
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor
 import com.android.systemui.statusbar.StatusBarState.SHADE
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.interruption.NotificationInterruptStateProviderImpl.MAX_HUN_WHEN_AGE_MS
@@ -57,7 +58,6 @@
 import com.android.systemui.statusbar.policy.HeadsUpManager
 import com.android.systemui.util.NotificationChannels
 import com.android.systemui.util.settings.GlobalSettings
-import com.android.systemui.util.settings.SystemSettings
 import com.android.systemui.util.time.SystemClock
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
@@ -98,8 +98,7 @@
         globalSettings.registerContentObserverSync(
             globalSettings.getUriFor(HEADS_UP_NOTIFICATIONS_ENABLED),
             /* notifyForDescendants = */ true,
-            observer
-        )
+            observer)
 
         // QQQ: Do we need to register for SETTING_HEADS_UP_TICKER? It seems unused.
 
@@ -171,10 +170,7 @@
 
 class PeekOldWhenSuppressor(private val systemClock: SystemClock) :
     VisualInterruptionFilter(
-        types = setOf(PEEK),
-        reason = "has old `when`",
-        uiEventId = HUN_SUPPRESSED_OLD_WHEN
-    ) {
+        types = setOf(PEEK), reason = "has old `when`", uiEventId = HUN_SUPPRESSED_OLD_WHEN) {
     private fun whenAge(entry: NotificationEntry) =
         systemClock.currentTimeMillis() - entry.sbn.notification.getWhen()
 
@@ -201,9 +197,7 @@
 
 class PulseLockscreenVisibilityPrivateSuppressor() :
     VisualInterruptionFilter(
-        types = setOf(PULSE),
-        reason = "hidden by lockscreen visibility override"
-    ) {
+        types = setOf(PULSE), reason = "hidden by lockscreen visibility override") {
     override fun shouldSuppress(entry: NotificationEntry) =
         entry.ranking.lockscreenVisibilityOverride == VISIBILITY_PRIVATE
 }
@@ -215,18 +209,13 @@
 
 class HunGroupAlertBehaviorSuppressor() :
     VisualInterruptionFilter(
-        types = setOf(PEEK, PULSE),
-        reason = "suppressive group alert behavior"
-    ) {
+        types = setOf(PEEK, PULSE), reason = "suppressive group alert behavior") {
     override fun shouldSuppress(entry: NotificationEntry) =
         entry.sbn.let { it.isGroup && it.notification.suppressAlertingDueToGrouping() }
 }
 
 class HunSilentNotificationSuppressor() :
-    VisualInterruptionFilter(
-        types = setOf(PEEK, PULSE),
-        reason = "notification isSilent"
-    ) {
+    VisualInterruptionFilter(types = setOf(PEEK, PULSE), reason = "notification isSilent") {
     override fun shouldSuppress(entry: NotificationEntry) =
         entry.sbn.let { Flags.notificationSilentFlag() && it.notification.isSilent }
 }
@@ -273,7 +262,7 @@
 class AvalancheSuppressor(
     private val avalancheProvider: AvalancheProvider,
     private val systemClock: SystemClock,
-    private val systemSettings: SystemSettings,
+    private val settingsInteractor: NotificationSettingsInteractor,
     private val packageManager: PackageManager,
     private val uiEventLogger: UiEventLogger,
     private val context: Context,
@@ -298,7 +287,7 @@
     // education HUNs.
     private var hasShownOnceForDebug = false
 
-    private fun shouldShowEdu() : Boolean {
+    private fun shouldShowEdu(): Boolean {
         val forceShowOnce = SystemProperties.get(FORCE_SHOW_AVALANCHE_EDU_ONCE, "").equals("1")
         return !hasSeenEdu || (forceShowOnce && !hasShownOnceForDebug)
     }
@@ -361,28 +350,25 @@
         return true
     }
 
-    /**
-     * Show avalanche education HUN from SystemUI.
-     */
+    /** Show avalanche education HUN from SystemUI. */
     private fun showEdu() {
         val res = context.resources
-        val titleStr = res.getString(
-            com.android.systemui.res.R.string.adaptive_notification_edu_hun_title)
-        val textStr = res.getString(
-            com.android.systemui.res.R.string.adaptive_notification_edu_hun_text)
-        val actionStr = res.getString(
-            com.android.systemui.res.R.string.go_to_adaptive_notification_settings)
+        val titleStr =
+            res.getString(com.android.systemui.res.R.string.adaptive_notification_edu_hun_title)
+        val textStr =
+            res.getString(com.android.systemui.res.R.string.adaptive_notification_edu_hun_text)
+        val actionStr =
+            res.getString(com.android.systemui.res.R.string.go_to_adaptive_notification_settings)
 
         val intent = Intent(Settings.ACTION_MANAGE_ADAPTIVE_NOTIFICATIONS)
-        val pendingIntent = PendingIntent.getActivity(
-            context, 0, intent,
-            PendingIntent.FLAG_IMMUTABLE
-        )
+        val pendingIntent =
+            PendingIntent.getActivity(context, 0, intent, PendingIntent.FLAG_IMMUTABLE)
 
         // Replace "System UI" app name with "Android System"
         val bundle = Bundle()
-        bundle.putString(Notification.EXTRA_SUBSTITUTE_APP_NAME,
-                context.getString(com.android.internal.R.string.android_system_label))
+        bundle.putString(
+            Notification.EXTRA_SUBSTITUTE_APP_NAME,
+            context.getString(com.android.internal.R.string.android_system_label))
 
         val builder =
             Notification.Builder(context, NotificationChannels.ALERTS)
@@ -400,14 +386,12 @@
 
         notificationManager.notify(SystemMessage.NOTE_ADAPTIVE_NOTIFICATIONS, builder.build())
         hasSeenEdu = true
-        hasShownOnceForDebug = true;
+        hasShownOnceForDebug = true
     }
 
     private fun calculateState(entry: NotificationEntry): State {
-        if (
-            entry.ranking.isConversation &&
-                entry.sbn.notification.getWhen() > avalancheProvider.startTime
-        ) {
+        if (entry.ranking.isConversation &&
+            entry.sbn.notification.getWhen() > avalancheProvider.startTime) {
             uiEventLogger.log(AvalancheEvent.AVALANCHE_SUPPRESSOR_HUN_ALLOWED_NEW_CONVERSATION)
             return State.ALLOW_CONVERSATION_AFTER_AVALANCHE
         }
@@ -440,10 +424,8 @@
             uiEventLogger.log(AvalancheEvent.AVALANCHE_SUPPRESSOR_HUN_ALLOWED_COLORIZED)
             return State.ALLOW_COLORIZED
         }
-        if (
-            packageManager.checkPermission(RECEIVE_EMERGENCY_BROADCAST, entry.sbn.packageName) ==
-                PERMISSION_GRANTED
-        ) {
+        if (packageManager.checkPermission(RECEIVE_EMERGENCY_BROADCAST, entry.sbn.packageName) ==
+            PERMISSION_GRANTED) {
             uiEventLogger.log(AvalancheEvent.AVALANCHE_SUPPRESSOR_HUN_ALLOWED_EMERGENCY)
             return State.ALLOW_EMERGENCY
         }
@@ -452,7 +434,6 @@
     }
 
     private fun isCooldownEnabled(): Boolean {
-        return systemSettings.getInt(Settings.System.NOTIFICATION_COOLDOWN_ENABLED, /* def */ 1) ==
-            1
+        return settingsInteractor.isCooldownEnabled.value
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt
index 1c476ce..c0d27cb 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImpl.kt
@@ -29,6 +29,7 @@
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.interruption.VisualInterruptionDecisionProvider.Decision
 import com.android.systemui.statusbar.notification.interruption.VisualInterruptionDecisionProvider.FullScreenIntentDecision
@@ -72,7 +73,8 @@
     private val packageManager: PackageManager,
     private val bubbles: Optional<Bubbles>,
     private val context: Context,
-    private val notificationManager: NotificationManager
+    private val notificationManager: NotificationManager,
+    private val settingsInteractor: NotificationSettingsInteractor
 ) : VisualInterruptionDecisionProvider {
 
     init {
@@ -183,8 +185,8 @@
 
         if (NotificationAvalancheSuppression.isEnabled) {
             addFilter(
-                AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
-                        uiEventLogger, context, notificationManager)
+                AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor,
+                    packageManager, uiEventLogger, context, notificationManager)
             )
             avalancheProvider.register()
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
index 35afda7..9f634be 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
@@ -19,6 +19,7 @@
 import android.content.Context;
 import android.util.AttributeSet;
 import android.util.Log;
+import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
 
@@ -89,10 +90,59 @@
         inflater.inflate(R.layout.status_bar_notification_row, parent, listenerExecutor, this);
     }
 
+    /**
+     * Inflates a new notificationView synchronously.
+     * This method is only for testing-purpose.
+     */
+    @VisibleForTesting
+    public ExpandableNotificationRow inflateSynchronously(@NonNull Context context,
+            @Nullable ViewGroup parent, @NonNull NotificationEntry entry) {
+        final LayoutInflater inflater = new BasicRowInflater(context);
+        inflater.setFactory2(makeRowInflater(entry));
+        final ExpandableNotificationRow inflate = (ExpandableNotificationRow) inflater.inflate(
+                R.layout.status_bar_notification_row,
+                parent /* root */,
+                false /* attachToRoot */);
+        return inflate;
+    }
+
     private RowAsyncLayoutInflater makeRowInflater(NotificationEntry entry) {
         return new RowAsyncLayoutInflater(entry, mSystemClock, mLogger);
     }
 
+    /**
+     * A {@link LayoutInflater} that is copy of BasicLayoutInflater.
+     */
+    private static class BasicRowInflater extends LayoutInflater {
+        private static final String[] sClassPrefixList =
+                {"android.widget.", "android.webkit.", "android.app."};
+        BasicRowInflater(Context context) {
+            super(context);
+        }
+
+        @Override
+        public LayoutInflater cloneInContext(Context newContext) {
+            return new BasicRowInflater(newContext);
+        }
+
+        @Override
+        protected View onCreateView(String name, AttributeSet attrs) throws ClassNotFoundException {
+            for (String prefix : sClassPrefixList) {
+                try {
+                    View view = createView(name, prefix, attrs);
+                    if (view != null) {
+                        return view;
+                    }
+                } catch (ClassNotFoundException e) {
+                    // In this case we want to let the base class take a crack
+                    // at it.
+                }
+            }
+
+            return super.onCreateView(name, attrs);
+        }
+    }
+
     @VisibleForTesting
     public static class RowAsyncLayoutInflater implements AsyncLayoutFactory {
         private final NotificationEntry mEntry;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/GroupHunAnimationFix.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/GroupHunAnimationFix.kt
index 5867612..3b30c86 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/GroupHunAnimationFix.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/GroupHunAnimationFix.kt
@@ -29,7 +29,7 @@
     val token: FlagToken
         get() = FlagToken(FLAG_NAME, isEnabled)
 
-    /** Are sections sorted by time? */
+    /** Return whether the fix is enabled */
     @JvmStatic
     inline val isEnabled
         get() = Flags.notificationGroupHunRemovalAnimationFix()
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index a072ea6..fb1c525 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -124,6 +124,7 @@
 import com.android.systemui.statusbar.notification.row.NotificationGuts;
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
 import com.android.systemui.statusbar.notification.row.NotificationSnooze;
+import com.android.systemui.statusbar.notification.shared.GroupHunAnimationFix;
 import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.notification.stack.ui.viewbinder.NotificationListViewBinder;
 import com.android.systemui.statusbar.phone.HeadsUpAppearanceController;
@@ -1987,6 +1988,10 @@
                 NotificationEntry entry = row.getEntry();
                 mHeadsUpAppearanceController.updateHeader(entry);
                 mHeadsUpAppearanceController.updateHeadsUpAndPulsingRoundness(entry);
+                if (GroupHunAnimationFix.isEnabled() && !animatingAway) {
+                    // invalidate list to make sure the row is sorted to the correct section
+                    mHeadsUpManager.onEntryAnimatingAwayEnded(entry);
+                }
             });
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
index 1a7bc16..e8a7840 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
@@ -91,6 +91,12 @@
     val expandFraction: Flow<Float> = shadeInteractor.anyExpansion.dumpValue("expandFraction")
 
     /**
+     * The amount [0-1] that quick settings has been opened. At 0, the shade may be open or closed;
+     * at 1, the quick settings are open.
+     */
+    val shadeToQsFraction: Flow<Float> = shadeInteractor.qsExpansion.dumpValue("shadeToQsFraction")
+
+    /**
      * The amount in px that the notification stack should scroll due to internal expansion. This
      * should only happen when a notification expansion hits the bottom of the screen, so it is
      * necessary to scroll up to keep expanding the notification.
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfaces.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfaces.java
index a0d4ca2..ae31151 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfaces.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfaces.java
@@ -261,8 +261,6 @@
 
     boolean isScreenFullyOff();
 
-    boolean isCameraAllowedByAdmin();
-
     boolean isGoingToSleep();
 
     void notifyBiometricAuthModeChanged();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacks.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacks.java
index c8a4450..5209d0f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacks.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacks.java
@@ -49,6 +49,7 @@
 import com.android.systemui.emergency.EmergencyGesture;
 import com.android.systemui.emergency.EmergencyGestureModule.EmergencyGestureIntentFactory;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.qs.QSHost;
 import com.android.systemui.qs.QSPanelController;
@@ -109,6 +110,7 @@
     private final Lazy<CameraLauncher> mCameraLauncherLazy;
     private final QuickSettingsController mQsController;
     private final QSHost mQSHost;
+    private final KeyguardInteractor mKeyguardInteractor;
     private static final VibrationAttributes HARDWARE_FEEDBACK_VIBRATION_ATTRIBUTES =
             VibrationAttributes.createForUsage(VibrationAttributes.USAGE_HARDWARE_FEEDBACK);
 
@@ -148,6 +150,7 @@
             UserTracker userTracker,
             QSHost qsHost,
             ActivityStarter activityStarter,
+            KeyguardInteractor keyguardInteractor,
             EmergencyGestureIntentFactory emergencyGestureIntentFactory) {
         mCentralSurfaces = centralSurfaces;
         mQsController = quickSettingsController;
@@ -176,7 +179,7 @@
         mCameraLauncherLazy = cameraLauncherLazy;
         mUserTracker = userTracker;
         mQSHost = qsHost;
-
+        mKeyguardInteractor = keyguardInteractor;
         mVibrateOnOpening = resources.getBoolean(R.bool.config_vibrateOnIconAnimation);
         mCameraLaunchGestureVibrationEffect = getCameraGestureVibrationEffect(
                 mVibratorOptional, resources);
@@ -351,6 +354,8 @@
             }
             return;
         }
+        mKeyguardInteractor.onCameraLaunchDetected(source);
+
         if (!mCentralSurfaces.isDeviceInteractive()) {
             mPowerManager.wakeUp(SystemClock.uptimeMillis(), PowerManager.WAKE_REASON_CAMERA_LAUNCH,
                     "com.android.systemui:CAMERA_GESTURE");
@@ -383,6 +388,7 @@
                 if (mStatusBarKeyguardViewManager.isBouncerShowing()) {
                     mStatusBarKeyguardViewManager.reset(true /* hide */);
                 }
+                mCentralSurfaces.startLaunchTransitionTimeout();
                 mCameraLauncherLazy.get().launchCamera(source,
                         mPanelExpansionInteractor.isFullyCollapsed());
                 mCentralSurfaces.updateScrimController();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesEmptyImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesEmptyImpl.kt
index 88d3e07..d4f2a93 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesEmptyImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesEmptyImpl.kt
@@ -34,72 +34,127 @@
  */
 abstract class CentralSurfacesEmptyImpl : CentralSurfaces {
     override val lifecycle = LifecycleRegistry(this)
+
     override fun updateIsKeyguard() = false
+
     override fun updateIsKeyguard(forceStateChange: Boolean) = false
+
     override fun getKeyguardMessageArea(): AuthKeyguardMessageArea? = null
+
     override fun isLaunchingActivityOverLockscreen() = false
+
     override fun isDismissingShadeForActivityLaunch() = false
+
     override fun onKeyguardViewManagerStatesUpdated() {}
+
     override fun getCommandQueuePanelsEnabled() = false
+
     override fun showWirelessChargingAnimation(batteryLevel: Int) {}
+
     override fun checkBarModes() {}
+
     override fun updateBubblesVisibility() {}
+
     override fun setInteracting(barWindow: Int, interacting: Boolean) {}
+
     override fun getDisplayWidth() = 0f
+
     override fun getDisplayHeight() = 0f
+
     override fun showKeyguard() {}
+
     override fun hideKeyguard() = false
+
     override fun showKeyguardImpl() {}
+
     override fun fadeKeyguardAfterLaunchTransition(
         beforeFading: Runnable?,
         endRunnable: Runnable?,
         cancelRunnable: Runnable?,
     ) {}
+
     override fun startLaunchTransitionTimeout() {}
+
     override fun hideKeyguardImpl(forceStateChange: Boolean) = false
+
     override fun keyguardGoingAway() {}
+
     override fun setKeyguardFadingAway(startTime: Long, delay: Long, fadeoutDuration: Long) {}
+
     override fun finishKeyguardFadingAway() {}
+
     override fun userActivity() {}
+
     override fun endAffordanceLaunch() {}
+
     override fun shouldKeyguardHideImmediately() = false
+
     override fun showBouncerWithDimissAndCancelIfKeyguard(
         performAction: OnDismissAction?,
         cancelAction: Runnable?,
     ) {}
+
     override fun getNavigationBarView(): NavigationBarView? = null
+
     override fun setBouncerShowing(bouncerShowing: Boolean) {}
+
     override fun isScreenFullyOff() = false
-    override fun isCameraAllowedByAdmin() = false
+
     override fun isGoingToSleep() = false
+
     override fun notifyBiometricAuthModeChanged() {}
+
     override fun setTransitionToFullShadeProgress(transitionToFullShadeProgress: Float) {}
+
     override fun setPrimaryBouncerHiddenFraction(expansion: Float) {}
+
     override fun updateScrimController() {}
+
     override fun shouldIgnoreTouch() = false
+
     override fun isDeviceInteractive() = false
+
     override fun handleExternalShadeWindowTouch(event: MotionEvent?) {}
+
     override fun handleCommunalHubTouch(event: MotionEvent?) {}
+
     override fun awakenDreams() {}
+
     override fun isBouncerShowing() = false
+
     override fun isBouncerShowingScrimmed() = false
+
     override fun updateNotificationPanelTouchState() {}
+
     override fun getRotation() = 0
+
     override fun setBarStateForTest(state: Int) {}
+
     override fun acquireGestureWakeLock(time: Long) {}
+
     override fun resendMessage(msg: Int) {}
+
     override fun resendMessage(msg: Any?) {}
+
     override fun setLastCameraLaunchSource(source: Int) {}
+
     override fun setLaunchCameraOnFinishedGoingToSleep(launch: Boolean) {}
+
     override fun setLaunchCameraOnFinishedWaking(launch: Boolean) {}
+
     override fun setLaunchEmergencyActionOnFinishedGoingToSleep(launch: Boolean) {}
+
     override fun setLaunchEmergencyActionOnFinishedWaking(launch: Boolean) {}
+
     override fun getQSPanelController(): QSPanelController? = null
+
     override fun getDisplayDensity() = 0f
+
     override fun setIsLaunchingActivityOverLockscreen(
         isLaunchingActivityOverLockscreen: Boolean,
         dismissShade: Boolean,
     ) {}
+
     override fun getAnimatorControllerFromNotification(
         associatedView: ExpandableNotificationRow?,
     ): ActivityTransitionAnimator.Controller? = null
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
index 462ae7a..461a38d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
@@ -47,7 +47,6 @@
 import android.app.TaskInfo;
 import android.app.UiModeManager;
 import android.app.WallpaperManager;
-import android.app.admin.DevicePolicyManager;
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
@@ -883,8 +882,6 @@
 
         // start old BaseStatusBar.start().
         mWindowManagerService = WindowManagerGlobal.getWindowManagerService();
-        mDevicePolicyManager = (DevicePolicyManager) mContext.getSystemService(
-                Context.DEVICE_POLICY_SERVICE);
 
         mAccessibilityManager = (AccessibilityManager)
                 mContext.getSystemService(Context.ACCESSIBILITY_SERVICE);
@@ -2627,6 +2624,7 @@
                 mStackScrollerController.updateSensitivenessForOccludedWakeup();
             }
             if (mLaunchCameraWhenFinishedWaking) {
+                startLaunchTransitionTimeout();
                 mCameraLauncherLazy.get().launchCamera(mLastCameraLaunchSource,
                         mShadeSurface.isFullyCollapsed());
                 mLaunchCameraWhenFinishedWaking = false;
@@ -2701,21 +2699,6 @@
     }
 
     @Override
-    public boolean isCameraAllowedByAdmin() {
-        if (mDevicePolicyManager.getCameraDisabled(null,
-                mLockscreenUserManager.getCurrentUserId())) {
-            return false;
-        } else if (mKeyguardStateController.isShowing()
-                && mStatusBarKeyguardViewManager.isSecure()) {
-            // Check if the admin has disabled the camera specifically for the keyguard
-            return (mDevicePolicyManager.getKeyguardDisabledFeatures(null,
-                    mLockscreenUserManager.getCurrentUserId())
-                    & DevicePolicyManager.KEYGUARD_DISABLE_SECURE_CAMERA) == 0;
-        }
-        return true;
-    }
-
-    @Override
     public boolean isGoingToSleep() {
         return mWakefulnessLifecycle.getWakefulness()
                 == WakefulnessLifecycle.WAKEFULNESS_GOING_TO_SLEEP;
@@ -2864,7 +2847,6 @@
 
     protected boolean mDeviceInteractive;
 
-    protected DevicePolicyManager mDevicePolicyManager;
     private final PowerManager mPowerManager;
     protected StatusBarKeyguardViewManager mStatusBarKeyguardViewManager;
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
index e08dbb9..25d9cc7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
@@ -117,7 +117,7 @@
 
         @Override
         public HeadsUpEntryPhone acquire() {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
+            NotificationThrottleHun.assertInLegacyMode();
             if (!mPoolObjects.isEmpty()) {
                 return mPoolObjects.pop();
             }
@@ -126,7 +126,7 @@
 
         @Override
         public boolean release(@NonNull HeadsUpEntryPhone instance) {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
+            NotificationThrottleHun.assertInLegacyMode();
             mPoolObjects.push(instance);
             return true;
         }
@@ -389,10 +389,13 @@
     //  OnReorderingAllowedListener:
 
     private final OnReorderingAllowedListener mOnReorderingAllowedListener = () -> {
-        mAnimationStateHandler.setHeadsUpGoingAwayAnimationsAllowed(false);
         if (NotificationThrottleHun.isEnabled()) {
             mAvalancheController.setEnableAtRuntime(true);
+            if (mEntriesToRemoveWhenReorderingAllowed.isEmpty()) {
+                return;
+            }
         }
+        mAnimationStateHandler.setHeadsUpGoingAwayAnimationsAllowed(false);
         for (NotificationEntry entry : mEntriesToRemoveWhenReorderingAllowed) {
             if (isHeadsUpEntry(entry.getKey())) {
                 // Maybe the heads-up was removed already
@@ -425,7 +428,7 @@
     @NonNull
     @Override
     protected HeadsUpEntry createHeadsUpEntry(NotificationEntry entry) {
-        if (NotificationsHeadsUpRefactor.isEnabled()) {
+        if (NotificationThrottleHun.isEnabled()) {
             return new HeadsUpEntryPhone(entry);
         } else {
             HeadsUpEntryPhone headsUpEntry = mEntryPool.acquire();
@@ -451,7 +454,7 @@
     @Override
     protected void onEntryRemoved(HeadsUpEntry headsUpEntry) {
         super.onEntryRemoved(headsUpEntry);
-        if (!NotificationsHeadsUpRefactor.isEnabled()) {
+        if (!NotificationThrottleHun.isEnabled()) {
             mEntryPool.release((HeadsUpEntryPhone) headsUpEntry);
         }
         updateTopHeadsUpFlow();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
index a963826..5ba5c06 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
@@ -54,6 +54,7 @@
                 // allowed again.
                 logDroppedHunsInBackground(getWaitingKeys().size)
                 clearNext()
+                headsUpEntryShowing = null
             }
             if (field != value) {
                 field = value
@@ -222,15 +223,19 @@
 
     /**
      * Returns duration based on
-     * 1) Whether HeadsUpEntry is the last one tracked byAvalancheController
+     * 1) Whether HeadsUpEntry is the last one tracked by AvalancheController
      * 2) The priority of the top HUN in the next batch Used by
      *    BaseHeadsUpManager.HeadsUpEntry.calculateFinishTime to shorten display duration.
      */
-    fun getDurationMs(entry: HeadsUpEntry, autoDismissMs: Int): Int {
+    fun getDurationMs(entry: HeadsUpEntry?, autoDismissMs: Int): Int {
         if (!isEnabled()) {
             // Use default duration, like we did before AvalancheController existed
             return autoDismissMs
         }
+        if (entry == null) {
+            // This should never happen
+            return autoDismissMs
+        }
         val showingList: MutableList<HeadsUpEntry> = mutableListOf()
         if (headsUpEntryShowing != null) {
             showingList.add(headsUpEntryShowing!!)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java
index dcd9cae..3786958 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BaseHeadsUpManager.java
@@ -171,7 +171,6 @@
         mLogger.logShowNotificationRequest(entry);
 
         Runnable runnable = () -> {
-            // TODO(b/315362456) log outside runnable too
             mLogger.logShowNotification(entry);
 
             // Add new entry and begin managing it
@@ -244,8 +243,10 @@
             return;
         }
         // TODO(b/328390331) move accessibility events to the view layer
-        headsUpEntry.mEntry.sendAccessibilityEvent(AccessibilityEvent.TYPE_WINDOW_CONTENT_CHANGED);
-
+        if (headsUpEntry.mEntry != null) {
+            headsUpEntry.mEntry.sendAccessibilityEvent(
+                    AccessibilityEvent.TYPE_WINDOW_CONTENT_CHANGED);
+        }
         if (shouldHeadsUpAgain) {
             headsUpEntry.updateEntry(true /* updatePostTime */, "updateNotification");
             if (headsUpEntry != null) {
@@ -334,6 +335,9 @@
     }
 
     protected boolean shouldHeadsUpBecomePinned(@NonNull NotificationEntry entry) {
+        if (entry == null) {
+            return false;
+        }
         final HeadsUpEntry headsUpEntry = getHeadsUpEntry(entry.getKey());
         if (headsUpEntry == null) {
             // This should not happen since shouldHeadsUpBecomePinned is always called after adding
@@ -344,6 +348,15 @@
     }
 
     protected boolean hasFullScreenIntent(@NonNull NotificationEntry entry) {
+        if (entry == null) {
+            return false;
+        }
+        if (entry.getSbn() == null) {
+            return false;
+        }
+        if (entry.getSbn().getNotification() == null) {
+            return false;
+        }
         return entry.getSbn().getNotification().fullScreenIntent != null;
     }
 
@@ -426,7 +439,7 @@
             onEntryRemoved(finalHeadsUpEntry);
             // TODO(b/328390331) move accessibility events to the view layer
             entry.sendAccessibilityEvent(AccessibilityEvent.TYPE_WINDOW_CONTENT_CHANGED);
-            if (NotificationsHeadsUpRefactor.isEnabled()) {
+            if (NotificationThrottleHun.isEnabled()) {
                 finalHeadsUpEntry.cancelAutoRemovalCallbacks("removeEntry");
             } else {
                 finalHeadsUpEntry.reset();
@@ -451,6 +464,15 @@
     }
 
     /**
+     * Called to notify the listeners that the HUN animating away animation has ended.
+     */
+    public void onEntryAnimatingAwayEnded(@NonNull NotificationEntry entry) {
+        for (OnHeadsUpChangedListener listener : mListeners) {
+            listener.onHeadsUpAnimatingAwayEnded(entry);
+        }
+    }
+
+    /**
      * Manager-specific logic, that should occur, when the entry is updated, and its posted time has
      * changed.
      *
@@ -499,6 +521,9 @@
         keySet.addAll(mAvalancheController.getWaitingKeys());
         for (String key : keySet) {
             HeadsUpEntry entry = getHeadsUpEntry(key);
+            if (entry.mEntry == null) {
+                continue;
+            }
             String packageName = entry.mEntry.getSbn().getPackageName();
             String snoozeKey = snoozeKey(packageName, mUser);
             mLogger.logPackageSnoozed(snoozeKey);
@@ -566,7 +591,7 @@
         pw.print("  now="); pw.println(mSystemClock.elapsedRealtime());
         pw.print("  mUser="); pw.println(mUser);
         for (HeadsUpEntry entry: mHeadsUpEntryMap.values()) {
-            pw.print("  HeadsUpEntry="); pw.println(entry.mEntry);
+            pw.println(entry.mEntry == null ? "null" : entry.mEntry);
         }
         int n = mSnoozedPackages.size();
         pw.println("  snoozed packages: " + n);
@@ -586,7 +611,7 @@
     private boolean hasPinnedNotificationInternal() {
         for (String key : mHeadsUpEntryMap.keySet()) {
             HeadsUpEntry entry = getHeadsUpEntry(key);
-            if (entry.mEntry.isRowPinned()) {
+            if (entry.mEntry != null && entry.mEntry.isRowPinned()) {
                 return true;
             }
         }
@@ -611,7 +636,7 @@
                 // when the user unpinned all of HUNs by moving one HUN, all of HUNs should not stay
                 // on the screen.
                 if (userUnPinned && headsUpEntry.mEntry != null) {
-                    if (headsUpEntry.mEntry.mustStayOnScreen()) {
+                    if (headsUpEntry.mEntry != null && headsUpEntry.mEntry.mustStayOnScreen()) {
                         headsUpEntry.mEntry.setHeadsUpIsVisible();
                     }
                 }
@@ -687,7 +712,7 @@
             return true;
         }
         return headsUpEntry == null || headsUpEntry.wasShownLongEnough()
-                || headsUpEntry.mEntry.isRowDismissed();
+                || (headsUpEntry.mEntry != null && headsUpEntry.mEntry.isRowDismissed());
     }
 
     /**
@@ -743,7 +768,7 @@
         @Nullable private Runnable mCancelRemoveRunnable;
 
         public HeadsUpEntry() {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
+            NotificationThrottleHun.assertInLegacyMode();
         }
 
         public HeadsUpEntry(NotificationEntry entry) {
@@ -754,7 +779,7 @@
 
         /** Attach a NotificationEntry. */
         public void setEntry(@NonNull final NotificationEntry entry) {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
+            NotificationThrottleHun.assertInLegacyMode();
             setEntry(entry, createRemoveRunnable(entry));
         }
 
@@ -866,6 +891,14 @@
         }
 
         public int compareNonTimeFields(HeadsUpEntry headsUpEntry) {
+            if (mEntry == null && headsUpEntry.mEntry == null) {
+                return 0;
+            } else if (headsUpEntry.mEntry == null) {
+                return -1;
+            } else if (mEntry == null) {
+                return 1;
+            }
+
             boolean selfFullscreen = hasFullScreenIntent(mEntry);
             boolean otherFullscreen = hasFullScreenIntent(headsUpEntry.mEntry);
             if (selfFullscreen && !otherFullscreen) {
@@ -892,6 +925,13 @@
         }
 
         public int compareTo(@NonNull HeadsUpEntry headsUpEntry) {
+            if (mEntry == null && headsUpEntry.mEntry == null) {
+                return 0;
+            } else if (headsUpEntry.mEntry == null) {
+                return -1;
+            } else if (mEntry == null) {
+                return 1;
+            }
             boolean isPinned = mEntry.isRowPinned();
             boolean otherPinned = headsUpEntry.mEntry.isRowPinned();
             if (isPinned && !otherPinned) {
@@ -936,7 +976,7 @@
         }
 
         public void reset() {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
+            NotificationThrottleHun.assertInLegacyMode();
             cancelAutoRemovalCallbacks("reset()");
             mEntry = null;
             mRemoveRunnable = null;
@@ -956,7 +996,7 @@
                     mLogger.logAutoRemoveCanceled(mEntry, reason);
                 }
             };
-            if (isHeadsUpEntry(this.mEntry.getKey())) {
+            if (mEntry != null && isHeadsUpEntry(mEntry.getKey())) {
                 mAvalancheController.update(this, runnable, reason + " cancelAutoRemovalCallbacks");
             } else {
                 // Just removed
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManager.kt
index 28a2a1f..fcf77d5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/HeadsUpManager.kt
@@ -42,6 +42,7 @@
      *   should be ranked higher and 0 if they are equal.
      */
     fun compare(a: NotificationEntry?, b: NotificationEntry?): Int
+
     /**
      * Extends the lifetime of the currently showing pulsing notification so that the pulse lasts
      * longer.
@@ -184,6 +185,8 @@
     fun unpinAll(userUnPinned: Boolean)
 
     fun updateNotification(key: String, shouldHeadsUpAgain: Boolean)
+
+    fun onEntryAnimatingAwayEnded(entry: NotificationEntry)
 }
 
 /** Sets the animation state of the HeadsUpManager. */
@@ -204,41 +207,77 @@
 /* No op impl of HeadsUpManager. */
 class HeadsUpManagerEmptyImpl @Inject constructor() : HeadsUpManager {
     override val allEntries = Stream.empty<NotificationEntry>()
+
     override fun addHeadsUpPhoneListener(listener: OnHeadsUpPhoneListenerChange) {}
+
     override fun addListener(listener: OnHeadsUpChangedListener) {}
+
     override fun addSwipedOutNotification(key: String) {}
+
     override fun canRemoveImmediately(key: String) = false
+
     override fun compare(a: NotificationEntry?, b: NotificationEntry?) = 0
+
     override fun dump(pw: PrintWriter, args: Array<out String>) {}
+
     override fun extendHeadsUp() {}
+
     override fun getEarliestRemovalTime(key: String?) = 0L
+
     override fun getTouchableRegion(): Region? = null
+
     override fun getTopEntry() = null
+
     override fun hasPinnedHeadsUp() = false
+
     override fun isHeadsUpEntry(key: String) = false
+
     override fun isHeadsUpAnimatingAwayValue() = false
+
     override fun isSnoozed(packageName: String) = false
+
     override fun isSticky(key: String?) = false
+
     override fun isTrackingHeadsUp() = false
+
     override fun onExpandingFinished() {}
+
     override fun releaseAllImmediately() {}
+
     override fun removeListener(listener: OnHeadsUpChangedListener) {}
+
     override fun removeNotification(key: String, releaseImmediately: Boolean) = false
+
     override fun removeNotification(key: String, releaseImmediately: Boolean, animate: Boolean) =
         false
+
     override fun setAnimationStateHandler(handler: AnimationStateHandler) {}
+
     override fun setExpanded(entry: NotificationEntry, expanded: Boolean) {}
+
     override fun setGutsShown(entry: NotificationEntry, gutsShown: Boolean) {}
+
     override fun setHeadsUpAnimatingAway(headsUpAnimatingAway: Boolean) {}
+
     override fun setRemoteInputActive(entry: NotificationEntry, remoteInputActive: Boolean) {}
+
     override fun setTrackingHeadsUp(tracking: Boolean) {}
+
     override fun setUser(user: Int) {}
+
     override fun setUserActionMayIndirectlyRemove(entry: NotificationEntry) {}
+
     override fun shouldSwallowClick(key: String): Boolean = false
+
     override fun showNotification(entry: NotificationEntry) {}
+
     override fun snooze() {}
+
     override fun unpinAll(userUnPinned: Boolean) {}
+
     override fun updateNotification(key: String, alert: Boolean) {}
+
+    override fun onEntryAnimatingAwayEnded(entry: NotificationEntry) {}
 }
 
 @Module
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/OnHeadsUpChangedListener.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/OnHeadsUpChangedListener.java
index 86998ab..de3bf04 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/OnHeadsUpChangedListener.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/OnHeadsUpChangedListener.java
@@ -48,4 +48,9 @@
      * @param isHeadsUp whether the notification is now a headsUp notification
      */
     default void onHeadsUpStateChanged(@NonNull NotificationEntry entry, boolean isHeadsUp) {}
+
+    /**
+     * Called on HUN disappearing animation ends
+     */
+    default void onHeadsUpAnimatingAwayEnded(@NonNull NotificationEntry entry) {}
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt
index cf9a78f..f693409 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt
@@ -20,6 +20,7 @@
 import android.os.UserManager.DISALLOW_CONFIG_LOCATION
 import android.os.UserManager.DISALLOW_MICROPHONE_TOGGLE
 import android.os.UserManager.DISALLOW_SHARE_LOCATION
+import com.android.systemui.Flags
 import com.android.systemui.qs.QsEventLogger
 import com.android.systemui.qs.pipeline.shared.TileSpec
 import com.android.systemui.qs.tileimpl.QSTileImpl
@@ -67,25 +68,18 @@
 import com.android.systemui.qs.tiles.viewmodel.QSTilePolicy
 import com.android.systemui.qs.tiles.viewmodel.QSTileUIConfig
 import com.android.systemui.qs.tiles.viewmodel.QSTileViewModel
+import com.android.systemui.qs.tiles.viewmodel.StubQSTileViewModel
 import com.android.systemui.res.R
 import dagger.Binds
 import dagger.Module
 import dagger.Provides
 import dagger.multibindings.IntoMap
 import dagger.multibindings.StringKey
+import javax.inject.Provider
 
 @Module
 interface PolicyModule {
 
-    /** Inject DndTile into tileMap in QSModule */
-    @Binds @IntoMap @StringKey(DndTile.TILE_SPEC) fun bindDndTile(dndTile: DndTile): QSTileImpl<*>
-
-    /** Inject ModesTile into tileMap in QSModule */
-    @Binds
-    @IntoMap
-    @StringKey(ModesTile.TILE_SPEC)
-    fun bindModesTile(modesTile: ModesTile): QSTileImpl<*>
-
     /** Inject WorkModeTile into tileMap in QSModule */
     @Binds
     @IntoMap
@@ -136,7 +130,19 @@
         const val CAMERA_TOGGLE_TILE_SPEC = "cameratoggle"
         const val MIC_TOGGLE_TILE_SPEC = "mictoggle"
         const val DND_TILE_SPEC = "dnd"
-        const val MODES_TILE_SPEC = "modes"
+
+        /** Inject DndTile or ModesTile into tileMap in QSModule based on feature flag */
+        @Provides
+        @IntoMap
+        @StringKey(DND_TILE_SPEC)
+        fun bindDndOrModesTile(
+            // Using providers to make sure that the unused tile isn't initialised at all if the
+            // flag is off.
+            dndTile: Provider<DndTile>,
+            modesTile: Provider<ModesTile>,
+        ): QSTileImpl<*> {
+            return if (android.app.Flags.modesUi()) modesTile.get() else dndTile.get()
+        }
 
         /** Inject flashlight config */
         @Provides
@@ -386,51 +392,51 @@
             return factory.create(MICROPHONE)
         }
 
-        /** Inject microphone toggle config */
+        /** Inject DND tile or Modes tile config based on feature flag */
         @Provides
         @IntoMap
         @StringKey(DND_TILE_SPEC)
-        fun provideDndTileConfig(uiEventLogger: QsEventLogger): QSTileConfig =
-            QSTileConfig(
-                tileSpec = TileSpec.create(DND_TILE_SPEC),
-                uiConfig =
-                    QSTileUIConfig.Resource(
-                        iconRes = R.drawable.qs_dnd_icon_off,
-                        labelRes = R.string.quick_settings_dnd_label,
-                    ),
-                instanceId = uiEventLogger.getNewInstanceId(),
-            )
-
-        @Provides
-        @IntoMap
-        @StringKey(MODES_TILE_SPEC)
-        fun provideModesTileConfig(uiEventLogger: QsEventLogger): QSTileConfig =
-            QSTileConfig(
-                tileSpec = TileSpec.create(MODES_TILE_SPEC),
-                uiConfig =
-                    QSTileUIConfig.Resource(
-                        iconRes = R.drawable.qs_dnd_icon_off,
-                        labelRes = R.string.quick_settings_modes_label,
-                    ),
-                instanceId = uiEventLogger.getNewInstanceId(),
-            )
+        fun provideDndOrModesTileConfig(uiEventLogger: QsEventLogger): QSTileConfig =
+            if (android.app.Flags.modesUi()) {
+                QSTileConfig(
+                    tileSpec = TileSpec.create(DND_TILE_SPEC),
+                    uiConfig =
+                        QSTileUIConfig.Resource(
+                            iconRes = R.drawable.qs_dnd_icon_off,
+                            labelRes = R.string.quick_settings_modes_label,
+                        ),
+                    instanceId = uiEventLogger.getNewInstanceId(),
+                )
+            } else {
+                QSTileConfig(
+                    tileSpec = TileSpec.create(DND_TILE_SPEC),
+                    uiConfig =
+                        QSTileUIConfig.Resource(
+                            iconRes = R.drawable.qs_dnd_icon_off,
+                            labelRes = R.string.quick_settings_dnd_label,
+                        ),
+                    instanceId = uiEventLogger.getNewInstanceId(),
+                )
+            }
 
         /** Inject ModesTile into tileViewModelMap in QSModule */
         @Provides
         @IntoMap
-        @StringKey(MODES_TILE_SPEC)
+        @StringKey(DND_TILE_SPEC)
         fun provideModesTileViewModel(
             factory: QSTileViewModelFactory.Static<ModesTileModel>,
             mapper: ModesTileMapper,
             stateInteractor: ModesTileDataInteractor,
             userActionInteractor: ModesTileUserActionInteractor
         ): QSTileViewModel =
-            factory.create(
-                TileSpec.create(MODES_TILE_SPEC),
-                userActionInteractor,
-                stateInteractor,
-                mapper,
-            )
+            if (android.app.Flags.modesUi() && Flags.qsNewTilesFuture())
+                factory.create(
+                    TileSpec.create(DND_TILE_SPEC),
+                    userActionInteractor,
+                    stateInteractor,
+                    mapper,
+                )
+            else StubQSTileViewModel
     }
 
     /** Inject FlashlightTile into tileMap in QSModule */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegate.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegate.kt
index 2b094d6..8aa989f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegate.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegate.kt
@@ -18,67 +18,155 @@
 
 import android.content.Intent
 import android.provider.Settings
+import android.util.Log
 import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
 import androidx.compose.ui.res.stringResource
+import androidx.lifecycle.DefaultLifecycleObserver
+import androidx.lifecycle.LifecycleOwner
 import com.android.compose.PlatformButton
 import com.android.compose.PlatformOutlinedButton
+import com.android.internal.annotations.VisibleForTesting
+import com.android.internal.jank.InteractionJankMonitor
+import com.android.systemui.animation.DialogCuj
 import com.android.systemui.animation.DialogTransitionAnimator
+import com.android.systemui.animation.Expandable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.dialog.ui.composable.AlertDialogContent
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.res.R
+import com.android.systemui.statusbar.phone.ComponentSystemUIDialog
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.phone.SystemUIDialogFactory
 import com.android.systemui.statusbar.phone.create
 import com.android.systemui.statusbar.policy.ui.dialog.composable.ModeTileGrid
 import com.android.systemui.statusbar.policy.ui.dialog.viewmodel.ModesDialogViewModel
+import com.android.systemui.util.Assert
 import javax.inject.Inject
+import javax.inject.Provider
+import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.withContext
 
+@SysUISingleton
 class ModesDialogDelegate
 @Inject
 constructor(
     private val sysuiDialogFactory: SystemUIDialogFactory,
     private val dialogTransitionAnimator: DialogTransitionAnimator,
     private val activityStarter: ActivityStarter,
-    private val viewModel: ModesDialogViewModel,
+    // Using a provider to avoid a circular dependency.
+    private val viewModel: Provider<ModesDialogViewModel>,
+    @Main private val mainCoroutineContext: CoroutineContext,
 ) : SystemUIDialog.Delegate {
+    // NOTE: This should only be accessed/written from the main thread.
+    @VisibleForTesting var currentDialog: ComponentSystemUIDialog? = null
+
     override fun createDialog(): SystemUIDialog {
-        return sysuiDialogFactory.create { dialog ->
-            AlertDialogContent(
-                title = { Text(stringResource(R.string.zen_modes_dialog_title)) },
-                content = { ModeTileGrid(viewModel) },
-                neutralButton = {
-                    PlatformOutlinedButton(
-                        onClick = {
-                            val animationController =
-                                dialogTransitionAnimator.createActivityTransitionController(
-                                    dialog.getButton(SystemUIDialog.BUTTON_NEUTRAL)
-                                )
-                            if (animationController == null) {
-                                // The controller will take care of dismissing for us after the
-                                // animation, but let's make sure we dismiss the dialog if we don't
-                                // animate it.
-                                dialog.dismiss()
-                            }
-                            activityStarter.startActivity(
-                                ZEN_MODE_SETTINGS_INTENT,
-                                true /* dismissShade */,
-                                animationController
-                            )
-                        }
-                    ) {
-                        Text(stringResource(R.string.zen_modes_dialog_settings))
+        Assert.isMainThread()
+        if (currentDialog != null) {
+            Log.w(TAG, "Dialog is already open, dismissing it and creating a new one.")
+            currentDialog?.dismiss()
+        }
+
+        currentDialog = sysuiDialogFactory.create() { ModesDialogContent(it) }
+        currentDialog
+            ?.lifecycle
+            ?.addObserver(
+                object : DefaultLifecycleObserver {
+                    override fun onStop(owner: LifecycleOwner) {
+                        Assert.isMainThread()
+                        currentDialog = null
                     }
-                },
-                positiveButton = {
-                    PlatformButton(onClick = { dialog.dismiss() }) {
-                        Text(stringResource(R.string.zen_modes_dialog_done))
-                    }
-                },
+                }
+            )
+
+        return currentDialog!!
+    }
+
+    @Composable
+    private fun ModesDialogContent(dialog: SystemUIDialog) {
+        AlertDialogContent(
+            title = { Text(stringResource(R.string.zen_modes_dialog_title)) },
+            content = { ModeTileGrid(viewModel.get()) },
+            neutralButton = {
+                PlatformOutlinedButton(onClick = { openSettings(dialog) }) {
+                    Text(stringResource(R.string.zen_modes_dialog_settings))
+                }
+            },
+            positiveButton = {
+                PlatformButton(onClick = { dialog.dismiss() }) {
+                    Text(stringResource(R.string.zen_modes_dialog_done))
+                }
+            },
+        )
+    }
+
+    private fun openSettings(dialog: SystemUIDialog) {
+        val animationController =
+            dialogTransitionAnimator.createActivityTransitionController(dialog)
+        if (animationController == null) {
+            // The controller will take care of dismissing for us after
+            // the animation, but let's make sure we dismiss the dialog
+            // if we don't animate it.
+            dialog.dismiss()
+        }
+        activityStarter.startActivity(
+            ZEN_MODE_SETTINGS_INTENT,
+            true /* dismissShade */,
+            animationController
+        )
+    }
+
+    suspend fun showDialog(expandable: Expandable? = null): SystemUIDialog {
+        // Dialogs shown by the DialogTransitionAnimator must be created and shown on the main
+        // thread, so we post it to the UI handler.
+        withContext(mainCoroutineContext) {
+            // Create the dialog if necessary
+            if (currentDialog == null) {
+                createDialog()
+            }
+
+            expandable
+                ?.dialogTransitionController(
+                    DialogCuj(InteractionJankMonitor.CUJ_SHADE_DIALOG_OPEN, INTERACTION_JANK_TAG)
+                )
+                ?.let { controller -> dialogTransitionAnimator.show(currentDialog!!, controller) }
+                ?: currentDialog!!.show()
+        }
+
+        return currentDialog!!
+    }
+
+    /**
+     * Launches the [intent] by animating from the dialog. If the dialog is not showing, just
+     * launches it normally without animating.
+     */
+    fun launchFromDialog(intent: Intent) {
+        Assert.isMainThread()
+        if (currentDialog == null) {
+            Log.w(
+                TAG,
+                "Cannot launch from dialog, the dialog is not present. " +
+                    "Will launch activity without animating."
             )
         }
+
+        val animationController =
+            currentDialog?.let { dialogTransitionAnimator.createActivityTransitionController(it) }
+        if (animationController == null) {
+            currentDialog?.dismiss()
+        }
+        activityStarter.startActivity(
+            intent,
+            true, /* dismissShade */
+            animationController,
+        )
     }
 
     companion object {
+        private const val TAG = "ModesDialogDelegate"
         private val ZEN_MODE_SETTINGS_INTENT = Intent(Settings.ACTION_ZEN_MODE_SETTINGS)
+        private const val INTERACTION_JANK_TAG = "configure_priority_modes"
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt
index 5bd26cc..7c1cb6a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt
@@ -29,7 +29,6 @@
     val text: String,
     val subtext: String,
     val enabled: Boolean,
-    val contentDescription: String,
     val onClick: () -> Unit,
     val onLongClick: () -> Unit,
 )
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt
index e84c8b6..5ffcb34 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt
@@ -17,16 +17,21 @@
 package com.android.systemui.statusbar.policy.ui.dialog.viewmodel
 
 import android.content.Context
+import android.content.Intent
+import android.provider.Settings.ACTION_AUTOMATIC_ZEN_RULE_SETTINGS
+import android.provider.Settings.EXTRA_AUTOMATIC_ZEN_RULE_ID
 import com.android.settingslib.notification.modes.ZenMode
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.policy.domain.interactor.ZenModeInteractor
+import com.android.systemui.statusbar.policy.ui.dialog.ModesDialogDelegate
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.scan
 
 /**
  * Viewmodel for the priority ("zen") modes dialog that can be opened from quick settings. It allows
@@ -39,15 +44,35 @@
     val context: Context,
     zenModeInteractor: ZenModeInteractor,
     @Background val bgDispatcher: CoroutineDispatcher,
+    private val dialogDelegate: ModesDialogDelegate,
 ) {
     // Modes that should be displayed in the dialog
-    // TODO(b/346519570): Include modes that have not been set up yet.
     private val visibleModes: Flow<List<ZenMode>> =
-        zenModeInteractor.modes.map {
-            it.filter { mode ->
-                mode.rule.isEnabled && (mode.isActive || mode.rule.isManualInvocationAllowed)
+        zenModeInteractor.modes
+            // While this is being collected (or in other words, while the dialog is open), we don't
+            // want a mode to disappear from the list if, for instance, the user deactivates it,
+            // since that can be confusing (similar to how we have visual stability for
+            // notifications while the shade is open).
+            // This ensures new modes are added to the list, and updates to modes already in the
+            // list are registered correctly.
+            .scan(listOf()) { prev, modes ->
+                val prevIds = prev.map { it.id }.toSet()
+
+                modes.filter { mode ->
+                    when {
+                        // Mode appeared previously -> keep it even if otherwise we may have
+                        // filtered it
+                        mode.id in prevIds -> true
+                        // Mode is enabled -> show if active (so user can toggle off), or if it
+                        // can be manually toggled on
+                        mode.rule.isEnabled -> mode.isActive || mode.rule.isManualInvocationAllowed
+                        // Mode was created as disabled, or disabled by the app that owns it ->
+                        // will be shown with a "Set up" text
+                        !mode.rule.isEnabled -> mode.status == ZenMode.Status.DISABLED_BY_OTHER
+                        else -> false
+                    }
+                }
             }
-        }
 
     val tiles: Flow<List<ModeTileViewModel>> =
         visibleModes
@@ -59,27 +84,40 @@
                         text = mode.rule.name,
                         subtext = getTileSubtext(mode),
                         enabled = mode.isActive,
-                        // TODO(b/346519570): This should be some combination of the above, e.g.
-                        //  "ON: Do Not Disturb, Until Mon 08:09"; see DndTile.
-                        contentDescription = "",
                         onClick = {
-                            if (mode.isActive) {
+                            if (!mode.rule.isEnabled) {
+                                openSettings(mode)
+                            } else if (mode.isActive) {
                                 zenModeInteractor.deactivateMode(mode)
                             } else {
-                                // TODO(b/346519570): Handle duration for DND mode.
-                                zenModeInteractor.activateMode(mode)
+                                if (mode.rule.isManualInvocationAllowed) {
+                                    // TODO(b/346519570): Handle duration for DND mode.
+                                    zenModeInteractor.activateMode(mode)
+                                }
                             }
                         },
-                        onLongClick = {
-                            // TODO(b/346519570): Open settings page for mode.
-                        }
+                        onLongClick = { openSettings(mode) }
                     )
                 }
             }
             .flowOn(bgDispatcher)
 
+    private fun openSettings(mode: ZenMode) {
+        val intent: Intent =
+            Intent(ACTION_AUTOMATIC_ZEN_RULE_SETTINGS)
+                .putExtra(EXTRA_AUTOMATIC_ZEN_RULE_ID, mode.id)
+
+        dialogDelegate.launchFromDialog(intent)
+    }
+
     private fun getTileSubtext(mode: ZenMode): String {
-        // TODO(b/346519570): Use ZenModeConfig.getDescription for manual DND
+        if (!mode.rule.isEnabled) {
+            return context.resources.getString(R.string.zen_mode_set_up)
+        }
+        if (!mode.rule.isManualInvocationAllowed && !mode.isActive) {
+            return context.resources.getString(R.string.zen_mode_no_manual_invocation)
+        }
+
         val on = context.resources.getString(R.string.zen_mode_on)
         val off = context.resources.getString(R.string.zen_mode_off)
         return mode.rule.triggerDescription ?: if (mode.isActive) on else off
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt
index 4869114..89227cf 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.domain.interactor.KeyguardStatusBarInteractor
 import com.android.systemui.statusbar.notification.domain.interactor.HeadsUpNotificationInteractor
+import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.policy.BatteryController
 import com.android.systemui.statusbar.policy.BatteryController.BatteryStateChangeCallback
 import javax.inject.Inject
@@ -33,6 +34,7 @@
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.stateIn
 
 /**
@@ -54,12 +56,20 @@
     keyguardStatusBarInteractor: KeyguardStatusBarInteractor,
     batteryController: BatteryController,
 ) {
+
+    private val showingHeadsUpStatusBar: Flow<Boolean> =
+        if (NotificationsHeadsUpRefactor.isEnabled) {
+            headsUpNotificationInteractor.showHeadsUpStatusBar
+        } else {
+            flowOf(false)
+        }
+
     /** True if this view should be visible and false otherwise. */
     val isVisible: StateFlow<Boolean> =
         combine(
                 sceneInteractor.currentScene,
                 keyguardInteractor.isDozing,
-                headsUpNotificationInteractor.showHeadsUpStatusBar,
+                showingHeadsUpStatusBar,
             ) { currentScene, isDozing, showHeadsUpStatusBar ->
                 currentScene == Scenes.Lockscreen && !isDozing && !showHeadsUpStatusBar
             }
diff --git a/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/BackGestureTutorialScreen.kt b/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/BackGestureTutorialScreen.kt
index 37833d8..94ff65e 100644
--- a/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/BackGestureTutorialScreen.kt
+++ b/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/BackGestureTutorialScreen.kt
@@ -21,7 +21,9 @@
 import android.graphics.PorterDuffColorFilter
 import androidx.activity.compose.BackHandler
 import androidx.annotation.StringRes
-import androidx.compose.foundation.background
+import androidx.compose.animation.animateColorAsState
+import androidx.compose.animation.core.LinearEasing
+import androidx.compose.animation.core.tween
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
@@ -40,6 +42,7 @@
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.draw.drawBehind
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.toArgb
 import androidx.compose.ui.input.pointer.pointerInteropFilter
@@ -63,6 +66,7 @@
 
 data class TutorialScreenColors(
     val backgroundColor: Color,
+    val successBackgroundColor: Color,
     val titleColor: Color,
     val animationProperties: LottieDynamicProperties
 )
@@ -100,6 +104,7 @@
     val onTertiaryFixed = LocalAndroidColorScheme.current.onTertiaryFixed
     val onTertiaryFixedVariant = LocalAndroidColorScheme.current.onTertiaryFixedVariant
     val tertiaryFixedDim = LocalAndroidColorScheme.current.tertiaryFixedDim
+    val surfaceContainer = MaterialTheme.colorScheme.surfaceContainer
     val dynamicProperties =
         rememberLottieDynamicProperties(
             rememberColorFilterProperty(".tertiaryFixedDim", tertiaryFixedDim),
@@ -108,9 +113,10 @@
             rememberColorFilterProperty(".onTertiaryFixedVariant", onTertiaryFixedVariant)
         )
     val screenColors =
-        remember(onTertiaryFixed, tertiaryFixedDim, dynamicProperties) {
+        remember(onTertiaryFixed, surfaceContainer, tertiaryFixedDim, dynamicProperties) {
             TutorialScreenColors(
                 backgroundColor = onTertiaryFixed,
+                successBackgroundColor = surfaceContainer,
                 titleColor = tertiaryFixedDim,
                 animationProperties = dynamicProperties,
             )
@@ -124,11 +130,19 @@
     onDoneButtonClicked: () -> Unit,
     screenColors: TutorialScreenColors
 ) {
+    val animatedColor by
+        animateColorAsState(
+            targetValue =
+                if (gestureDone) screenColors.successBackgroundColor
+                else screenColors.backgroundColor,
+            animationSpec = tween(durationMillis = 150, easing = LinearEasing),
+            label = "backgroundColor"
+        )
     Column(
         verticalArrangement = Arrangement.Center,
         modifier =
             Modifier.fillMaxSize()
-                .background(color = screenColors.backgroundColor)
+                .drawBehind { drawRect(animatedColor) }
                 .padding(start = 48.dp, top = 124.dp, end = 48.dp, bottom = 48.dp)
     ) {
         Row(modifier = Modifier.fillMaxWidth().weight(1f)) {
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 8ab5bc6..169f865 100644
--- a/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensorImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/util/sensors/ProximitySensorImpl.java
@@ -26,8 +26,8 @@
 import com.android.systemui.util.concurrency.DelayableExecutor;
 import com.android.systemui.util.concurrency.Execution;
 
-import java.util.ArrayList;
-import java.util.List;
+import java.util.Set;
+import java.util.concurrent.CopyOnWriteArraySet;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import javax.inject.Inject;
@@ -64,7 +64,7 @@
     ThresholdSensor mSecondaryThresholdSensor;
     private final DelayableExecutor mDelayableExecutor;
     private final Execution mExecution;
-    private final List<ThresholdSensor.Listener> mListeners = new ArrayList<>();
+    private final Set<Listener> mListeners = new CopyOnWriteArraySet<>();
     private String mTag = null;
     @VisibleForTesting protected boolean mPaused;
     private ThresholdSensorEvent mLastPrimaryEvent;
@@ -246,7 +246,7 @@
     public void unregister(ThresholdSensor.Listener listener) {
         mExecution.assertIsMainThread();
         mListeners.remove(listener);
-        if (mListeners.size() == 0) {
+        if (mListeners.isEmpty()) {
             unregisterInternal();
         }
     }
@@ -296,8 +296,7 @@
         }
         if (mLastEvent != null) {
             ThresholdSensorEvent lastEvent = mLastEvent;  // Listeners can null out mLastEvent.
-            List<ThresholdSensor.Listener> listeners = new ArrayList<>(mListeners);
-            listeners.forEach(proximitySensorListener ->
+            mListeners.forEach(proximitySensorListener ->
                     proximitySensorListener.onThresholdCrossed(lastEvent));
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt b/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt
index b5934ec..9125a91 100644
--- a/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt
@@ -336,7 +336,7 @@
      * @param name to look up in the table
      * @return the corresponding value, or null if not present
      */
-    fun getString(name: String): String
+    fun getString(name: String): String?
 
     /**
      * Store a name/value pair into the database.
@@ -385,15 +385,15 @@
      * an integer.
      *
      * @param name The name of the setting to retrieve.
-     * @param def Value to return if the setting is not defined.
-     * @return The setting's current value, or 'def' if it is not defined or not a valid integer.
+     * @param default Value to return if the setting is not defined.
+     * @return The setting's current value, or default if it is not defined or not a valid integer.
      */
-    fun getInt(name: String, def: Int): Int {
+    fun getInt(name: String, default: Int): Int {
         val v = getString(name)
         return try {
-            v.toInt()
+            v?.toInt() ?: default
         } catch (e: NumberFormatException) {
-            def
+            default
         }
     }
 
@@ -412,7 +412,7 @@
      */
     @Throws(SettingNotFoundException::class)
     fun getInt(name: String): Int {
-        val v = getString(name)
+        val v = getString(name) ?: throw SettingNotFoundException(name)
         return try {
             v.toInt()
         } catch (e: NumberFormatException) {
@@ -441,11 +441,11 @@
      * boolean.
      *
      * @param name The name of the setting to retrieve.
-     * @param def Value to return if the setting is not defined.
-     * @return The setting's current value, or 'def' if it is not defined or not a valid boolean.
+     * @param default Value to return if the setting is not defined.
+     * @return The setting's current value, or default if it is not defined or not a valid boolean.
      */
-    fun getBool(name: String, def: Boolean): Boolean {
-        return getInt(name, if (def) 1 else 0) != 0
+    fun getBool(name: String, default: Boolean): Boolean {
+        return getInt(name, if (default) 1 else 0) != 0
     }
 
     /**
@@ -579,13 +579,12 @@
     companion object {
         /** Convert a string to a long, or uses a default if the string is malformed or null */
         @JvmStatic
-        fun parseLongOrUseDefault(valString: String, def: Long): Long {
-            val value: Long
-            value =
+        fun parseLongOrUseDefault(valString: String?, default: Long): Long {
+            val value: Long =
                 try {
-                    valString.toLong()
+                    valString?.toLong() ?: default
                 } catch (e: NumberFormatException) {
-                    def
+                    default
                 }
             return value
         }
diff --git a/packages/SystemUI/src/com/android/systemui/util/settings/UserSettingsProxy.kt b/packages/SystemUI/src/com/android/systemui/util/settings/UserSettingsProxy.kt
index 848a6e6..ac7c1ce 100644
--- a/packages/SystemUI/src/com/android/systemui/util/settings/UserSettingsProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/settings/UserSettingsProxy.kt
@@ -354,12 +354,12 @@
      * @param name to look up in the table
      * @return the corresponding value, or null if not present
      */
-    override fun getString(name: String): String {
+    override fun getString(name: String): String? {
         return getStringForUser(name, userId)
     }
 
     /** See [getString]. */
-    fun getStringForUser(name: String, userHandle: Int): String
+    fun getStringForUser(name: String, userHandle: Int): String?
 
     /**
      * Store a name/value pair into the database. Values written by this method will be overridden
@@ -388,17 +388,17 @@
         overrideableByRestore: Boolean
     ): Boolean
 
-    override fun getInt(name: String, def: Int): Int {
-        return getIntForUser(name, def, userId)
+    override fun getInt(name: String, default: Int): Int {
+        return getIntForUser(name, default, userId)
     }
 
     /** Similar implementation to [getInt] for the specified [userHandle]. */
-    fun getIntForUser(name: String, def: Int, userHandle: Int): Int {
+    fun getIntForUser(name: String, default: Int, userHandle: Int): Int {
         val v = getStringForUser(name, userHandle)
         return try {
-            v.toInt()
+            v?.toInt() ?: default
         } catch (e: NumberFormatException) {
-            def
+            default
         }
     }
 
@@ -408,7 +408,7 @@
     /** Similar implementation to [getInt] for the specified [userHandle]. */
     @Throws(SettingNotFoundException::class)
     fun getIntForUser(name: String, userHandle: Int): Int {
-        val v = getStringForUser(name, userHandle)
+        val v = getStringForUser(name, userHandle) ?: throw SettingNotFoundException(name)
         return try {
             v.toInt()
         } catch (e: NumberFormatException) {
diff --git a/packages/SystemUI/src/com/android/systemui/volume/CsdWarningAction.kt b/packages/SystemUI/src/com/android/systemui/volume/CsdWarningAction.kt
new file mode 100644
index 0000000..a77acb5
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/volume/CsdWarningAction.kt
@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2024 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.volume
+
+import android.app.PendingIntent
+import android.app.PendingIntent.FLAG_IMMUTABLE
+import android.app.PendingIntent.FLAG_UPDATE_CURRENT
+import android.content.Context
+import android.content.Intent
+
+/**
+ * label: Notification action label text. intent: The Intent used to start Activity or Broadcast.
+ * isActivity: Defines if the pending intent should start an activity. Default is to broadcast
+ */
+data class CsdWarningAction(
+    val label: String? = null,
+    val intent: Intent? = null,
+    val isActivity: Boolean = false,
+) {
+    fun toPendingIntent(context: Context): PendingIntent? {
+        if (label == null || intent == null) {
+            return null
+        }
+        if (isActivity) {
+            return PendingIntent.getActivity(
+                context,
+                0,
+                intent,
+                PendingIntent.FLAG_UPDATE_CURRENT or PendingIntent.FLAG_IMMUTABLE
+            )
+        }
+        return PendingIntent.getBroadcast(context, 0, intent, FLAG_IMMUTABLE)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/volume/CsdWarningDialog.java b/packages/SystemUI/src/com/android/systemui/volume/CsdWarningDialog.java
index bb230e6..a63660b 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/CsdWarningDialog.java
+++ b/packages/SystemUI/src/com/android/systemui/volume/CsdWarningDialog.java
@@ -30,7 +30,6 @@
 import android.media.AudioManager;
 import android.provider.Settings;
 import android.util.Log;
-import android.util.Pair;
 import android.view.KeyEvent;
 import android.view.WindowManager;
 
@@ -109,7 +108,7 @@
     private long mShowTime;
 
     @VisibleForTesting public int mCachedMediaStreamVolume;
-    private Optional<ImmutableList<Pair<String, Intent>>> mActionIntents;
+    private Optional<ImmutableList<CsdWarningAction>> mActionIntents;
     private final BroadcastDispatcher mBroadcastDispatcher;
 
     /**
@@ -121,7 +120,7 @@
         CsdWarningDialog create(
                 int csdWarning,
                 Runnable onCleanup,
-                Optional<ImmutableList<Pair<String, Intent>>> actionIntents);
+                Optional<ImmutableList<CsdWarningAction>> actionIntents);
     }
 
     @AssistedInject
@@ -132,7 +131,7 @@
             NotificationManager notificationManager,
             @Background DelayableExecutor delayableExecutor,
             @Assisted Runnable onCleanup,
-            @Assisted Optional<ImmutableList<Pair<String, Intent>>> actionIntents,
+            @Assisted Optional<ImmutableList<CsdWarningAction>> actionIntents,
             BroadcastDispatcher broadcastDispatcher) {
         super(context);
         mCsdWarning = csdWarning;
@@ -351,39 +350,45 @@
         if (Flags.sounddoseCustomization()
                 && mActionIntents.isPresent()
                 && !mActionIntents.get().isEmpty()) {
-            ImmutableList<Pair<String, Intent>> actionIntentsList = mActionIntents.get();
-            for (Pair<String, Intent> intentPair : actionIntentsList) {
-                if (intentPair != null && intentPair.first != null && intentPair.second != null) {
-                    PendingIntent pendingActionIntent =
-                            PendingIntent.getBroadcast(mContext, 0, intentPair.second,
-                                    FLAG_IMMUTABLE);
-                    builder.addAction(0, intentPair.first, pendingActionIntent);
-                    // Register receiver to undo volume only when
-                    // notification conaining the undo action would be sent.
-                    if (intentPair.first == mContext.getString(R.string.volume_undo_action)) {
-                        final IntentFilter filterUndo = new IntentFilter(
-                                VolumeDialog.ACTION_VOLUME_UNDO);
-                        mBroadcastDispatcher.registerReceiver(mReceiverUndo,
-                                filterUndo,
-                                /* executor = default */ null,
-                                /* user = default */ null,
-                                Context.RECEIVER_NOT_EXPORTED,
-                                /* permission = default */ null);
+            ImmutableList<CsdWarningAction> actionIntentsList = mActionIntents.get();
+            for (CsdWarningAction action : actionIntentsList) {
+                if (action.getLabel() == null || action.getIntent() == null) {
+                    Log.w(TAG, "Null action intent received. Skipping addition to notification");
+                    continue;
+                }
+                PendingIntent pendingActionIntent = action.toPendingIntent(mContext);
+                if (pendingActionIntent == null) {
+                    Log.w(TAG, "Null pending intent received. Skipping addition to notification");
+                    continue;
+                }
+                builder.addAction(0, action.getLabel(), pendingActionIntent);
 
-                        // Register receiver to learn if notification has been dismissed.
-                        // This is required to unregister receivers to prevent leak.
-                        Intent dismissIntent = new Intent(DISMISS_CSD_NOTIFICATION)
-                                .setPackage(mContext.getPackageName());
-                        PendingIntent pendingDismissIntent = PendingIntent.getBroadcast(mContext,
-                                0, dismissIntent, FLAG_IMMUTABLE);
-                        mBroadcastDispatcher.registerReceiver(mReceiverDismissNotification,
-                                new IntentFilter(DISMISS_CSD_NOTIFICATION),
-                                /* executor = default */ null,
-                                /* user = default */ null,
-                                Context.RECEIVER_NOT_EXPORTED,
-                                /* permission = default */ null);
-                        builder.setDeleteIntent(pendingDismissIntent);
-                    }
+                // Register receiver to undo volume only when
+                // notification conaining the undo action would be sent.
+                if (action.getLabel().equals(mContext.getString(R.string.volume_undo_action))) {
+                    final IntentFilter filterUndo = new IntentFilter(
+                            VolumeDialog.ACTION_VOLUME_UNDO);
+                    mBroadcastDispatcher.registerReceiver(mReceiverUndo,
+                            filterUndo,
+                            /* executor = default */ null,
+                            /* user = default */ null,
+                            Context.RECEIVER_NOT_EXPORTED,
+                            /* permission = default */ null);
+
+                    // Register receiver to learn if notification has been dismissed.
+                    // This is required to unregister receivers to prevent leak.
+                    Intent dismissIntent = new Intent(DISMISS_CSD_NOTIFICATION)
+                            .setPackage(mContext.getPackageName());
+                    PendingIntent pendingDismissIntent = PendingIntent.getBroadcast(
+                            mContext,
+                            0, dismissIntent, FLAG_IMMUTABLE);
+                    mBroadcastDispatcher.registerReceiver(mReceiverDismissNotification,
+                            new IntentFilter(DISMISS_CSD_NOTIFICATION),
+                            /* executor = default */ null,
+                            /* user = default */ null,
+                            Context.RECEIVER_NOT_EXPORTED,
+                            /* permission = default */ null);
+                    builder.setDeleteIntent(pendingDismissIntent);
                 }
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java b/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java
index 0770d89..e56f6b3 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java
@@ -51,7 +51,6 @@
 import android.content.ContentResolver;
 import android.content.Context;
 import android.content.DialogInterface;
-import android.content.Intent;
 import android.content.pm.PackageManager;
 import android.content.res.ColorStateList;
 import android.content.res.Configuration;
@@ -79,7 +78,6 @@
 import android.provider.Settings.Global;
 import android.text.InputFilter;
 import android.util.Log;
-import android.util.Pair;
 import android.util.Slog;
 import android.util.SparseBooleanArray;
 import android.view.ContextThemeWrapper;
@@ -322,8 +320,8 @@
     private final VolumePanelFlag mVolumePanelFlag;
     private final VolumeDialogInteractor mInteractor;
     // Optional actions for soundDose
-    private Optional<ImmutableList<Pair<String, Intent>>> mCsdWarningNotificationActions =
-            Optional.of(ImmutableList.of());
+    private Optional<ImmutableList<CsdWarningAction>>
+            mCsdWarningNotificationActions = Optional.of(ImmutableList.of());
 
     public VolumeDialogImpl(
             Context context,
@@ -2231,7 +2229,7 @@
     }
 
     public void setCsdWarningNotificationActionIntents(
-            ImmutableList<Pair<String, Intent>> actionIntent) {
+            ImmutableList<CsdWarningAction> actionIntent) {
         mCsdWarningNotificationActions = Optional.of(actionIntent);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java b/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java
index dc2b80c..68d12f6 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java
+++ b/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java
@@ -16,6 +16,8 @@
 
 package com.android.systemui.volume;
 
+import static com.android.settingslib.flags.Flags.volumeDialogAudioSharingFix;
+
 import android.content.Context;
 import android.content.res.Configuration;
 import android.os.Handler;
@@ -26,6 +28,7 @@
 import com.android.systemui.qs.tiles.DndTile;
 import com.android.systemui.res.R;
 import com.android.systemui.statusbar.policy.ConfigurationController;
+import com.android.systemui.volume.domain.interactor.AudioSharingInteractor;
 
 import java.io.PrintWriter;
 
@@ -41,23 +44,29 @@
     private boolean mEnabled;
     private final Context mContext;
     private VolumeDialogComponent mVolumeComponent;
+    private AudioSharingInteractor mAudioSharingInteractor;
 
     @Inject
-    public VolumeUI(Context context, VolumeDialogComponent volumeDialogComponent) {
+    public VolumeUI(Context context, VolumeDialogComponent volumeDialogComponent,
+            AudioSharingInteractor audioSharingInteractor) {
         mContext = context;
         mVolumeComponent = volumeDialogComponent;
+        mAudioSharingInteractor = audioSharingInteractor;
     }
 
     @Override
     public void start() {
         boolean enableVolumeUi = mContext.getResources().getBoolean(R.bool.enable_volume_ui);
         boolean enableSafetyWarning =
-            mContext.getResources().getBoolean(R.bool.enable_safety_warning);
+                mContext.getResources().getBoolean(R.bool.enable_safety_warning);
         mEnabled = enableVolumeUi || enableSafetyWarning;
         if (!mEnabled) return;
 
         mVolumeComponent.setEnableDialogs(enableVolumeUi, enableSafetyWarning);
         setDefaultVolumeController();
+        if (volumeDialogAudioSharingFix()) {
+            mAudioSharingInteractor.handlePrimaryGroupChange();
+        }
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt b/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt
index 0c1bc21..efaca7a 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt
@@ -75,7 +75,6 @@
         @Provides
         @SysUISingleton
         fun provideAudioSharingRepository(
-            @Application context: Context,
             contentResolver: ContentResolver,
             localBluetoothManager: LocalBluetoothManager?,
             @Application coroutineScope: CoroutineScope,
diff --git a/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractor.kt
index aba3015..2170c36 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractor.kt
@@ -17,18 +17,28 @@
 package com.android.systemui.volume.domain.interactor
 
 import android.bluetooth.BluetoothCsipSetCoordinator
+import android.media.AudioManager.STREAM_MUSIC
 import androidx.annotation.IntRange
 import com.android.settingslib.volume.data.repository.AudioSharingRepository
 import com.android.settingslib.volume.data.repository.AudioSharingRepository.Companion.AUDIO_SHARING_VOLUME_MAX
 import com.android.settingslib.volume.data.repository.AudioSharingRepository.Companion.AUDIO_SHARING_VOLUME_MIN
+import com.android.settingslib.volume.domain.interactor.AudioVolumeInteractor
+import com.android.settingslib.volume.shared.model.AudioStream
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
 import javax.inject.Inject
+import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 
 interface AudioSharingInteractor {
     /** Audio sharing secondary headset volume changes. */
@@ -45,6 +55,16 @@
         @IntRange(from = AUDIO_SHARING_VOLUME_MIN.toLong(), to = AUDIO_SHARING_VOLUME_MAX.toLong())
         level: Int
     )
+
+    /**
+     * Handle primary group change in audio sharing.
+     *
+     * Once the primary group is changed, we need to sync its volume to STREAM_MUSIC to make sure
+     * the volume adjustment during audio sharing can be kept after the sharing ends.
+     *
+     * TODO(b/355396988) Migrate to audio framework solution once it is in place.
+     */
+    fun handlePrimaryGroupChange()
 }
 
 @SysUISingleton
@@ -52,26 +72,60 @@
 @Inject
 constructor(
     @Application private val coroutineScope: CoroutineScope,
+    @Background private val backgroundCoroutineContext: CoroutineContext,
+    private val audioVolumeInteractor: AudioVolumeInteractor,
     private val audioSharingRepository: AudioSharingRepository
 ) : AudioSharingInteractor {
 
     override val volume: Flow<Int?> =
         combine(audioSharingRepository.secondaryGroupId, audioSharingRepository.volumeMap) {
-            secondaryGroupId,
-            volumeMap ->
-            if (secondaryGroupId == BluetoothCsipSetCoordinator.GROUP_ID_INVALID) null
-            else volumeMap.getOrDefault(secondaryGroupId, DEFAULT_VOLUME)
-        }
+                secondaryGroupId,
+                volumeMap ->
+                if (secondaryGroupId == BluetoothCsipSetCoordinator.GROUP_ID_INVALID) null
+                else volumeMap.getOrDefault(secondaryGroupId, DEFAULT_VOLUME)
+            }
+            .distinctUntilChanged()
 
     override val volumeMin: Int = AUDIO_SHARING_VOLUME_MIN
 
     override val volumeMax: Int = AUDIO_SHARING_VOLUME_MAX
 
-    override fun setStreamVolume(level: Int) {
+    override fun setStreamVolume(
+        @IntRange(from = AUDIO_SHARING_VOLUME_MIN.toLong(), to = AUDIO_SHARING_VOLUME_MAX.toLong())
+        level: Int
+    ) {
         coroutineScope.launch { audioSharingRepository.setSecondaryVolume(level) }
     }
 
+    override fun handlePrimaryGroupChange() {
+        coroutineScope.launch {
+            audioSharingRepository.primaryGroupId
+                .map { primaryGroupId -> audioSharingRepository.volumeMap.value[primaryGroupId] }
+                .filterNotNull()
+                .distinctUntilChanged()
+                .collect {
+                    // Once primary device change, we need to update the STREAM_MUSIC volume to get
+                    // align with the primary device's volume
+                    setMusicStreamVolume(it)
+                }
+        }
+    }
+
+    private suspend fun setMusicStreamVolume(volume: Int) {
+        withContext(backgroundCoroutineContext) {
+            val musicStream =
+                audioVolumeInteractor.getAudioStream(AudioStream(STREAM_MUSIC)).first()
+            val musicVolume =
+                Math.round(
+                    volume.toFloat() * (musicStream.maxVolume - musicStream.minVolume) /
+                        (AUDIO_SHARING_VOLUME_MAX - AUDIO_SHARING_VOLUME_MIN)
+                )
+            audioVolumeInteractor.setVolume(AudioStream(STREAM_MUSIC), musicVolume)
+        }
+    }
+
     private companion object {
+        const val TAG = "AudioSharingInteractor"
         const val DEFAULT_VOLUME = 20
     }
 }
@@ -82,7 +136,12 @@
     override val volumeMin: Int = EMPTY_VOLUME
     override val volumeMax: Int = EMPTY_VOLUME
 
-    override fun setStreamVolume(level: Int) {}
+    override fun setStreamVolume(
+        @IntRange(from = AUDIO_SHARING_VOLUME_MIN.toLong(), to = AUDIO_SHARING_VOLUME_MAX.toLong())
+        level: Int
+    ) {}
+
+    override fun handlePrimaryGroupChange() {}
 
     private companion object {
         const val EMPTY_VOLUME = 0
diff --git a/packages/SystemUI/tests/src/com/android/systemui/accessibility/FullscreenMagnificationControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/accessibility/FullscreenMagnificationControllerTest.java
index cbd535b..530ae15 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/accessibility/FullscreenMagnificationControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/accessibility/FullscreenMagnificationControllerTest.java
@@ -25,6 +25,7 @@
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -36,7 +37,10 @@
 import android.content.Context;
 import android.content.pm.ActivityInfo;
 import android.graphics.Rect;
+import android.graphics.drawable.GradientDrawable;
+import android.hardware.display.DisplayManager;
 import android.os.RemoteException;
+import android.platform.test.annotations.EnableFlags;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.Display;
@@ -54,6 +58,7 @@
 import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
+import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.res.R;
 
@@ -61,6 +66,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -76,6 +82,12 @@
     private static final long WAIT_TIMEOUT_S = 5L * HW_TIMEOUT_MULTIPLIER;
     private static final long ANIMATION_TIMEOUT_MS =
             5L * ANIMATION_DURATION_MS * HW_TIMEOUT_MULTIPLIER;
+
+    private static final String UNIQUE_DISPLAY_ID_PRIMARY = "000";
+    private static final String UNIQUE_DISPLAY_ID_SECONDARY = "111";
+    private static final int CORNER_RADIUS_PRIMARY = 10;
+    private static final int CORNER_RADIUS_SECONDARY = 20;
+
     private FullscreenMagnificationController mFullscreenMagnificationController;
     private SurfaceControlViewHost mSurfaceControlViewHost;
     private ValueAnimator mShowHideBorderAnimator;
@@ -83,10 +95,35 @@
     private TestableWindowManager mWindowManager;
     @Mock
     private IWindowManager mIWindowManager;
+    @Mock
+    private DisplayManager mDisplayManager;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+        mContext = spy(mContext);
+        Display display = mock(Display.class);
+        when(display.getUniqueId()).thenReturn(UNIQUE_DISPLAY_ID_PRIMARY);
+        when(mContext.getDisplayNoVerify()).thenReturn(display);
+
+        // Override the resources to Display Primary
+        mContext.getOrCreateTestableResources()
+                .addOverride(
+                        com.android.internal.R.dimen.rounded_corner_radius,
+                        CORNER_RADIUS_PRIMARY);
+        mContext.getOrCreateTestableResources()
+                .addOverride(com.android.internal.R.dimen.rounded_corner_radius_adjustment, 0);
+        mContext.getOrCreateTestableResources()
+                .addOverride(com.android.internal.R.dimen.rounded_corner_radius_top, 0);
+        mContext.getOrCreateTestableResources()
+                .addOverride(
+                        com.android.internal.R.dimen.rounded_corner_radius_top_adjustment, 0);
+        mContext.getOrCreateTestableResources()
+                .addOverride(com.android.internal.R.dimen.rounded_corner_radius_bottom, 0);
+        mContext.getOrCreateTestableResources()
+                .addOverride(
+                        com.android.internal.R.dimen.rounded_corner_radius_bottom_adjustment, 0);
+
         getInstrumentation().runOnMainSync(() -> mSurfaceControlViewHost =
                 spy(new SurfaceControlViewHost(mContext, mContext.getDisplay(),
                         new InputTransferToken(), "FullscreenMagnification")));
@@ -101,6 +138,7 @@
                 mContext,
                 mContext.getMainThreadHandler(),
                 mContext.getMainExecutor(),
+                mDisplayManager,
                 mContext.getSystemService(AccessibilityManager.class),
                 mContext.getSystemService(WindowManager.class),
                 mIWindowManager,
@@ -259,6 +297,87 @@
         verify(mSurfaceControlViewHost).relayout(newWidth, newHeight);
     }
 
+    @EnableFlags(Flags.FLAG_UPDATE_CORNER_RADIUS_ON_DISPLAY_CHANGED)
+    @Test
+    public void enableFullscreenMagnification_applyPrimaryCornerRadius()
+            throws InterruptedException {
+        CountDownLatch transactionCommittedLatch = new CountDownLatch(1);
+        CountDownLatch animationEndLatch = new CountDownLatch(1);
+        mTransaction.addTransactionCommittedListener(
+                Runnable::run, transactionCommittedLatch::countDown);
+        mShowHideBorderAnimator.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                animationEndLatch.countDown();
+            }
+        });
+
+        getInstrumentation().runOnMainSync(() ->
+                //Enable fullscreen magnification
+                mFullscreenMagnificationController
+                        .onFullscreenMagnificationActivationChanged(true));
+        assertWithMessage("Failed to wait for transaction committed")
+                .that(transactionCommittedLatch.await(WAIT_TIMEOUT_S, TimeUnit.SECONDS))
+                .isTrue();
+        assertWithMessage("Failed to wait for animation to be finished")
+                .that(animationEndLatch.await(ANIMATION_TIMEOUT_MS, TimeUnit.MILLISECONDS))
+                .isTrue();
+
+        // Verify the initial corner radius is applied
+        GradientDrawable backgroundDrawable =
+                (GradientDrawable) mSurfaceControlViewHost.getView().getBackground();
+        assertThat(backgroundDrawable.getCornerRadius()).isEqualTo(CORNER_RADIUS_PRIMARY);
+    }
+
+    @EnableFlags(Flags.FLAG_UPDATE_CORNER_RADIUS_ON_DISPLAY_CHANGED)
+    @Test
+    public void onDisplayChanged_updateCornerRadiusToSecondary() throws InterruptedException {
+        CountDownLatch transactionCommittedLatch = new CountDownLatch(1);
+        CountDownLatch animationEndLatch = new CountDownLatch(1);
+        mTransaction.addTransactionCommittedListener(
+                Runnable::run, transactionCommittedLatch::countDown);
+        mShowHideBorderAnimator.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                animationEndLatch.countDown();
+            }
+        });
+
+        getInstrumentation().runOnMainSync(() ->
+                //Enable fullscreen magnification
+                mFullscreenMagnificationController
+                        .onFullscreenMagnificationActivationChanged(true));
+        assertWithMessage("Failed to wait for transaction committed")
+                .that(transactionCommittedLatch.await(WAIT_TIMEOUT_S, TimeUnit.SECONDS))
+                .isTrue();
+        assertWithMessage("Failed to wait for animation to be finished")
+                .that(animationEndLatch.await(ANIMATION_TIMEOUT_MS, TimeUnit.MILLISECONDS))
+                .isTrue();
+
+        ArgumentCaptor<DisplayManager.DisplayListener> displayListenerCaptor =
+                ArgumentCaptor.forClass(DisplayManager.DisplayListener.class);
+        verify(mDisplayManager).registerDisplayListener(displayListenerCaptor.capture(), any());
+
+        Display newDisplay = mock(Display.class);
+        when(newDisplay.getUniqueId()).thenReturn(UNIQUE_DISPLAY_ID_SECONDARY);
+        when(mContext.getDisplayNoVerify()).thenReturn(newDisplay);
+        // Override the resources to Display Secondary
+        mContext.getOrCreateTestableResources()
+                .removeOverride(com.android.internal.R.dimen.rounded_corner_radius);
+        mContext.getOrCreateTestableResources()
+                .addOverride(
+                        com.android.internal.R.dimen.rounded_corner_radius,
+                        CORNER_RADIUS_SECONDARY);
+        getInstrumentation().runOnMainSync(() ->
+                displayListenerCaptor.getValue().onDisplayChanged(Display.DEFAULT_DISPLAY));
+        waitForIdleSync();
+        // Verify the corner radius is updated
+        GradientDrawable backgroundDrawable2 =
+                (GradientDrawable) mSurfaceControlViewHost.getView().getBackground();
+        assertThat(backgroundDrawable2.getCornerRadius()).isEqualTo(CORNER_RADIUS_SECONDARY);
+    }
+
+
     private ValueAnimator newNullTargetObjectAnimator() {
         final ValueAnimator animator =
                 ObjectAnimator.ofFloat(/* target= */ null, View.ALPHA, 0f, 1f);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
index 6dcea14..9df653f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
@@ -74,7 +74,6 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
-import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.withArgCaptor
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -92,6 +91,7 @@
 import org.mockito.Mock
 import org.mockito.Mockito
 import org.mockito.junit.MockitoJUnit
+import org.mockito.kotlin.any
 import org.mockito.kotlin.whenever
 import platform.test.runner.parameterized.ParameterizedAndroidJunit4
 import platform.test.runner.parameterized.Parameters
@@ -1379,6 +1379,28 @@
     }
 
     @Test
+    @EnableFlags(FLAG_BP_TALKBACK)
+    fun no_hint_for_talkback_guidance_after_auth() = runGenericTest {
+        val hint by collectLastValue(kosmos.promptViewModel.accessibilityHint)
+
+        kosmos.promptViewModel.showAuthenticated(testCase.authenticatedModality, 0)
+        kosmos.promptViewModel.confirmAuthenticated()
+
+        // Touches should fall outside of sensor area
+        whenever(kosmos.udfpsUtils.getTouchInNativeCoordinates(any(), any(), any()))
+            .thenReturn(Point(0, 0))
+        whenever(kosmos.udfpsUtils.onTouchOutsideOfSensorArea(any(), any(), any(), any(), any()))
+            .thenReturn("Direction")
+
+        kosmos.promptViewModel.onAnnounceAccessibilityHint(
+            obtainMotionEvent(MotionEvent.ACTION_HOVER_ENTER),
+            true
+        )
+
+        assertThat(hint.isNullOrBlank()).isTrue()
+    }
+
+    @Test
     @EnableFlags(FLAG_CUSTOM_BIOMETRIC_PROMPT, FLAG_CONSTRAINT_BP)
     fun descriptionOverriddenByVerticalListContentView() =
         runGenericTest(description = "test description", contentView = promptContentView) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractorTest.kt
index 63b4ff7..72e0726 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractorTest.kt
@@ -18,6 +18,7 @@
 
 import android.content.res.Configuration
 import android.graphics.Rect
+import android.util.LayoutDirection
 import android.view.Surface
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -34,6 +35,8 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @ExperimentalCoroutinesApi
 @SmallTest
@@ -70,6 +73,28 @@
         }
 
     @Test
+    fun directionalDimensionPixelSize() =
+        testScope.runTest {
+            val resourceId = 1001
+            val pixelSize = 501
+            configurationRepository.setDimensionPixelSize(resourceId, pixelSize)
+
+            val config: Configuration = mock()
+            val dimensionPixelSize by
+                collectLastValue(
+                    underTest.directionalDimensionPixelSize(LayoutDirection.LTR, resourceId)
+                )
+
+            whenever(config.layoutDirection).thenReturn(LayoutDirection.LTR)
+            configurationRepository.onConfigurationChange(config)
+            assertThat(dimensionPixelSize).isEqualTo(pixelSize)
+
+            whenever(config.layoutDirection).thenReturn(LayoutDirection.RTL)
+            configurationRepository.onConfigurationChange(config)
+            assertThat(dimensionPixelSize).isEqualTo(-pixelSize)
+        }
+
+    @Test
     fun dimensionPixelSizes() =
         testScope.runTest {
             val resourceId1 = 1001
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/touch/CommunalTouchHandlerTest.java b/packages/SystemUI/tests/src/com/android/systemui/dreams/touch/CommunalTouchHandlerTest.java
similarity index 82%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/touch/CommunalTouchHandlerTest.java
rename to packages/SystemUI/tests/src/com/android/systemui/dreams/touch/CommunalTouchHandlerTest.java
index 7936ccc..c2c94a8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/touch/CommunalTouchHandlerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/dreams/touch/CommunalTouchHandlerTest.java
@@ -23,6 +23,9 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.graphics.Rect;
+import android.graphics.Region;
+import android.util.LayoutDirection;
 import android.view.GestureDetector;
 import android.view.MotionEvent;
 
@@ -68,10 +71,12 @@
         AtomicReference reference = new AtomicReference<>(null);
         when(mLifecycle.getInternalScopeRef()).thenReturn(reference);
         when(mLifecycle.getCurrentState()).thenReturn(Lifecycle.State.CREATED);
+
         mTouchHandler = new CommunalTouchHandler(
                 Optional.of(mCentralSurfaces),
                 INITIATION_WIDTH,
                 mKosmos.getCommunalInteractor(),
+                mKosmos.getConfigurationInteractor(),
                 mLifecycle
                 );
     }
@@ -127,4 +132,26 @@
                 .onScroll(motionEvent1, motionEvent2, 1, 1))
                 .isTrue();
     }
+
+    @Test
+    public void testTouchInitiationArea() {
+        final int right = 80;
+        final int bottom = 100;
+        final Rect bounds = new Rect(0, 0, right, bottom);
+
+        {
+            final Region region = new Region();
+            mTouchHandler.mLayoutDirectionCallback.accept(LayoutDirection.LTR);
+            mTouchHandler.getTouchInitiationRegion(bounds, region, null);
+            assertThat(region.getBounds()).isEqualTo(
+                    new Rect(right - INITIATION_WIDTH, 0, right, bottom));
+        }
+
+        {
+            final Region region = new Region();
+            mTouchHandler.mLayoutDirectionCallback.accept(LayoutDirection.RTL);
+            mTouchHandler.getTouchInitiationRegion(bounds, region, null);
+            assertThat(region.getBounds()).isEqualTo(new Rect(0, 0, INITIATION_WIDTH, bottom));
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagsClassicDebugTest.kt b/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagsClassicDebugTest.kt
index 6c4a730..3388a78 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagsClassicDebugTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/flags/FeatureFlagsClassicDebugTest.kt
@@ -68,7 +68,7 @@
     @Mock private lateinit var systemProperties: SystemPropertiesHelper
     @Mock private lateinit var resources: Resources
     @Mock private lateinit var restarter: Restarter
-    private val userTracker = FakeUserTracker()
+    private lateinit var userTracker: FakeUserTracker
     private val flagMap = mutableMapOf<String, Flag<*>>()
     private lateinit var broadcastReceiver: BroadcastReceiver
     private lateinit var clearCacheAction: Consumer<String>
@@ -82,6 +82,9 @@
         MockitoAnnotations.initMocks(this)
         flagMap.put(teamfoodableFlagA.name, teamfoodableFlagA)
         flagMap.put(releasedFlagB.name, releasedFlagB)
+
+        userTracker = FakeUserTracker(onCreateCurrentUserContext = { mockContext })
+
         mFeatureFlagsClassicDebug =
             FeatureFlagsClassicDebug(
                 flagManager,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
index 0043173..6b60740 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModelTest.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.keyboard.shortcut.ui.viewmodel
 
+import android.app.role.RoleManager
+import android.app.role.mockRoleManager
 import android.view.KeyEvent
 import android.view.KeyboardShortcutGroup
 import android.view.KeyboardShortcutInfo
@@ -45,7 +47,9 @@
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.model.sysUiState
+import com.android.systemui.settings.fakeUserTracker
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SHORTCUT_HELPER_SHOWING
+import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
@@ -77,12 +81,15 @@
     private val testScope = kosmos.testScope
     private val testHelper = kosmos.shortcutHelperTestHelper
     private val sysUiState = kosmos.sysUiState
+    private val fakeUserTracker = kosmos.fakeUserTracker
+    private val mockRoleManager = kosmos.mockRoleManager
     private val viewModel = kosmos.shortcutHelperViewModel
 
     @Before
     fun setUp() {
         fakeSystemSource.setGroups(TestShortcuts.systemGroups)
         fakeMultiTaskingSource.setGroups(TestShortcuts.multitaskingGroups)
+        fakeCurrentAppsSource.setGroups(TestShortcuts.currentAppGroups)
     }
 
     @Test
@@ -216,21 +223,21 @@
         }
 
     @Test
-    fun shortcutsUiState_featureActive_emitsActiveWithFirstCategorySelectedByDefault() =
+    fun shortcutsUiState_noCurrentAppCategory_defaultSelectedCategoryIsSystem() =
         testScope.runTest {
+            fakeCurrentAppsSource.setGroups(emptyList())
+
             val uiState by collectLastValue(viewModel.shortcutsUiState)
 
             testHelper.showFromActivity()
 
             val activeUiState = uiState as ShortcutsUiState.Active
-            assertThat(activeUiState.defaultSelectedCategory)
-                .isEqualTo(activeUiState.shortcutCategories.first().type)
+            assertThat(activeUiState.defaultSelectedCategory).isEqualTo(System)
         }
 
     @Test
-    fun shortcutsUiState_featureActive_emitsActiveWithCurrentAppsCategorySelectedWhenPresent() =
+    fun shortcutsUiState_currentAppCategoryPresent_currentAppIsDefaultSelected() =
         testScope.runTest {
-            fakeCurrentAppsSource.setGroups(TestShortcuts.currentAppGroups)
             val uiState by collectLastValue(viewModel.shortcutsUiState)
 
             testHelper.showFromActivity()
@@ -241,6 +248,24 @@
         }
 
     @Test
+    fun shortcutsUiState_currentAppIsLauncher_defaultSelectedCategoryIsSystem() =
+        testScope.runTest {
+            whenever(
+                    mockRoleManager.getRoleHoldersAsUser(
+                        RoleManager.ROLE_HOME,
+                        fakeUserTracker.userHandle
+                    )
+                )
+                .thenReturn(listOf(TestShortcuts.currentAppPackageName))
+            val uiState by collectLastValue(viewModel.shortcutsUiState)
+
+            testHelper.showFromActivity()
+
+            val activeUiState = uiState as ShortcutsUiState.Active
+            assertThat(activeUiState.defaultSelectedCategory).isEqualTo(System)
+        }
+
+    @Test
     fun shortcutsUiState_userTypedQuery_filtersMatchingShortcutLabels() =
         testScope.runTest {
             fakeSystemSource.setGroups(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt
index 506c5ae..29cd9a2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt
@@ -30,6 +30,7 @@
 import android.view.SurfaceControlViewHost
 import androidx.test.filters.SmallTest
 import com.android.internal.widget.LockPatternUtils
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.SystemUIAppComponentFactoryBase
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.DialogTransitionAnimator
@@ -96,7 +97,8 @@
     @Mock private lateinit var previewSurfacePackage: SurfaceControlViewHost.SurfacePackage
     @Mock private lateinit var launchAnimator: DialogTransitionAnimator
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
-    @Mock private lateinit var logger: KeyguardQuickAffordancesMetricsLogger
+    @Mock private lateinit var logger: KeyguardQuickAffordancesLogger
+    @Mock private lateinit var metricsLogger: KeyguardQuickAffordancesMetricsLogger
 
     private lateinit var dockManager: DockManagerFake
     private lateinit var biometricSettingsRepository: FakeBiometricSettingsRepository
@@ -199,6 +201,7 @@
                 repository = { quickAffordanceRepository },
                 launchAnimator = launchAnimator,
                 logger = logger,
+                metricsLogger = metricsLogger,
                 devicePolicyManager = devicePolicyManager,
                 dockManager = dockManager,
                 biometricSettingsRepository = biometricSettingsRepository,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
index e68a4a5..9de7528 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
@@ -70,10 +70,10 @@
 import android.view.RemoteAnimationTarget;
 import android.view.View;
 import android.view.ViewRootImpl;
-import android.view.WindowManager;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.internal.logging.InstanceId;
 import com.android.internal.logging.UiEventLogger;
@@ -167,7 +167,7 @@
     private @Mock BroadcastDispatcher mBroadcastDispatcher;
     private @Mock DismissCallbackRegistry mDismissCallbackRegistry;
     private @Mock DumpManager mDumpManager;
-    private @Mock WindowManager mWindowManager;
+    private @Mock ViewCaptureAwareWindowManager mWindowManager;
     private @Mock IActivityManager mActivityManager;
     private @Mock ConfigurationController mConfigurationController;
     private @Mock PowerManager mPowerManager;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt
index 7560a97..e3bdcd7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorParameterizedTest.kt
@@ -23,6 +23,7 @@
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.SmallTest
 import com.android.internal.widget.LockPatternUtils
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.ActivityTransitionAnimator
 import com.android.systemui.animation.DialogTransitionAnimator
@@ -232,7 +233,8 @@
     @Mock private lateinit var expandable: Expandable
     @Mock private lateinit var launchAnimator: DialogTransitionAnimator
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
-    @Mock private lateinit var logger: KeyguardQuickAffordancesMetricsLogger
+    @Mock private lateinit var logger: KeyguardQuickAffordancesLogger
+    @Mock private lateinit var metricsLogger: KeyguardQuickAffordancesMetricsLogger
 
     private lateinit var underTest: KeyguardQuickAffordanceInteractor
     private lateinit var testScope: TestScope
@@ -327,6 +329,7 @@
                 repository = { quickAffordanceRepository },
                 launchAnimator = launchAnimator,
                 logger = logger,
+                metricsLogger = metricsLogger,
                 devicePolicyManager = devicePolicyManager,
                 dockManager = dockManager,
                 biometricSettingsRepository = biometricSettingsRepository,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorSceneContainerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorSceneContainerTest.kt
index fd1bf54..591ce1a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorSceneContainerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorSceneContainerTest.kt
@@ -23,6 +23,7 @@
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.SmallTest
 import com.android.internal.widget.LockPatternUtils
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.ActivityTransitionAnimator
 import com.android.systemui.animation.DialogTransitionAnimator
@@ -232,7 +233,8 @@
     @Mock private lateinit var expandable: Expandable
     @Mock private lateinit var launchAnimator: DialogTransitionAnimator
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
-    @Mock private lateinit var logger: KeyguardQuickAffordancesMetricsLogger
+    @Mock private lateinit var logger: KeyguardQuickAffordancesLogger
+    @Mock private lateinit var metricsLogger: KeyguardQuickAffordancesMetricsLogger
 
     private lateinit var underTest: KeyguardQuickAffordanceInteractor
     private lateinit var testScope: TestScope
@@ -327,6 +329,7 @@
                 repository = { quickAffordanceRepository },
                 launchAnimator = launchAnimator,
                 logger = logger,
+                metricsLogger = metricsLogger,
                 devicePolicyManager = devicePolicyManager,
                 dockManager = dockManager,
                 biometricSettingsRepository = biometricSettingsRepository,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
index 3b96be4..fc7f693 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
@@ -23,6 +23,7 @@
 import androidx.test.filters.SmallTest
 import com.android.internal.logging.testing.UiEventLoggerFake
 import com.android.internal.widget.LockPatternUtils
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.animation.Expandable
@@ -99,7 +100,8 @@
     @Mock private lateinit var activityStarter: ActivityStarter
     @Mock private lateinit var launchAnimator: DialogTransitionAnimator
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
-    @Mock private lateinit var logger: KeyguardQuickAffordancesMetricsLogger
+    @Mock private lateinit var logger: KeyguardQuickAffordancesLogger
+    @Mock private lateinit var metricsLogger: KeyguardQuickAffordancesMetricsLogger
     @Mock private lateinit var broadcastDispatcher: BroadcastDispatcher
     @Mock private lateinit var accessibilityManager: AccessibilityManagerWrapper
 
@@ -237,6 +239,7 @@
                         repository = { quickAffordanceRepository },
                         launchAnimator = launchAnimator,
                         logger = logger,
+                        metricsLogger = metricsLogger,
                         devicePolicyManager = devicePolicyManager,
                         dockManager = dockManager,
                         biometricSettingsRepository = biometricSettingsRepository,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt
index e89abf6..77977f3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt
@@ -23,6 +23,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.internal.widget.LockPatternUtils
+import com.android.keyguard.logging.KeyguardQuickAffordancesLogger
 import com.android.systemui.Flags as AConfigFlags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.DialogTransitionAnimator
@@ -93,7 +94,8 @@
     @Mock private lateinit var lockPatternUtils: LockPatternUtils
     @Mock private lateinit var keyguardStateController: KeyguardStateController
     @Mock private lateinit var launchAnimator: DialogTransitionAnimator
-    @Mock private lateinit var logger: KeyguardQuickAffordancesMetricsLogger
+    @Mock private lateinit var logger: KeyguardQuickAffordancesLogger
+    @Mock private lateinit var metricsLogger: KeyguardQuickAffordancesMetricsLogger
     @Mock private lateinit var shadeInteractor: ShadeInteractor
     @Mock
     private lateinit var aodToLockscreenTransitionViewModel: AodToLockscreenTransitionViewModel
@@ -299,6 +301,7 @@
                         repository = { quickAffordanceRepository },
                         launchAnimator = launchAnimator,
                         logger = logger,
+                        metricsLogger = metricsLogger,
                         devicePolicyManager = devicePolicyManager,
                         dockManager = dockManager,
                         biometricSettingsRepository = biometricSettingsRepository,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerImplTest.java
index 196bbb9..413aa55 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/navigationbar/NavigationBarControllerImplTest.java
@@ -307,7 +307,7 @@
 
         mNavigationBarController.mIsLargeScreen = false;
         mNavigationBarController.mIsPhone = true;
-        assertFalse(mNavigationBarController.supportsTaskbar());
+        assertTrue(mNavigationBarController.supportsTaskbar());
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/footer/ui/viewmodel/FooterActionsViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/footer/ui/viewmodel/FooterActionsViewModelTest.kt
index 31652a5..f90e1e9 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/footer/ui/viewmodel/FooterActionsViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/footer/ui/viewmodel/FooterActionsViewModelTest.kt
@@ -47,6 +47,7 @@
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runTest
@@ -60,7 +61,8 @@
 @RunWith(AndroidJUnit4::class)
 @RunWithLooper
 class FooterActionsViewModelTest : SysuiTestCase() {
-    private val testScope = TestScope()
+    private val testDispatcher = StandardTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
     private lateinit var utils: FooterActionsTestUtils
 
     private val themedContext = ContextThemeWrapper(context, R.style.Theme_SystemUI_QuickSettings)
@@ -127,7 +129,7 @@
     fun userSwitcher() = runTest {
         val picture: Drawable = mock()
         val userInfoController = FakeUserInfoController(FakeInfo(picture = picture))
-        val settings = FakeGlobalSettings()
+        val settings = FakeGlobalSettings(testDispatcher)
         val userId = 42
         val userSwitcherControllerWrapper =
             MockUserSwitcherControllerWrapper(currentUserName = "foo")
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tileimpl/QSTileViewImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tileimpl/QSTileViewImplTest.kt
index 988769f..90ffaf1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tileimpl/QSTileViewImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tileimpl/QSTileViewImplTest.kt
@@ -25,6 +25,7 @@
 import android.view.ContextThemeWrapper
 import android.view.View
 import android.view.accessibility.AccessibilityNodeInfo
+import android.widget.Button
 import android.widget.TextView
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -381,6 +382,47 @@
     }
 
     @Test
+    fun testNonSwitchA11yClass_longClickActionHasCorrectLabel() {
+        val state =
+            QSTile.State().apply {
+                expandedAccessibilityClassName = Button::class.java.name
+                handlesLongClick = true
+            }
+        tileView.changeState(state)
+        val info = AccessibilityNodeInfo(tileView)
+        tileView.onInitializeAccessibilityNodeInfo(info)
+
+        assertThat(
+                info.actionList
+                    .find {
+                        it.id == AccessibilityNodeInfo.AccessibilityAction.ACTION_LONG_CLICK.id
+                    }
+                    ?.label
+            )
+            .isEqualTo(context.getString(R.string.accessibility_long_click_tile))
+    }
+
+    @Test
+    fun testNonSwitchA11yClass_disabledByPolicy_noLongClickAction() {
+        val state =
+            QSTile.State().apply {
+                expandedAccessibilityClassName = Button::class.java.name
+                handlesLongClick = true
+                disabledByPolicy = true
+            }
+        tileView.changeState(state)
+        val info = AccessibilityNodeInfo(tileView)
+        tileView.onInitializeAccessibilityNodeInfo(info)
+
+        assertThat(
+                info.actionList.find {
+                    it.id == AccessibilityNodeInfo.AccessibilityAction.ACTION_LONG_CLICK.id
+                }
+            )
+            .isNull()
+    }
+
+    @Test
     fun onStateChange_longPressEffectActive_withInvalidDuration_doesNotInitializeEffect() {
         val state = QSTile.State() // A state that handles longPress
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
index 27b6ea6..a5de7cd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
@@ -27,7 +27,6 @@
 import com.android.internal.logging.MetricsLogger
 import com.android.settingslib.notification.data.repository.FakeZenModeRepository
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.classifier.FalsingManagerFake
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.plugins.statusbar.StatusBarStateController
@@ -47,10 +46,9 @@
 import com.android.systemui.util.settings.FakeSettings
 import com.android.systemui.util.settings.SecureSettings
 import com.google.common.truth.Truth.assertThat
-import kotlin.coroutines.EmptyCoroutineContext
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.After
@@ -82,13 +80,14 @@
 
     @Mock private lateinit var qsTileConfigProvider: QSTileConfigProvider
 
-    @Mock private lateinit var dialogTransitionAnimator: DialogTransitionAnimator
-
     @Mock private lateinit var dialogDelegate: ModesDialogDelegate
 
+    private val testDispatcher = UnconfinedTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
+
     private val inputHandler = FakeQSTileIntentUserInputHandler()
     private val zenModeRepository = FakeZenModeRepository()
-    private val tileDataInteractor = ModesTileDataInteractor(zenModeRepository)
+    private val tileDataInteractor = ModesTileDataInteractor(zenModeRepository, testDispatcher)
     private val mapper =
         ModesTileMapper(
             context.orCreateTestableResources
@@ -100,9 +99,6 @@
             context.theme,
         )
 
-    private val testDispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(testDispatcher)
-
     private lateinit var userActionInteractor: ModesTileUserActionInteractor
     private lateinit var secureSettings: SecureSettings
     private lateinit var testableLooper: TestableLooper
@@ -131,9 +127,7 @@
 
         userActionInteractor =
             ModesTileUserActionInteractor(
-                EmptyCoroutineContext,
                 inputHandler,
-                dialogTransitionAnimator,
                 dialogDelegate,
             )
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/SaveImageInBackgroundTaskTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/SaveImageInBackgroundTaskTest.kt
index 24e8b18..5e07aef 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/SaveImageInBackgroundTaskTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/SaveImageInBackgroundTaskTest.kt
@@ -228,7 +228,7 @@
             )
         val quickSharePendingIntent =
             quickShareAction.actionIntent.intent.extras!!.getParcelable(
-                ScreenshotController.EXTRA_ACTION_INTENT,
+                SmartActionsReceiver.EXTRA_ACTION_INTENT,
                 PendingIntent::class.java
             )
 
@@ -266,7 +266,7 @@
         assertEquals(
             immutablePendingIntent,
             quickShareAction.actionIntent.intent.extras!!.getParcelable(
-                ScreenshotController.EXTRA_ACTION_INTENT,
+                SmartActionsReceiver.EXTRA_ACTION_INTENT,
                 PendingIntent::class.java
             )
         )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java
index 471fdc0..9dc5cfe 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java
@@ -16,8 +16,8 @@
 
 package com.android.systemui.screenshot;
 
-import static com.android.systemui.screenshot.ScreenshotController.EXTRA_ACTION_TYPE;
-import static com.android.systemui.screenshot.ScreenshotController.EXTRA_ID;
+import static com.android.systemui.screenshot.SmartActionsReceiver.EXTRA_ACTION_TYPE;
+import static com.android.systemui.screenshot.SmartActionsReceiver.EXTRA_ID;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -57,7 +57,7 @@
         MockitoAnnotations.initMocks(this);
         mSmartActionsReceiver = new SmartActionsReceiver(mMockScreenshotSmartActions);
         mIntent = new Intent(mContext, SmartActionsReceiver.class)
-                .putExtra(ScreenshotController.EXTRA_ACTION_INTENT, mMockPendingIntent);
+                .putExtra(SmartActionsReceiver.EXTRA_ACTION_INTENT, mMockPendingIntent);
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
index 8d3a29a..a295981 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
@@ -42,10 +42,10 @@
 @SmallTest
 class TakeScreenshotExecutorTest : SysuiTestCase() {
 
-    private val controller = mock<ScreenshotController>()
+    private val controller = mock<LegacyScreenshotController>()
     private val notificationsController0 = mock<ScreenshotNotificationsController>()
     private val notificationsController1 = mock<ScreenshotNotificationsController>()
-    private val controllerFactory = mock<ScreenshotController.Factory>()
+    private val controllerFactory = mock<InteractiveScreenshotHandler.Factory>()
     private val callback = mock<TakeScreenshotService.RequestCallback>()
     private val notificationControllerFactory = mock<ScreenshotNotificationsController.Factory>()
 
@@ -287,7 +287,7 @@
     fun onCloseSystemDialogsReceived_controllerHasPendingTransitions() =
         testScope.runTest {
             setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
-            whenever(controller.isPendingSharedTransition).thenReturn(true)
+            whenever(controller.isPendingSharedTransition()).thenReturn(true)
             val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
index 9986205..a8d5008 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
@@ -17,6 +17,7 @@
 package com.android.systemui.screenshot.appclips;
 
 import static android.app.Activity.RESULT_OK;
+import static android.app.ActivityManager.RunningTaskInfo;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD;
 
 import static com.android.systemui.screenshot.appclips.AppClipsEvent.SCREENSHOT_FOR_NOTE_ACCEPTED;
@@ -32,7 +33,6 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.app.ActivityTaskManager.RootTaskInfo;
 import android.app.IActivityTaskManager;
 import android.app.assist.AssistContent;
 import android.content.ComponentName;
@@ -103,7 +103,7 @@
     private static final String BACKLINKS_TASK_APP_NAME = "Backlinks app";
     private static final String BACKLINKS_TASK_PACKAGE_NAME = "backlinksTaskPackageName";
 
-    private static final RootTaskInfo TASK_THAT_SUPPORTS_BACKLINKS =
+    private static final RunningTaskInfo TASK_THAT_SUPPORTS_BACKLINKS =
             createTaskInfoForBacklinksTask();
     private static final AssistContent ASSIST_CONTENT_FOR_BACKLINKS_TASK =
             createAssistContentForBacklinksTask();
@@ -233,6 +233,10 @@
         assertThat(backlinksData.getText().toString()).isEqualTo(BACKLINKS_TASK_APP_NAME);
         assertThat(backlinksData.getCompoundDrawablesRelative()[0]).isEqualTo(FAKE_DRAWABLE);
 
+        // Verify dropdown icon is not shown and there are no click listeners on text view.
+        assertThat(backlinksData.getCompoundDrawablesRelative()[2]).isNull();
+        assertThat(backlinksData.hasOnClickListeners()).isFalse();
+
         CheckBox backlinksIncludeData = mActivity.findViewById(R.id.backlinks_include_data);
         assertThat(backlinksIncludeData.getVisibility()).isEqualTo(View.VISIBLE);
         assertThat(backlinksIncludeData.getText().toString())
@@ -258,20 +262,71 @@
         assertThat(backlinksData.getVisibility()).isEqualTo(View.GONE);
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_APP_CLIPS_BACKLINKS)
+    public void appClipsLaunched_backlinks_multipleBacklinksAvailable_defaultShown()
+            throws RemoteException {
+        // Set up mocking for multiple backlinks.
+        ResolveInfo resolveInfo1 = createBacklinksTaskResolveInfo();
+
+        int taskId2 = BACKLINKS_TASK_ID + 2;
+        String package2 = BACKLINKS_TASK_PACKAGE_NAME + 2;
+        String appName2 = BACKLINKS_TASK_APP_NAME + 2;
+
+        ResolveInfo resolveInfo2 = createBacklinksTaskResolveInfo();
+        ActivityInfo activityInfo2 = resolveInfo2.activityInfo;
+        activityInfo2.name = appName2;
+        activityInfo2.packageName = package2;
+        activityInfo2.applicationInfo.packageName = package2;
+        RunningTaskInfo runningTaskInfo2 = createTaskInfoForBacklinksTask();
+        runningTaskInfo2.taskId = taskId2;
+        runningTaskInfo2.topActivity = new ComponentName(package2, "backlinksClass");
+        runningTaskInfo2.topActivityInfo = resolveInfo2.activityInfo;
+        runningTaskInfo2.baseIntent = new Intent().setComponent(runningTaskInfo2.topActivity);
+
+        when(mAtmService.getTasks(eq(Integer.MAX_VALUE), eq(false), eq(false),
+                mDisplayIdCaptor.capture()))
+                .thenReturn(List.of(TASK_THAT_SUPPORTS_BACKLINKS, runningTaskInfo2));
+        when(mPackageManager.resolveActivity(any(Intent.class), anyInt())).thenReturn(resolveInfo1,
+                resolveInfo1, resolveInfo1, resolveInfo2, resolveInfo2, resolveInfo2);
+        when(mPackageManager.loadItemIcon(any(), any())).thenReturn(FAKE_DRAWABLE);
+
+        // Using same AssistContent data for both tasks.
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, BACKLINKS_TASK_ID);
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, taskId2);
+
+        // Mocking complete, trigger backlinks.
+        launchActivity();
+        waitForIdleSync();
+
+        // Verify default backlink shown to user and text view has on click listener.
+        TextView backlinksData = mActivity.findViewById(R.id.backlinks_data);
+        assertThat(backlinksData.getText().toString()).isEqualTo(BACKLINKS_TASK_APP_NAME);
+        assertThat(backlinksData.hasOnClickListeners()).isTrue();
+
+        // Verify dropdown icon is not null.
+        assertThat(backlinksData.getCompoundDrawablesRelative()[2]).isNotNull();
+    }
+
     private void setUpMocksForBacklinks() throws RemoteException {
-        when(mAtmService.getAllRootTaskInfosOnDisplay(mDisplayIdCaptor.capture()))
+        when(mAtmService.getTasks(eq(Integer.MAX_VALUE), eq(false), eq(false),
+                mDisplayIdCaptor.capture()))
                 .thenReturn(List.of(TASK_THAT_SUPPORTS_BACKLINKS));
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(ASSIST_CONTENT_FOR_BACKLINKS_TASK);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, BACKLINKS_TASK_ID);
         when(mPackageManager
                 .resolveActivity(any(Intent.class), anyInt()))
                 .thenReturn(createBacklinksTaskResolveInfo());
         when(mPackageManager.loadItemIcon(any(), any())).thenReturn(FAKE_DRAWABLE);
     }
 
+    private void mockForAssistContent(AssistContent expected, int taskId) {
+        doAnswer(invocation -> {
+            AssistContentRequester.Callback callback = invocation.getArgument(1);
+            callback.onAssistContentAvailable(expected);
+            return null;
+        }).when(mAssistContentRequester).requestAssistContent(eq(taskId), any());
+    }
+
     private void launchActivity() {
         launchActivity(createResultReceiver(FAKE_CONSUMER));
     }
@@ -319,8 +374,8 @@
         return resolveInfo;
     }
 
-    private static RootTaskInfo createTaskInfoForBacklinksTask() {
-        RootTaskInfo taskInfo = new RootTaskInfo();
+    private static RunningTaskInfo createTaskInfoForBacklinksTask() {
+        RunningTaskInfo taskInfo = new RunningTaskInfo();
         taskInfo.taskId = BACKLINKS_TASK_ID;
         taskInfo.isVisible = true;
         taskInfo.isRunning = true;
@@ -328,7 +383,6 @@
         taskInfo.topActivity = new ComponentName(BACKLINKS_TASK_PACKAGE_NAME, "backlinksClass");
         taskInfo.topActivityInfo = createBacklinksTaskResolveInfo().activityInfo;
         taskInfo.baseIntent = new Intent().setComponent(taskInfo.topActivity);
-        taskInfo.childTaskIds = new int[]{BACKLINKS_TASK_ID + 1};
         taskInfo.configuration.windowConfiguration.setActivityType(ACTIVITY_TYPE_STANDARD);
         return taskInfo;
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java
index 193d29c..178547e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsViewModelTest.java
@@ -37,7 +37,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import android.app.ActivityTaskManager.RootTaskInfo;
+import android.app.ActivityManager.RunningTaskInfo;
 import android.app.IActivityTaskManager;
 import android.app.assist.AssistContent;
 import android.content.ClipData;
@@ -107,7 +107,7 @@
         mPackageManagerIntentCaptor = ArgumentCaptor.forClass(Intent.class);
 
         // Set up mocking for backlinks.
-        when(mAtmService.getAllRootTaskInfosOnDisplay(DEFAULT_DISPLAY))
+        when(mAtmService.getTasks(Integer.MAX_VALUE, false, false, DEFAULT_DISPLAY))
                 .thenReturn(List.of(createTaskInfoForBacklinksTask()));
         when(mPackageManager.resolveActivity(mPackageManagerIntentCaptor.capture(), anyInt()))
                 .thenReturn(createBacklinksTaskResolveInfo());
@@ -190,11 +190,7 @@
         Uri expectedUri = Uri.parse("https://developers.android.com");
         AssistContent contentWithUri = new AssistContent();
         contentWithUri.setWebUri(expectedUri);
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithUri);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(contentWithUri, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -203,7 +199,7 @@
         assertThat(queriedIntent.getData()).isEqualTo(expectedUri);
         assertThat(queriedIntent.getAction()).isEqualTo(ACTION_VIEW);
 
-        InternalBacklinksData result = mViewModel.getBacklinksLiveData().getValue();
+        InternalBacklinksData result = mViewModel.mSelectedBacklinksLiveData.getValue();
         assertThat(result.getAppIcon()).isEqualTo(FAKE_DRAWABLE);
         ClipData clipData = result.getClipData();
         ClipDescription resultDescription = clipData.getDescription();
@@ -211,6 +207,8 @@
         assertThat(resultDescription.getMimeType(0)).isEqualTo(MIMETYPE_TEXT_URILIST);
         assertThat(clipData.getItemCount()).isEqualTo(1);
         assertThat(clipData.getItemAt(0).getUri()).isEqualTo(expectedUri);
+
+        assertThat(result).isEqualTo(mViewModel.getBacklinksLiveData().getValue().get(0));
     }
 
     @Test
@@ -218,12 +216,8 @@
         Uri expectedUri = Uri.parse("https://developers.android.com");
         AssistContent contentWithUri = new AssistContent();
         contentWithUri.setWebUri(expectedUri);
+        mockForAssistContent(contentWithUri, BACKLINKS_TASK_ID);
         resetPackageManagerMockingForUsingFallbackBacklinks();
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithUri);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -236,11 +230,7 @@
         Intent expectedIntent = new Intent().setPackage(BACKLINKS_TASK_PACKAGE_NAME);
         AssistContent contentWithAppProvidedIntent = new AssistContent();
         contentWithAppProvidedIntent.setIntent(expectedIntent);
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithAppProvidedIntent);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(contentWithAppProvidedIntent, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -248,7 +238,7 @@
         Intent queriedIntent = mPackageManagerIntentCaptor.getValue();
         assertThat(queriedIntent.getPackage()).isEqualTo(expectedIntent.getPackage());
 
-        InternalBacklinksData result = mViewModel.getBacklinksLiveData().getValue();
+        InternalBacklinksData result = mViewModel.mSelectedBacklinksLiveData.getValue();
         assertThat(result.getAppIcon()).isEqualTo(FAKE_DRAWABLE);
         ClipData clipData = result.getClipData();
         ClipDescription resultDescription = clipData.getDescription();
@@ -263,12 +253,8 @@
         Intent expectedIntent = new Intent().setPackage(BACKLINKS_TASK_PACKAGE_NAME);
         AssistContent contentWithAppProvidedIntent = new AssistContent();
         contentWithAppProvidedIntent.setIntent(expectedIntent);
+        mockForAssistContent(contentWithAppProvidedIntent, BACKLINKS_TASK_ID);
         resetPackageManagerMockingForUsingFallbackBacklinks();
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(contentWithAppProvidedIntent);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -278,11 +264,7 @@
 
     @Test
     public void triggerBacklinks_shouldUpdateBacklinks_withMainLauncherIntent() {
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(EMPTY_ASSIST_CONTENT);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(EMPTY_ASSIST_CONTENT, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
@@ -298,15 +280,12 @@
     @Test
     public void triggerBacklinks_withNonResolvableMainLauncherIntent_noBacklinksAvailable() {
         reset(mPackageManager);
-        doAnswer(invocation -> {
-            AssistContentRequester.Callback callback = invocation.getArgument(1);
-            callback.onAssistContentAvailable(EMPTY_ASSIST_CONTENT);
-            return null;
-        }).when(mAssistContentRequester).requestAssistContent(eq(BACKLINKS_TASK_ID), any());
+        mockForAssistContent(EMPTY_ASSIST_CONTENT, BACKLINKS_TASK_ID);
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
 
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue()).isNull();
         assertThat(mViewModel.getBacklinksLiveData().getValue()).isNull();
     }
 
@@ -314,14 +293,15 @@
     public void triggerBacklinks_nonStandardActivityIgnored_noBacklinkAvailable()
             throws RemoteException {
         reset(mAtmService);
-        RootTaskInfo taskInfo = createTaskInfoForBacklinksTask();
+        RunningTaskInfo taskInfo = createTaskInfoForBacklinksTask();
         taskInfo.configuration.windowConfiguration.setActivityType(ACTIVITY_TYPE_HOME);
-        when(mAtmService.getAllRootTaskInfosOnDisplay(DEFAULT_DISPLAY))
+        when(mAtmService.getTasks(Integer.MAX_VALUE, false, false, DEFAULT_DISPLAY))
                 .thenReturn(List.of(taskInfo));
 
         mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
         waitForIdleSync();
 
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue()).isNull();
         assertThat(mViewModel.getBacklinksLiveData().getValue()).isNull();
     }
 
@@ -330,9 +310,68 @@
         mViewModel.triggerBacklinks(Set.of(BACKLINKS_TASK_ID), DEFAULT_DISPLAY);
         waitForIdleSync();
 
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue()).isNull();
         assertThat(mViewModel.getBacklinksLiveData().getValue()).isNull();
     }
 
+    @Test
+    public void triggerBacklinks_multipleAppsOnScreen_multipleBacklinksAvailable()
+            throws RemoteException {
+        // Set up mocking for multiple backlinks.
+        reset(mAtmService, mPackageManager);
+        RunningTaskInfo runningTaskInfo1 = createTaskInfoForBacklinksTask();
+        ResolveInfo resolveInfo1 = createBacklinksTaskResolveInfo();
+
+        int taskId2 = BACKLINKS_TASK_ID + 2;
+        String package2 = BACKLINKS_TASK_PACKAGE_NAME + 2;
+        String appName2 = BACKLINKS_TASK_APP_NAME + 2;
+
+        ResolveInfo resolveInfo2 = createBacklinksTaskResolveInfo();
+        ActivityInfo activityInfo2 = resolveInfo2.activityInfo;
+        activityInfo2.name = appName2;
+        activityInfo2.packageName = package2;
+        activityInfo2.applicationInfo.packageName = package2;
+        RunningTaskInfo runningTaskInfo2 = createTaskInfoForBacklinksTask();
+        runningTaskInfo2.taskId = taskId2;
+        runningTaskInfo2.topActivity = new ComponentName(package2, "backlinksClass");
+        runningTaskInfo2.topActivityInfo = resolveInfo2.activityInfo;
+        runningTaskInfo2.baseIntent = new Intent().setComponent(runningTaskInfo2.topActivity);
+
+        // For each task, the logic queries PM 3 times, twice for verifying if an app can be
+        // launched via launcher and once with the data provided in backlink intent.
+        when(mPackageManager.resolveActivity(any(), anyInt())).thenReturn(resolveInfo1,
+                resolveInfo1, resolveInfo1, resolveInfo2, resolveInfo2, resolveInfo2);
+        when(mPackageManager.loadItemIcon(any(), any())).thenReturn(FAKE_DRAWABLE);
+        when(mAtmService.getTasks(Integer.MAX_VALUE, false, false, DEFAULT_DISPLAY))
+                .thenReturn(List.of(runningTaskInfo1, runningTaskInfo2));
+
+        // Using app provided web uri for the first backlink.
+        Uri expectedUri = Uri.parse("https://developers.android.com");
+        AssistContent contentWithUri = new AssistContent();
+        contentWithUri.setWebUri(expectedUri);
+        mockForAssistContent(contentWithUri, BACKLINKS_TASK_ID);
+
+        // Using app provided intent for the second backlink.
+        Intent expectedIntent = new Intent().setPackage(package2);
+        AssistContent contentWithAppProvidedIntent = new AssistContent();
+        contentWithAppProvidedIntent.setIntent(expectedIntent);
+        mockForAssistContent(contentWithAppProvidedIntent, taskId2);
+
+        // Set up complete, trigger the backlinks action.
+        mViewModel.triggerBacklinks(Collections.emptySet(), DEFAULT_DISPLAY);
+        waitForIdleSync();
+
+        // Verify two backlinks are received and the first backlink is set as default selected.
+        assertThat(mViewModel.mSelectedBacklinksLiveData.getValue().getClipData().getItemAt(
+                0).getUri()).isEqualTo(expectedUri);
+        List<InternalBacklinksData> actualBacklinks = mViewModel.getBacklinksLiveData().getValue();
+        assertThat(actualBacklinks).hasSize(2);
+        assertThat(actualBacklinks.get(0).getClipData().getItemAt(0).getUri())
+                .isEqualTo(expectedUri);
+        assertThat(actualBacklinks.get(1).getClipData().getItemAt(0).getIntent())
+                .isEqualTo(expectedIntent);
+    }
+
     private void resetPackageManagerMockingForUsingFallbackBacklinks() {
         ResolveInfo backlinksTaskResolveInfo = createBacklinksTaskResolveInfo();
         reset(mPackageManager);
@@ -350,7 +389,7 @@
     }
 
     private void verifyMainLauncherBacklinksIntent() {
-        InternalBacklinksData result = mViewModel.getBacklinksLiveData().getValue();
+        InternalBacklinksData result = mViewModel.mSelectedBacklinksLiveData.getValue();
         assertThat(result.getAppIcon()).isEqualTo(FAKE_DRAWABLE);
 
         ClipData clipData = result.getClipData();
@@ -368,6 +407,14 @@
                 new ComponentName(BACKLINKS_TASK_PACKAGE_NAME, BACKLINKS_TASK_APP_NAME));
     }
 
+    private void mockForAssistContent(AssistContent expected, int taskId) {
+        doAnswer(invocation -> {
+            AssistContentRequester.Callback callback = invocation.getArgument(1);
+            callback.onAssistContentAvailable(expected);
+            return null;
+        }).when(mAssistContentRequester).requestAssistContent(eq(taskId), any());
+    }
+
     private static ResolveInfo createBacklinksTaskResolveInfo() {
         ActivityInfo activityInfo = new ActivityInfo();
         activityInfo.applicationInfo = new ApplicationInfo();
@@ -379,8 +426,8 @@
         return resolveInfo;
     }
 
-    private static RootTaskInfo createTaskInfoForBacklinksTask() {
-        RootTaskInfo taskInfo = new RootTaskInfo();
+    private static RunningTaskInfo createTaskInfoForBacklinksTask() {
+        RunningTaskInfo taskInfo = new RunningTaskInfo();
         taskInfo.taskId = BACKLINKS_TASK_ID;
         taskInfo.isVisible = true;
         taskInfo.isRunning = true;
@@ -388,7 +435,6 @@
         taskInfo.topActivity = new ComponentName(BACKLINKS_TASK_PACKAGE_NAME, "backlinksClass");
         taskInfo.topActivityInfo = createBacklinksTaskResolveInfo().activityInfo;
         taskInfo.baseIntent = new Intent().setComponent(taskInfo.topActivity);
-        taskInfo.childTaskIds = new int[]{BACKLINKS_TASK_ID + 1};
         taskInfo.configuration.windowConfiguration.setActivityType(ACTIVITY_TYPE_STANDARD);
         return taskInfo;
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt
index 967df39..5de31d8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt
@@ -429,6 +429,7 @@
     fun gestureExclusionZone_setAfterInit() =
         with(kosmos) {
             testScope.runTest {
+                whenever(containerView.layoutDirection).thenReturn(View.LAYOUT_DIRECTION_LTR)
                 goToScene(CommunalScenes.Communal)
 
                 assertThat(containerView.systemGestureExclusionRects)
@@ -450,10 +451,37 @@
         }
 
     @Test
+    @DisableFlags(FLAG_GLANCEABLE_HUB_BACK_GESTURE)
+    fun gestureExclusionZone_setAfterInit_rtl() =
+        with(kosmos) {
+            testScope.runTest {
+                whenever(containerView.layoutDirection).thenReturn(View.LAYOUT_DIRECTION_RTL)
+                goToScene(CommunalScenes.Communal)
+
+                assertThat(containerView.systemGestureExclusionRects)
+                    .containsExactly(
+                        Rect(
+                            /* left= */ 0,
+                            /* top= */ TOP_SWIPE_REGION_WIDTH,
+                            /* right= */ CONTAINER_WIDTH,
+                            /* bottom= */ CONTAINER_HEIGHT - BOTTOM_SWIPE_REGION_WIDTH
+                        ),
+                        Rect(
+                            /* left= */ 0,
+                            /* top= */ 0,
+                            /* right= */ CONTAINER_WIDTH,
+                            /* bottom= */ CONTAINER_HEIGHT
+                        )
+                    )
+            }
+        }
+
+    @Test
     @EnableFlags(FLAG_GLANCEABLE_HUB_BACK_GESTURE)
     fun gestureExclusionZone_setAfterInit_backGestureEnabled() =
         with(kosmos) {
             testScope.runTest {
+                whenever(containerView.layoutDirection).thenReturn(View.LAYOUT_DIRECTION_LTR)
                 goToScene(CommunalScenes.Communal)
 
                 assertThat(containerView.systemGestureExclusionRects)
@@ -475,6 +503,32 @@
         }
 
     @Test
+    @EnableFlags(FLAG_GLANCEABLE_HUB_BACK_GESTURE)
+    fun gestureExclusionZone_setAfterInit_backGestureEnabled_rtl() =
+        with(kosmos) {
+            testScope.runTest {
+                whenever(containerView.layoutDirection).thenReturn(View.LAYOUT_DIRECTION_RTL)
+                goToScene(CommunalScenes.Communal)
+
+                assertThat(containerView.systemGestureExclusionRects)
+                    .containsExactly(
+                        Rect(
+                            /* left= */ FAKE_INSETS.left,
+                            /* top= */ TOP_SWIPE_REGION_WIDTH,
+                            /* right= */ CONTAINER_WIDTH - FAKE_INSETS.right,
+                            /* bottom= */ CONTAINER_HEIGHT - BOTTOM_SWIPE_REGION_WIDTH
+                        ),
+                        Rect(
+                            /* left= */ FAKE_INSETS.left,
+                            /* top= */ 0,
+                            /* right= */ CONTAINER_WIDTH,
+                            /* bottom= */ CONTAINER_HEIGHT
+                        )
+                    )
+            }
+        }
+
+    @Test
     fun gestureExclusionZone_unsetWhenShadeOpen() =
         with(kosmos) {
             testScope.runTest {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
index 06a883c..b7ce336 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
@@ -44,7 +44,6 @@
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.fragments.FragmentHostManager;
-import com.android.systemui.keyguard.data.repository.FakeCommandQueue;
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
@@ -190,7 +189,6 @@
                 mKosmos.getKeyguardTransitionInteractor();
         KeyguardInteractor keyguardInteractor = new KeyguardInteractor(
                 mKeyguardRepository,
-                new FakeCommandQueue(),
                 powerInteractor,
                 new FakeKeyguardBouncerRepository(),
                 new ConfigurationInteractor(configurationRepository),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt
index 69cc9d5..30326a57 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt
@@ -22,6 +22,7 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.shared.settings.data.repository.FakeSecureSettingsRepository
+import com.android.systemui.shared.settings.data.repository.FakeSystemSettingsRepository
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
@@ -38,18 +39,21 @@
 
     private lateinit var testScope: TestScope
     private lateinit var secureSettingsRepository: FakeSecureSettingsRepository
+    private lateinit var systemSettingsRepository: FakeSystemSettingsRepository
 
     @Before
     fun setUp() {
         val testDispatcher = StandardTestDispatcher()
         testScope = TestScope(testDispatcher)
         secureSettingsRepository = FakeSecureSettingsRepository()
+        systemSettingsRepository = FakeSystemSettingsRepository()
 
         underTest =
             NotificationSettingsRepository(
                 scope = testScope.backgroundScope,
                 backgroundDispatcher = testDispatcher,
                 secureSettingsRepository = secureSettingsRepository,
+                systemSettingsRepository = systemSettingsRepository,
             )
     }
 
@@ -100,4 +104,22 @@
             )
             assertThat(historyEnabled).isEqualTo(false)
         }
+
+    @Test
+    fun testGetIsCooldownEnabled() =
+        testScope.runTest {
+            val cooldownEnabled by collectLastValue(underTest.isCooldownEnabled)
+
+            systemSettingsRepository.setInt(
+                name = Settings.System.NOTIFICATION_COOLDOWN_ENABLED,
+                value = 1,
+            )
+            assertThat(cooldownEnabled).isEqualTo(true)
+
+            systemSettingsRepository.setInt(
+                name = Settings.System.NOTIFICATION_COOLDOWN_ENABLED,
+                value = 0,
+            )
+            assertThat(cooldownEnabled).isEqualTo(false)
+        }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt
index a7f36c3..f9509d2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderImplTest.kt
@@ -66,7 +66,8 @@
             packageManager,
             Optional.of(bubbles),
             context,
-            notificationManager
+            notificationManager,
+            settingsInteractor
         )
     }
 
@@ -101,7 +102,7 @@
         whenever(avalancheProvider.startTime).thenReturn(whenAgo(10))
 
         val avalancheSuppressor = AvalancheSuppressor(
-            avalancheProvider, systemClock, systemSettings, packageManager,
+            avalancheProvider, systemClock, settingsInteractor, packageManager,
             uiEventLogger, context, notificationManager
         )
         avalancheSuppressor.hasSeenEdu = false
@@ -125,7 +126,7 @@
         whenever(avalancheProvider.startTime).thenReturn(whenAgo(10))
 
         val avalancheSuppressor = AvalancheSuppressor(
-            avalancheProvider, systemClock, systemSettings, packageManager,
+            avalancheProvider, systemClock, settingsInteractor, packageManager,
             uiEventLogger, context, notificationManager
         )
         avalancheSuppressor.hasSeenEdu = true
@@ -147,7 +148,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
@@ -167,7 +168,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
@@ -187,7 +188,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
@@ -205,7 +206,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
@@ -223,7 +224,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
@@ -241,7 +242,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
@@ -259,7 +260,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             assertFsiNotSuppressed()
@@ -271,7 +272,7 @@
         avalancheProvider.startTime = whenAgo(10)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
@@ -300,7 +301,7 @@
         setAllowedEmergencyPkg(true)
 
         withFilter(
-            AvalancheSuppressor(avalancheProvider, systemClock, systemSettings, packageManager,
+            AvalancheSuppressor(avalancheProvider, systemClock, settingsInteractor, packageManager,
                     uiEventLogger, context, notificationManager)
         ) {
             ensurePeekState()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
index d5ab62b..9d3d9c1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
@@ -61,6 +61,7 @@
 import com.android.systemui.log.core.LogLevel
 import com.android.systemui.res.R
 import com.android.systemui.settings.FakeUserTracker
+import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor
 import com.android.systemui.statusbar.FakeStatusBarStateController
 import com.android.systemui.statusbar.NotificationEntryHelper.modifyRanking
 import com.android.systemui.statusbar.StatusBarState.KEYGUARD
@@ -88,6 +89,7 @@
 import com.android.wm.shell.bubbles.Bubbles
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertTrue
+import kotlinx.coroutines.flow.MutableStateFlow
 import org.junit.Assert.assertEquals
 import org.junit.Before
 import org.junit.Test
@@ -134,6 +136,7 @@
     protected val avalancheProvider: AvalancheProvider = mock()
     protected val bubbles: Bubbles = mock()
     lateinit var systemSettings: SystemSettings
+    protected val settingsInteractor: NotificationSettingsInteractor = mock()
     protected val packageManager: PackageManager = mock()
     protected val notificationManager: NotificationManager = mock()
     protected abstract val provider: VisualInterruptionDecisionProvider
@@ -164,7 +167,7 @@
         userTracker.set(listOf(user), /* currentUserIndex = */ 0)
         systemSettings = FakeSettings()
         whenever(bubbles.canShowBubbleNotification()).thenReturn(true)
-
+        whenever(settingsInteractor.isCooldownEnabled).thenReturn(MutableStateFlow(true))
         provider.start()
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
index f4cebd7..7fd9c9f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor
 import com.android.systemui.statusbar.notification.NotifPipelineFlags
 import com.android.systemui.statusbar.policy.BatteryController
 import com.android.systemui.statusbar.policy.DeviceProvisionedController
@@ -61,7 +62,8 @@
         packageManager: PackageManager,
         bubbles: Optional<Bubbles>,
         context: Context,
-        notificationManager: NotificationManager
+        notificationManager: NotificationManager,
+        settingsInteractor: NotificationSettingsInteractor
     ): VisualInterruptionDecisionProvider {
         return if (VisualInterruptionRefactor.isEnabled) {
             VisualInterruptionDecisionProviderImpl(
@@ -85,7 +87,8 @@
                 packageManager,
                 bubbles,
                 context,
-                notificationManager
+                notificationManager,
+                settingsInteractor
             )
         } else {
             NotificationInterruptStateProviderWrapper(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
index c36a046..3df4a67 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
@@ -20,6 +20,7 @@
 import static com.android.systemui.log.LogBufferHelperKt.logcatLogBuffer;
 import static com.android.systemui.statusbar.StatusBarState.KEYGUARD;
 import static com.android.systemui.statusbar.StatusBarState.SHADE;
+import static com.android.systemui.statusbar.notification.row.NotificationRowContentBinder.FLAG_CONTENT_VIEW_ALL;
 import static com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout.ROWS_ALL;
 
 import static kotlinx.coroutines.flow.FlowKt.emptyFlow;
@@ -96,9 +97,12 @@
 import com.android.systemui.statusbar.notification.init.NotificationsController;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
+import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.shared.GroupHunAnimationFix;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController.NotificationPanelEvent;
 import com.android.systemui.statusbar.notification.stack.NotificationSwipeHelper.NotificationCallback;
 import com.android.systemui.statusbar.notification.stack.ui.viewbinder.NotificationListViewBinder;
+import com.android.systemui.statusbar.phone.HeadsUpAppearanceController;
 import com.android.systemui.statusbar.phone.HeadsUpTouchHelper;
 import com.android.systemui.statusbar.phone.KeyguardBypassController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
@@ -192,6 +196,8 @@
 
     private NotificationStackScrollLayoutController mController;
 
+    private NotificationTestHelper mNotificationTestHelper;
+
     @Before
     public void setUp() {
         allowTestableLooperAsMainThread();
@@ -199,6 +205,11 @@
 
         when(mNotificationSwipeHelperBuilder.build()).thenReturn(mNotificationSwipeHelper);
         when(mKeyguardTransitionRepo.getTransitions()).thenReturn(emptyFlow());
+        mNotificationTestHelper = new NotificationTestHelper(
+                mContext,
+                mDependency,
+                TestableLooper.get(this));
+        mNotificationTestHelper.setDefaultInflationFlags(FLAG_CONTENT_VIEW_ALL);
     }
 
     @Test
@@ -222,6 +233,42 @@
     }
 
     @Test
+    @EnableFlags(GroupHunAnimationFix.FLAG_NAME)
+    public void changeHeadsUpAnimatingAwayToTrue_onEntryAnimatingAwayEndedNotCalled()
+            throws Exception {
+        // Before: bind an ExpandableNotificationRow,
+        initController(/* viewIsAttached= */ true);
+        mController.setHeadsUpAppearanceController(mock(HeadsUpAppearanceController.class));
+        NotificationListContainer listContainer = mController.getNotificationListContainer();
+        ExpandableNotificationRow row = mNotificationTestHelper.createRow();
+        listContainer.bindRow(row);
+
+        // When: call setHeadsUpAnimatingAway to change set mHeadsupDisappearRunning to true
+        row.setHeadsUpAnimatingAway(true);
+
+        // Then: mHeadsUpManager.onEntryAnimatingAwayEnded is not called
+        verify(mHeadsUpManager, never()).onEntryAnimatingAwayEnded(row.getEntry());
+    }
+
+    @Test
+    @EnableFlags(GroupHunAnimationFix.FLAG_NAME)
+    public void changeHeadsUpAnimatingAwayToFalse_onEntryAnimatingAwayEndedCalled()
+            throws Exception {
+        // Before: bind an ExpandableNotificationRow, set its mHeadsupDisappearRunning to true
+        initController(/* viewIsAttached= */ true);
+        mController.setHeadsUpAppearanceController(mock(HeadsUpAppearanceController.class));
+        NotificationListContainer listContainer = mController.getNotificationListContainer();
+        ExpandableNotificationRow row = mNotificationTestHelper.createRow();
+        listContainer.bindRow(row);
+        row.setHeadsUpAnimatingAway(true);
+
+        // When: call setHeadsUpAnimatingAway to change set mHeadsupDisappearRunning to false
+        row.setHeadsUpAnimatingAway(false);
+
+        // Then: mHeadsUpManager.onEntryAnimatingAwayEnded is called
+        verify(mHeadsUpManager).onEntryAnimatingAwayEnded(row.getEntry());
+    }
+    @Test
     public void testOnDensityOrFontScaleChanged_reInflatesFooterViews() {
         initController(/* viewIsAttached= */ true);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
index e46906f..4762527 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
@@ -66,7 +66,8 @@
             repository = powerRepository,
             falsingCollector = mock(),
             screenOffAnimationController = mock(),
-            statusBarStateController = mock()
+            statusBarStateController = mock(),
+            cameraGestureHelper = mock(),
         )
 
     private val configurationRepository =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java
index 5e5586d..d9e9495 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java
@@ -42,6 +42,7 @@
 import com.android.systemui.assist.AssistManager;
 import com.android.systemui.emergency.EmergencyGestureModule.EmergencyGestureIntentFactory;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.qs.QSHost;
 import com.android.systemui.recents.ScreenPinningRequest;
@@ -103,6 +104,7 @@
     @Mock private QSHost mQSHost;
     @Mock private ActivityStarter mActivityStarter;
     @Mock private EmergencyGestureIntentFactory mEmergencyGestureIntentFactory;
+    @Mock private KeyguardInteractor mKeyguardInteractor;
 
     CentralSurfacesCommandQueueCallbacks mSbcqCallbacks;
 
@@ -140,6 +142,7 @@
                 mUserTracker,
                 mQSHost,
                 mActivityStarter,
+                mKeyguardInteractor,
                 mEmergencyGestureIntentFactory);
 
         when(mUserTracker.getUserHandle()).thenReturn(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
index d2540a6..bd9cccd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesImplTest.java
@@ -150,6 +150,7 @@
 import com.android.systemui.shade.ShadeControllerImpl;
 import com.android.systemui.shade.ShadeExpansionStateManager;
 import com.android.systemui.shade.ShadeLogger;
+import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.KeyboardShortcutListSearch;
 import com.android.systemui.statusbar.KeyboardShortcuts;
@@ -343,7 +344,7 @@
     @Mock private NotificationManager mNotificationManager;
     @Mock private GlanceableHubContainerController mGlanceableHubContainerController;
     @Mock private EmergencyGestureIntentFactory mEmergencyGestureIntentFactory;
-
+    @Mock private NotificationSettingsInteractor mNotificationSettingsInteractor;
     private ShadeController mShadeController;
     private final FakeSystemClock mFakeSystemClock = new FakeSystemClock();
     private final FakeGlobalSettings mFakeGlobalSettings = new FakeGlobalSettings();
@@ -403,7 +404,8 @@
                         mPackageManager,
                         Optional.of(mBubbles),
                         mContext,
-                        mNotificationManager);
+                        mNotificationManager,
+                        mNotificationSettingsInteractor);
         mVisualInterruptionDecisionProvider.start();
 
         mContext.addMockSystemService(TrustManager.class, mock(TrustManager.class));
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java
index cf87afb..7f33c23 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardClockPositionAlgorithmTest.java
@@ -299,27 +299,7 @@
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
-    public void notifPaddingMakesUpToFullMarginInSplitShade_refactorFlagOff_usesResource() {
-        int keyguardSplitShadeTopMargin = 100;
-        int largeScreenHeaderHeightResource = 70;
-        when(mResources.getDimensionPixelSize(R.dimen.keyguard_split_shade_top_margin))
-                .thenReturn(keyguardSplitShadeTopMargin);
-        when(mResources.getDimensionPixelSize(R.dimen.large_screen_shade_header_height))
-                .thenReturn(largeScreenHeaderHeightResource);
-        mClockPositionAlgorithm.loadDimens(mContext, mResources);
-        givenLockScreen();
-        mIsSplitShade = true;
-        // WHEN the position algorithm is run
-        positionClock();
-        // THEN the notif padding makes up lacking margin (margin - header height).
-        int expectedPadding = keyguardSplitShadeTopMargin - largeScreenHeaderHeightResource;
-        assertThat(mClockPosition.stackScrollerPadding).isEqualTo(expectedPadding);
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_CENTRALIZED_STATUS_BAR_HEIGHT_FIX)
-    public void notifPaddingMakesUpToFullMarginInSplitShade_refactorFlagOn_usesHelper() {
+    public void notifPaddingMakesUpToFullMarginInSplitShade_usesHelper() {
         int keyguardSplitShadeTopMargin = 100;
         int largeScreenHeaderHeightHelper = 50;
         int largeScreenHeaderHeightResource = 70;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt
index 917e5b8..0641e17 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt
@@ -62,7 +62,7 @@
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        settings = FakeGlobalSettings()
+        settings = FakeGlobalSettings(testContext)
 
         whenever(telephonyManager.emergencyCallbackMode).thenReturn(false)
         whenever(subscriptionManager.activeSubscriptionIdList).thenReturn(intArrayOf())
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegateTest.kt
new file mode 100644
index 0000000..bf0a39b
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegateTest.kt
@@ -0,0 +1,124 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+@file:OptIn(ExperimentalCoroutinesApi::class)
+
+package com.android.systemui.statusbar.policy.ui.dialog
+
+import android.app.Dialog
+import android.content.Intent
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.animation.ActivityTransitionAnimator
+import com.android.systemui.animation.mockActivityTransitionAnimatorController
+import com.android.systemui.animation.mockDialogTransitionAnimator
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.mainCoroutineContext
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.plugins.activityStarter
+import com.android.systemui.runOnMainThreadAndWaitForIdleSync
+import com.android.systemui.statusbar.phone.SystemUIDialog
+import com.android.systemui.statusbar.phone.systemUIDialogFactory
+import com.android.systemui.statusbar.policy.ui.dialog.viewmodel.modesDialogViewModel
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ModesDialogDelegateTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val activityStarter = kosmos.activityStarter
+    private val mockDialogTransitionAnimator = kosmos.mockDialogTransitionAnimator
+    private val mockAnimationController = kosmos.mockActivityTransitionAnimatorController
+    private lateinit var underTest: ModesDialogDelegate
+
+    @Before
+    fun setup() {
+        whenever(
+                mockDialogTransitionAnimator.createActivityTransitionController(
+                    any<SystemUIDialog>(),
+                    eq(null)
+                )
+            )
+            .thenReturn(mockAnimationController)
+
+        underTest =
+            ModesDialogDelegate(
+                kosmos.systemUIDialogFactory,
+                mockDialogTransitionAnimator,
+                activityStarter,
+                { kosmos.modesDialogViewModel },
+                kosmos.mainCoroutineContext,
+            )
+    }
+
+    @Test
+    fun launchFromDialog_whenDialogNotOpen() {
+        val intent: Intent = mock()
+
+        runOnMainThreadAndWaitForIdleSync { underTest.launchFromDialog(intent) }
+
+        verify(activityStarter)
+            .startActivity(eq(intent), eq(true), eq<ActivityTransitionAnimator.Controller?>(null))
+    }
+
+    @Test
+    fun launchFromDialog_whenDialogOpen() =
+        testScope.runTest {
+            val intent: Intent = mock()
+            lateinit var dialog: Dialog
+
+            runOnMainThreadAndWaitForIdleSync {
+                kosmos.applicationCoroutineScope.launch { dialog = underTest.showDialog() }
+                runCurrent()
+                underTest.launchFromDialog(intent)
+            }
+
+            verify(mockDialogTransitionAnimator)
+                .createActivityTransitionController(any<Dialog>(), eq(null))
+            verify(activityStarter).startActivity(eq(intent), eq(true), eq(mockAnimationController))
+
+            runOnMainThreadAndWaitForIdleSync { dialog.dismiss() }
+        }
+
+    @Test
+    fun dismiss_clearsDialogReference() {
+        val dialog = runOnMainThreadAndWaitForIdleSync { underTest.createDialog() }
+
+        assertThat(underTest.currentDialog).isEqualTo(dialog)
+
+        runOnMainThreadAndWaitForIdleSync {
+            dialog.show()
+            dialog.dismiss()
+        }
+
+        assertThat(underTest.currentDialog).isNull()
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/settings/SettingsProxyTest.kt b/packages/SystemUI/tests/src/com/android/systemui/util/settings/SettingsProxyTest.kt
index 5ac6110..b0acd03 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/settings/SettingsProxyTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/settings/SettingsProxyTest.kt
@@ -275,6 +275,18 @@
     }
 
     @Test
+    fun getInt_keyMalformed_returnDefaultValue() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThat(mSettings.getInt(TEST_SETTING, 5)).isEqualTo(5)
+    }
+
+    @Test
+    fun getInt_keyMalformed_throwException() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThrows(SettingNotFoundException::class.java) { mSettings.getInt(TEST_SETTING) }
+    }
+
+    @Test
     fun getBool_keyPresent_returnValidValue() {
         mSettings.putBool(TEST_SETTING, true)
         assertThat(mSettings.getBool(TEST_SETTING)).isTrue()
@@ -323,6 +335,18 @@
     }
 
     @Test
+    fun getLong_keyMalformed_throwException() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThrows(SettingNotFoundException::class.java) { mSettings.getLong(TEST_SETTING) }
+    }
+
+    @Test
+    fun getLong_keyMalformed_returnDefaultValue() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThat(mSettings.getLong(TEST_SETTING, 2L)).isEqualTo(2L)
+    }
+
+    @Test
     fun getFloat_keyPresent_returnValidValue() {
         mSettings.putFloat(TEST_SETTING, 2.5F)
         assertThat(mSettings.getFloat(TEST_SETTING)).isEqualTo(2.5F)
@@ -346,6 +370,18 @@
         assertThat(mSettings.getFloat(TEST_SETTING, 2.5F)).isEqualTo(2.5F)
     }
 
+    @Test
+    fun getFloat_keyMalformed_throwException() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThrows(SettingNotFoundException::class.java) { mSettings.getFloat(TEST_SETTING) }
+    }
+
+    @Test
+    fun getFloat_keyMalformed_returnDefaultValue() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThat(mSettings.getFloat(TEST_SETTING, 2.5F)).isEqualTo(2.5F)
+    }
+
     private class FakeSettingsProxy(val testDispatcher: CoroutineDispatcher) : SettingsProxy {
 
         private val mContentResolver = mock(ContentResolver::class.java)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/settings/UserSettingsProxyTest.kt b/packages/SystemUI/tests/src/com/android/systemui/util/settings/UserSettingsProxyTest.kt
index 5f7420d..ead9939 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/settings/UserSettingsProxyTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/settings/UserSettingsProxyTest.kt
@@ -23,6 +23,7 @@
 import android.os.Handler
 import android.os.Looper
 import android.provider.Settings
+import android.provider.Settings.SettingNotFoundException
 import android.testing.TestableLooper
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -433,6 +434,18 @@
     }
 
     @Test
+    fun getInt_keyMalformed_returnDefaultValue() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThat(mSettings.getInt(TEST_SETTING, 5)).isEqualTo(5)
+    }
+
+    @Test
+    fun getInt_keyMalformed_throwException() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThrows(SettingNotFoundException::class.java) { mSettings.getInt(TEST_SETTING) }
+    }
+
+    @Test
     fun getIntForUser_multipleUsers__validResult() {
         mSettings.putIntForUser(TEST_SETTING, 1, MAIN_USER_ID)
         mSettings.putIntForUser(TEST_SETTING, 2, SECONDARY_USER_ID)
@@ -501,6 +514,18 @@
     }
 
     @Test
+    fun getLong_keyMalformed_throwException() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThrows(SettingNotFoundException::class.java) { mSettings.getLong(TEST_SETTING) }
+    }
+
+    @Test
+    fun getLong_keyMalformed_returnDefaultValue() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThat(mSettings.getLong(TEST_SETTING, 2L)).isEqualTo(2L)
+    }
+
+    @Test
     fun getLongForUser_multipleUsers__validResult() {
         mSettings.putLongForUser(TEST_SETTING, 1L, MAIN_USER_ID)
         mSettings.putLongForUser(TEST_SETTING, 2L, SECONDARY_USER_ID)
@@ -535,6 +560,18 @@
     }
 
     @Test
+    fun getFloat_keyMalformed_throwException() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThrows(SettingNotFoundException::class.java) { mSettings.getFloat(TEST_SETTING) }
+    }
+
+    @Test
+    fun getFloat_keyMalformed_returnDefaultValue() {
+        mSettings.putString(TEST_SETTING, "nan")
+        assertThat(mSettings.getFloat(TEST_SETTING, 2.5F)).isEqualTo(2.5F)
+    }
+
+    @Test
     fun getFloatForUser_multipleUsers__validResult() {
         mSettings.putFloatForUser(TEST_SETTING, 1F, MAIN_USER_ID)
         mSettings.putFloatForUser(TEST_SETTING, 2F, SECONDARY_USER_ID)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/CsdWarningDialogTest.java b/packages/SystemUI/tests/src/com/android/systemui/volume/CsdWarningDialogTest.java
index 49aedcc..bebf1cf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/volume/CsdWarningDialogTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/CsdWarningDialogTest.java
@@ -36,7 +36,6 @@
 import android.media.AudioManager;
 import android.platform.test.annotations.EnableFlags;
 import android.testing.TestableLooper;
-import android.util.Pair;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -70,6 +69,8 @@
     private CsdWarningDialog mDialog;
     private static final String DISMISS_CSD_NOTIFICATION =
             "com.android.systemui.volume.DISMISS_CSD_NOTIFICATION";
+    private final Optional<ImmutableList<CsdWarningAction>> mEmptyActions =
+            Optional.of(ImmutableList.of());
 
     @Before
     public void setup() {
@@ -87,7 +88,7 @@
         // instantiate directly instead of via factory; we don't want executor to be @Background
         mDialog = new CsdWarningDialog(CSD_WARNING_DOSE_REACHED_1X, mContext,
                 mAudioManager, mNotificationManager, executor, null,
-                Optional.of(ImmutableList.of(new Pair("", new Intent()))),
+                mEmptyActions,
                 mFakeBroadcastDispatcher);
 
         mDialog.show();
@@ -104,7 +105,7 @@
         FakeExecutor executor =  new FakeExecutor(new FakeSystemClock());
         mDialog = new CsdWarningDialog(CSD_WARNING_DOSE_REPEATED_5X, mContext,
                 mAudioManager, mNotificationManager, executor, null,
-                Optional.of(ImmutableList.of(new Pair("", new Intent()))),
+                mEmptyActions,
                 mFakeBroadcastDispatcher);
 
         mDialog.show();
@@ -121,7 +122,7 @@
                 .setPackage(mContext.getPackageName());
         mDialog = new CsdWarningDialog(CSD_WARNING_DOSE_REPEATED_5X, mContext,
                 mAudioManager, mNotificationManager, executor, null,
-                Optional.of(ImmutableList.of(new Pair("Undo", undoIntent))),
+                Optional.of(ImmutableList.of(new CsdWarningAction("Undo", undoIntent, false))),
                 mFakeBroadcastDispatcher);
 
         when(mAudioManager.getStreamVolume(AudioManager.STREAM_MUSIC)).thenReturn(25);
@@ -148,7 +149,7 @@
                 .setPackage(mContext.getPackageName());
         mDialog = new CsdWarningDialog(CSD_WARNING_DOSE_REPEATED_5X, mContext,
                 mAudioManager, mNotificationManager, executor, null,
-                Optional.of(ImmutableList.of(new Pair("Undo", undoIntent))),
+                Optional.of(ImmutableList.of(new CsdWarningAction("Undo", undoIntent, false))),
                 mFakeBroadcastDispatcher);
         Intent dismissIntent = new Intent(DISMISS_CSD_NOTIFICATION)
                 .setPackage(mContext.getPackageName());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
index 3f5dc82..8b7d921 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
@@ -252,7 +252,6 @@
     }
 
     @Test
-    @EnableFlags(Flags.FLAG_VOLUME_DIALOG_AUDIO_SHARING_FIX)
     public void handleAudioSharingStreamVolumeChanges_updateState() {
         ArgumentCaptor<VolumeDialogController.State> stateCaptor =
                 ArgumentCaptor.forClass(VolumeDialogController.State.class);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java
index b5cbf59..caa1779 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java
@@ -44,7 +44,6 @@
 import static org.mockito.Mockito.when;
 
 import android.app.KeyguardManager;
-import android.content.Intent;
 import android.content.res.Configuration;
 import android.graphics.Bitmap;
 import android.graphics.Canvas;
@@ -57,7 +56,6 @@
 import android.provider.Settings;
 import android.testing.TestableLooper;
 import android.util.Log;
-import android.util.Pair;
 import android.view.Gravity;
 import android.view.InputDevice;
 import android.view.MotionEvent;
@@ -166,7 +164,7 @@
             new CsdWarningDialog.Factory() {
                 @Override
                 public CsdWarningDialog create(int warningType, Runnable onCleanup,
-                        Optional<ImmutableList<Pair<String, Intent>>> actionIntents) {
+                        Optional<ImmutableList<CsdWarningAction>> actionIntents) {
                     return mCsdWarningDialog;
                 }
             };
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index 60b5b5d..2457eb7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -98,6 +98,8 @@
 import androidx.annotation.Nullable;
 import androidx.test.filters.SmallTest;
 
+import com.android.app.viewcapture.ViewCapture;
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.internal.colorextraction.ColorExtractor;
 import com.android.internal.logging.UiEventLogger;
 import com.android.internal.statusbar.IStatusBarService;
@@ -121,6 +123,7 @@
 import com.android.systemui.shade.ShadeController;
 import com.android.systemui.shade.ShadeWindowLogger;
 import com.android.systemui.shade.domain.interactor.ShadeInteractor;
+import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.statusbar.NotificationEntryHelper;
 import com.android.systemui.statusbar.NotificationLockscreenUserManager;
@@ -210,6 +213,7 @@
 import java.util.Optional;
 import java.util.concurrent.Executor;
 
+import kotlin.Lazy;
 import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
 import platform.test.runner.parameterized.Parameters;
 
@@ -342,6 +346,8 @@
     private Icon mAppBubbleIcon;
     @Mock
     private Display mDefaultDisplay;
+    @Mock
+    private Lazy<ViewCapture> mLazyViewCapture;
 
     private final KosmosJavaAdapter mKosmos = new KosmosJavaAdapter(this);
     private ShadeInteractor mShadeInteractor;
@@ -407,7 +413,8 @@
         mNotificationShadeWindowController = new NotificationShadeWindowControllerImpl(
                 mContext,
                 new FakeWindowRootViewComponent.Factory(mNotificationShadeWindowView),
-                mWindowManager,
+                new ViewCaptureAwareWindowManager(mWindowManager, mLazyViewCapture,
+                        /* isViewCaptureEnabled= */ false),
                 mActivityManager,
                 mDozeParameters,
                 mStatusBarStateController,
@@ -480,7 +487,8 @@
                         mock(PackageManager.class),
                         Optional.of(mock(Bubbles.class)),
                         mContext,
-                        mock(NotificationManager.class)
+                        mock(NotificationManager.class),
+                        mock(NotificationSettingsInteractor.class)
                         );
         interruptionDecisionProvider.start();
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt b/packages/SystemUI/tests/utils/src/android/app/role/RoleManagerKosmos.kt
similarity index 67%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
rename to packages/SystemUI/tests/utils/src/android/app/role/RoleManagerKosmos.kt
index 37c9552..356bc86 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/android/app/role/RoleManagerKosmos.kt
@@ -14,11 +14,11 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.domain.interactor
+package android.app.role
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
+import com.android.systemui.util.mockito.mock
 
-val Kosmos.partitionedGridLayout by
-    Kosmos.Fixture { PartitionedGridLayout(partitionedGridViewModel) }
+val Kosmos.mockRoleManager: RoleManager by Kosmos.Fixture { mock() }
+
+var Kosmos.roleManager: RoleManager by Kosmos.Fixture { mockRoleManager }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java b/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java
index a124b34..27a2cab 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/SysuiTestCase.java
@@ -271,10 +271,14 @@
     }
 
     protected void waitForIdleSync() {
-        if (mHandler == null) {
-            mHandler = new Handler(Looper.getMainLooper());
+        if (isRobolectricTest()) {
+            mRealInstrumentation.waitForIdleSync();
+        } else {
+            if (mHandler == null) {
+                mHandler = new Handler(Looper.getMainLooper());
+            }
+            waitForIdleSync(mHandler);
         }
-        waitForIdleSync(mHandler);
     }
 
     protected void waitForUiOffloadThread() {
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/TestMocksModule.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/TestMocksModule.kt
index 0b6b816..5063140 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/TestMocksModule.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/TestMocksModule.kt
@@ -33,6 +33,7 @@
 import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.biometrics.AuthController
 import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor
+import com.android.systemui.camera.CameraGestureHelper
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
 import com.android.systemui.communal.domain.interactor.CommunalSettingsInteractor
@@ -143,6 +144,7 @@
     @get:Provides val primaryBouncerInteractor: PrimaryBouncerInteractor = mock(),
     @get:Provides val keyguardStateController: KeyguardStateController = mock(),
     @get:Provides val globalSettings: GlobalSettings = mock(),
+    @get:Provides val cameraGestureHelper: CameraGestureHelper = mock(),
 
     // log buffers
     @get:[Provides BroadcastDispatcherLog]
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/animation/ActivityTransitionAnimatorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/animation/ActivityTransitionAnimatorKosmos.kt
index b23767e..5ac41ec 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/animation/ActivityTransitionAnimatorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/animation/ActivityTransitionAnimatorKosmos.kt
@@ -18,6 +18,10 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testCase
+import com.android.systemui.util.mockito.mock
+
+val Kosmos.mockActivityTransitionAnimatorController by
+    Kosmos.Fixture { mock<ActivityTransitionAnimator.Controller>() }
 
 val Kosmos.activityTransitionAnimator by
     Kosmos.Fixture {
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/camera/CameraGestureHelperKosmos.kt
similarity index 61%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/camera/CameraGestureHelperKosmos.kt
index 37c9552..931567e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/camera/CameraGestureHelperKosmos.kt
@@ -14,11 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.domain.interactor
+package com.android.systemui.camera
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.whenever
+import org.mockito.Mockito.mock
 
-val Kosmos.partitionedGridLayout by
-    Kosmos.Fixture { PartitionedGridLayout(partitionedGridViewModel) }
+val Kosmos.cameraGestureHelper: CameraGestureHelper by
+    Kosmos.Fixture<CameraGestureHelper> {
+        mock(CameraGestureHelper::class.java).also { helper ->
+            whenever(helper.canCameraGestureBeLaunched(any())).thenReturn(true)
+        }
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/qs/QSLongPressEffectKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/qs/QSLongPressEffectKosmos.kt
index 28355e1..ca748b66 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/qs/QSLongPressEffectKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/qs/QSLongPressEffectKosmos.kt
@@ -16,10 +16,16 @@
 
 package com.android.systemui.haptics.qs
 
-import com.android.systemui.classifier.falsingManager
 import com.android.systemui.haptics.vibratorHelper
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.log.core.FakeLogBuffer
 import com.android.systemui.statusbar.policy.keyguardStateController
 
 val Kosmos.qsLongPressEffect by
-    Kosmos.Fixture { QSLongPressEffect(vibratorHelper, keyguardStateController, falsingManager) }
+    Kosmos.Fixture {
+        QSLongPressEffect(
+            vibratorHelper,
+            keyguardStateController,
+            FakeLogBuffer.Factory.create(),
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
index c423b62..c2a03d4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyboard.shortcut
 
+import android.app.role.mockRoleManager
 import android.content.applicationContext
 import android.content.res.mainResources
 import android.hardware.input.fakeInputManager
@@ -41,6 +42,7 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.model.sysUiState
 import com.android.systemui.settings.displayTracker
+import com.android.systemui.settings.fakeUserTracker
 
 var Kosmos.shortcutHelperAppCategoriesShortcutsSource: KeyboardShortcutGroupsSource by
     Kosmos.Fixture {
@@ -117,6 +119,8 @@
 val Kosmos.shortcutHelperViewModel by
     Kosmos.Fixture {
         ShortcutHelperViewModel(
+            mockRoleManager,
+            fakeUserTracker,
             applicationCoroutineScope,
             testDispatcher,
             shortcutHelperStateInteractor,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
index 87143ef..727de9e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
@@ -22,6 +22,7 @@
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
 import com.android.systemui.keyguard.shared.model.BiometricUnlockModel
 import com.android.systemui.keyguard.shared.model.BiometricUnlockSource
+import com.android.systemui.keyguard.shared.model.CameraLaunchSourceModel
 import com.android.systemui.keyguard.shared.model.DismissAction
 import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.KeyguardDone
@@ -138,6 +139,8 @@
     private val _canIgnoreAuthAndReturnToGone = MutableStateFlow(false)
     override val canIgnoreAuthAndReturnToGone = _canIgnoreAuthAndReturnToGone.asStateFlow()
 
+    override val onCameraLaunchDetected = MutableStateFlow(CameraLaunchSourceModel())
+
     override fun setQuickSettingsVisible(isVisible: Boolean) {
         _isQuickSettingsVisible.value = isVisible
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
index b5ca964..a95609e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
@@ -21,7 +21,6 @@
 import com.android.systemui.common.ui.data.repository.FakeConfigurationRepository
 import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.flags.FakeFeatureFlags
-import com.android.systemui.keyguard.data.repository.FakeCommandQueue
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionStep
@@ -49,7 +48,6 @@
     fun create(
         featureFlags: FakeFeatureFlags = FakeFeatureFlags(),
         repository: FakeKeyguardRepository = FakeKeyguardRepository(),
-        commandQueue: FakeCommandQueue = FakeCommandQueue(),
         bouncerRepository: FakeKeyguardBouncerRepository = FakeKeyguardBouncerRepository(),
         configurationRepository: FakeConfigurationRepository = FakeConfigurationRepository(),
         shadeRepository: FakeShadeRepository = FakeShadeRepository(),
@@ -87,7 +85,6 @@
                 }
         return WithDependencies(
             repository = repository,
-            commandQueue = commandQueue,
             featureFlags = featureFlags,
             bouncerRepository = bouncerRepository,
             configurationRepository = configurationRepository,
@@ -95,7 +92,6 @@
             powerInteractor = powerInteractor,
             KeyguardInteractor(
                 repository = repository,
-                commandQueue = commandQueue,
                 powerInteractor = powerInteractor,
                 bouncerRepository = bouncerRepository,
                 configurationInteractor = ConfigurationInteractor(configurationRepository),
@@ -112,7 +108,6 @@
 
     data class WithDependencies(
         val repository: FakeKeyguardRepository,
-        val commandQueue: FakeCommandQueue,
         val featureFlags: FakeFeatureFlags,
         val bouncerRepository: FakeKeyguardBouncerRepository,
         val configurationRepository: FakeConfigurationRepository,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
index 81d8f0b..5ab56e9 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
@@ -18,7 +18,6 @@
 
 import com.android.systemui.bouncer.data.repository.keyguardBouncerRepository
 import com.android.systemui.common.ui.domain.interactor.configurationInteractor
-import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.keyguard.data.repository.keyguardRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
@@ -31,7 +30,6 @@
     Kosmos.Fixture {
         KeyguardInteractor(
             repository = keyguardRepository,
-            commandQueue = fakeCommandQueue,
             powerInteractor = powerInteractor,
             bouncerRepository = keyguardBouncerRepository,
             configurationInteractor = configurationInteractor,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/gesture/data/GestureRepositoryKosmos.kt
similarity index 62%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/gesture/data/GestureRepositoryKosmos.kt
index 37c9552..9bd346e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/gesture/data/GestureRepositoryKosmos.kt
@@ -14,11 +14,12 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.domain.interactor
+package com.android.systemui.keyguard.gesture.data
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.navigationbar.gestural.data.respository.GestureRepository
+import com.android.systemui.navigationbar.gestural.data.respository.GestureRepositoryImpl
 
-val Kosmos.partitionedGridLayout by
-    Kosmos.Fixture { PartitionedGridLayout(partitionedGridViewModel) }
+val Kosmos.gestureRepository: GestureRepository by
+    Kosmos.Fixture { GestureRepositoryImpl(testDispatcher) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/gesture/domain/GestureInteractorKosmos.kt
similarity index 60%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/gesture/domain/GestureInteractorKosmos.kt
index 37c9552..658aaa6 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/gesture/domain/GestureInteractorKosmos.kt
@@ -14,11 +14,14 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.domain.interactor
+package com.android.systemui.keyguard.gesture.domain
 
+import com.android.systemui.keyguard.gesture.data.gestureRepository
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.navigationbar.gestural.domain.GestureInteractor
 
-val Kosmos.partitionedGridLayout by
-    Kosmos.Fixture { PartitionedGridLayout(partitionedGridViewModel) }
+val Kosmos.gestureInteractor: GestureInteractor by
+    Kosmos.Fixture {
+        GestureInteractor(gestureRepository = gestureRepository, scope = applicationCoroutineScope)
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinderKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinderKosmos.kt
index 6eb8a49..2919d3f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinderKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/binder/AlternateBouncerViewBinderKosmos.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.common.ui.domain.interactor.configurationInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryUdfpsInteractor
 import com.android.systemui.deviceentry.ui.viewmodel.AlternateBouncerUdfpsAccessibilityOverlayViewModel
+import com.android.systemui.keyguard.dismissCallbackRegistry
 import com.android.systemui.keyguard.ui.SwipeUpAnywhereGestureHandler
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerDependencies
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerMessageAreaViewModel
@@ -49,6 +50,7 @@
             alternateBouncerDependencies = { alternateBouncerDependencies },
             windowManager = { windowManager },
             layoutInflater = { mockedLayoutInflater },
+            dismissCallbackRegistry = dismissCallbackRegistry,
         )
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModelKosmos.kt
new file mode 100644
index 0000000..6c644ee
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerToLockscreenTransitionViewModelKosmos.kt
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+@file:OptIn(ExperimentalCoroutinesApi::class)
+
+package com.android.systemui.keyguard.ui.viewmodel
+
+import com.android.systemui.keyguard.ui.keyguardTransitionAnimationFlow
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+
+val Kosmos.alternateBouncerToLockscreenTransitionViewModel by Fixture {
+    AlternateBouncerToLockscreenTransitionViewModel(
+        animationFlow = keyguardTransitionAnimationFlow,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt
index 3c5baa5..82860fc 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt
@@ -40,6 +40,8 @@
         keyguardTransitionInteractor = keyguardTransitionInteractor,
         notificationsKeyguardInteractor = notificationsKeyguardInteractor,
         alternateBouncerToGoneTransitionViewModel = alternateBouncerToGoneTransitionViewModel,
+        alternateBouncerToLockscreenTransitionViewModel =
+            alternateBouncerToLockscreenTransitionViewModel,
         aodToGoneTransitionViewModel = aodToGoneTransitionViewModel,
         aodToLockscreenTransitionViewModel = aodToLockscreenTransitionViewModel,
         aodToOccludedTransitionViewModel = aodToOccludedTransitionViewModel,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorFactory.kt
index d92ace9..9a07c4e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorFactory.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorFactory.kt
@@ -42,6 +42,7 @@
                     falsingCollector,
                     screenOffAnimationController,
                     statusBarStateController,
+                    mock(),
                 )
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorKosmos.kt
index 8486691..d50091e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/power/domain/interactor/PowerInteractorKosmos.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.power.domain.interactor
 
+import com.android.systemui.camera.cameraGestureHelper
 import com.android.systemui.classifier.falsingCollector
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.plugins.statusbar.statusBarStateController
@@ -29,5 +30,6 @@
             falsingCollector = falsingCollector,
             screenOffAnimationController = screenOffAnimationController,
             statusBarStateController = statusBarStateController,
+            cameraGestureHelper = { cameraGestureHelper },
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
index 5568c6c..34e99d3 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/GridLayoutTypeInteractorKosmos.kt
@@ -20,24 +20,13 @@
 import com.android.systemui.qs.panels.data.repository.gridLayoutTypeRepository
 import com.android.systemui.qs.panels.shared.model.GridLayoutType
 import com.android.systemui.qs.panels.shared.model.InfiniteGridLayoutType
-import com.android.systemui.qs.panels.shared.model.PartitionedGridLayoutType
 import com.android.systemui.qs.panels.ui.compose.GridLayout
 
 val Kosmos.gridLayoutTypeInteractor by
     Kosmos.Fixture { GridLayoutTypeInteractor(gridLayoutTypeRepository) }
 
 val Kosmos.gridLayoutMap: Map<GridLayoutType, GridLayout> by
-    Kosmos.Fixture {
-        mapOf(
-            Pair(PartitionedGridLayoutType, partitionedGridLayout),
-            Pair(InfiniteGridLayoutType, infiniteGridLayout)
-        )
-    }
+    Kosmos.Fixture { mapOf(Pair(InfiniteGridLayoutType, infiniteGridLayout)) }
 
 var Kosmos.gridConsistencyInteractorsMap: Map<GridLayoutType, GridTypeConsistencyInteractor> by
-    Kosmos.Fixture {
-        mapOf(
-            Pair(PartitionedGridLayoutType, noopGridConsistencyInteractor),
-            Pair(InfiniteGridLayoutType, infiniteGridConsistencyInteractor)
-        )
-    }
+    Kosmos.Fixture { mapOf(Pair(InfiniteGridLayoutType, infiniteGridConsistencyInteractor)) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
index 6625bb5..9481fca 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/TileGridViewModelKosmos.kt
@@ -20,7 +20,7 @@
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.qs.panels.domain.interactor.gridLayoutMap
 import com.android.systemui.qs.panels.domain.interactor.gridLayoutTypeInteractor
-import com.android.systemui.qs.panels.domain.interactor.partitionedGridLayout
+import com.android.systemui.qs.panels.domain.interactor.infiniteGridLayout
 import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
 
 val Kosmos.tileGridViewModel by
@@ -29,7 +29,7 @@
             gridLayoutTypeInteractor,
             gridLayoutMap,
             currentTilesInteractor,
-            partitionedGridLayout,
+            infiniteGridLayout,
             applicationCoroutineScope,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorKosmos.kt
new file mode 100644
index 0000000..2ecfb45
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorKosmos.kt
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2024 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.qs.tiles.impl.modes.domain.interactor
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.tiles.base.actions.qsTileIntentUserInputHandler
+import com.android.systemui.statusbar.policy.ui.dialog.modesDialogDelegate
+import javax.inject.Provider
+
+val Kosmos.modesTileUserActionInteractor: ModesTileUserActionInteractor by
+    Kosmos.Fixture {
+        ModesTileUserActionInteractor(
+            qsTileIntentUserInputHandler,
+            Provider { modesDialogDelegate }.get(),
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryKosmos.kt
index a75d2bc..ecfc168 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryKosmos.kt
@@ -20,6 +20,7 @@
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.shared.settings.data.repository.secureSettingsRepository
+import com.android.systemui.shared.settings.data.repository.systemSettingsRepository
 
 val Kosmos.notificationSettingsRepository by
     Kosmos.Fixture {
@@ -27,5 +28,6 @@
             scope = testScope.backgroundScope,
             backgroundDispatcher = testDispatcher,
             secureSettingsRepository = secureSettingsRepository,
+            systemSettingsRepository = systemSettingsRepository,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shared/settings/data/repository/SystemSettingsRepositoryKosmos.kt
similarity index 67%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/shared/settings/data/repository/SystemSettingsRepositoryKosmos.kt
index 37c9552..01f19ae 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/PartitionedGridLayoutKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shared/settings/data/repository/SystemSettingsRepositoryKosmos.kt
@@ -14,11 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.panels.domain.interactor
+package com.android.systemui.shared.settings.data.repository
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.qs.panels.ui.compose.PartitionedGridLayout
-import com.android.systemui.qs.panels.ui.viewmodel.partitionedGridViewModel
 
-val Kosmos.partitionedGridLayout by
-    Kosmos.Fixture { PartitionedGridLayout(partitionedGridViewModel) }
+var Kosmos.systemSettingsRepository: SystemSettingsRepository by
+    Kosmos.Fixture { fakeSystemSettingsRepository }
+val Kosmos.fakeSystemSettingsRepository by Kosmos.Fixture { FakeSystemSettingsRepository() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
index b8dec31..0b309b5 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
@@ -319,14 +319,9 @@
         // NOTE: This flag is read when the ExpandableNotificationRow is inflated, so it needs to be
         //  set, but we do not want to override an existing value that is needed by a specific test.
 
-        val rowFuture: SettableFuture<ExpandableNotificationRow> = SettableFuture.create()
         val rowInflaterTask =
             RowInflaterTask(mFakeSystemClock, Mockito.mock(RowInflaterTaskLogger::class.java))
-        rowInflaterTask.inflate(context, null, entry, MoreExecutors.directExecutor()) { inflatedRow
-            ->
-            rowFuture.set(inflatedRow)
-        }
-        val row = rowFuture.get(1, TimeUnit.SECONDS)
+        val row = rowInflaterTask.inflateSynchronously(context, null, entry)
 
         entry.row = row
         mIconManager.createIcons(entry)
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegateKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegateKosmos.kt
new file mode 100644
index 0000000..99bb479
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/ui/dialog/ModesDialogDelegateKosmos.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 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.policy.ui.dialog
+
+import com.android.systemui.animation.dialogTransitionAnimator
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.mainCoroutineContext
+import com.android.systemui.plugins.activityStarter
+import com.android.systemui.statusbar.phone.systemUIDialogFactory
+import com.android.systemui.statusbar.policy.ui.dialog.viewmodel.modesDialogViewModel
+import com.android.systemui.util.mockito.mock
+
+val Kosmos.mockModesDialogDelegate by Kosmos.Fixture { mock<ModesDialogDelegate>() }
+
+var Kosmos.modesDialogDelegate: ModesDialogDelegate by
+    Kosmos.Fixture {
+        ModesDialogDelegate(
+            systemUIDialogFactory,
+            dialogTransitionAnimator,
+            activityStarter,
+            { modesDialogViewModel },
+            mainCoroutineContext,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelKosmos.kt
new file mode 100644
index 0000000..00020f8
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelKosmos.kt
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2024 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.policy.ui.dialog.viewmodel
+
+import android.content.mockedContext
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.statusbar.policy.domain.interactor.zenModeInteractor
+import com.android.systemui.statusbar.policy.ui.dialog.modesDialogDelegate
+import javax.inject.Provider
+
+val Kosmos.modesDialogViewModel: ModesDialogViewModel by
+    Kosmos.Fixture {
+        ModesDialogViewModel(
+            mockedContext,
+            zenModeInteractor,
+            testDispatcher,
+            Provider { modesDialogDelegate }.get(),
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettings.java b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettings.java
index 476b7d8..65f4122 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettings.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettings.java
@@ -38,6 +38,11 @@
 
     public static final Uri CONTENT_URI = Uri.parse("content://settings/fake_global");
 
+    /**
+     * @deprecated Please use FakeGlobalSettings(testDispatcher) to provide the same dispatcher used
+     * by main test scope.
+     */
+    @Deprecated
     public FakeGlobalSettings() {
         mDispatcher = StandardTestDispatcher(/* scheduler = */ null, /* name = */ null);
     }
@@ -46,6 +51,7 @@
         mDispatcher = dispatcher;
     }
 
+    @NonNull
     @Override
     public ContentResolver getContentResolver() {
         throw new UnsupportedOperationException(
@@ -53,6 +59,7 @@
                         + "GlobalSettings.registerContentObserver helpful instead.");
     }
 
+    @NonNull
     @Override
     public CoroutineDispatcher getBackgroundDispatcher() {
         return mDispatcher;
@@ -60,7 +67,7 @@
 
     @Override
     public void registerContentObserverSync(Uri uri, boolean notifyDescendants,
-            ContentObserver settingsObserver) {
+            @NonNull ContentObserver settingsObserver) {
         List<ContentObserver> observers;
         mContentObserversAllUsers.putIfAbsent(uri.toString(), new ArrayList<>());
         observers = mContentObserversAllUsers.get(uri.toString());
@@ -68,25 +75,26 @@
     }
 
     @Override
-    public void unregisterContentObserverSync(ContentObserver settingsObserver) {
+    public void unregisterContentObserverSync(@NonNull ContentObserver settingsObserver) {
         for (Map.Entry<String, List<ContentObserver>> entry :
                 mContentObserversAllUsers.entrySet()) {
             entry.getValue().remove(settingsObserver);
         }
     }
 
+    @NonNull
     @Override
-    public Uri getUriFor(String name) {
+    public Uri getUriFor(@NonNull String name) {
         return Uri.withAppendedPath(CONTENT_URI, name);
     }
 
     @Override
-    public String getString(String name) {
+    public String getString(@NonNull String name) {
         return mValues.get(getUriFor(name).toString());
     }
 
     @Override
-    public boolean putString(String name, String value) {
+    public boolean putString(@NonNull String name, @NonNull String value) {
         return putString(name, value, null, false);
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt
index df6fc41..35fa2af 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt
@@ -18,5 +18,6 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testDispatcher
 
-val Kosmos.fakeGlobalSettings: FakeGlobalSettings by Fixture { FakeGlobalSettings() }
+val Kosmos.fakeGlobalSettings: FakeGlobalSettings by Fixture { FakeGlobalSettings(testDispatcher) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettings.java b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettings.java
index e35da11..3f0318b 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettings.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettings.java
@@ -26,7 +26,9 @@
 import android.util.Pair;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.VisibleForTesting;
 
+import com.android.systemui.settings.FakeUserTracker;
 import com.android.systemui.settings.UserTracker;
 
 import kotlinx.coroutines.CoroutineDispatcher;
@@ -42,49 +44,68 @@
             new HashMap<>();
     private final Map<String, List<ContentObserver>> mContentObserversAllUsers = new HashMap<>();
     private final CoroutineDispatcher mDispatcher;
+    private final UserTracker mUserTracker;
 
     public static final Uri CONTENT_URI = Uri.parse("content://settings/fake");
     @UserIdInt
     private int mUserId = UserHandle.USER_CURRENT;
 
+    /**
+     * @deprecated Please use FakeSettings(testDispatcher) to provide the same dispatcher used
+     * by main test scope.
+     */
+    @Deprecated
     public FakeSettings() {
         mDispatcher = StandardTestDispatcher(/* scheduler = */ null, /* name = */ null);
+        mUserTracker = new FakeUserTracker();
     }
 
     public FakeSettings(CoroutineDispatcher dispatcher) {
         mDispatcher = dispatcher;
+        mUserTracker = new FakeUserTracker();
     }
 
-    public FakeSettings(String initialKey, String initialValue) {
-        mDispatcher = StandardTestDispatcher(/* scheduler = */ null, /* name = */ null);
+    public FakeSettings(CoroutineDispatcher dispatcher, UserTracker userTracker) {
+        mDispatcher = dispatcher;
+        mUserTracker = userTracker;
+    }
+
+    @VisibleForTesting
+    FakeSettings(String initialKey, String initialValue) {
+        this();
         putString(initialKey, initialValue);
     }
 
-    public FakeSettings(Map<String, String> initialValues) {
-        mDispatcher = StandardTestDispatcher(/* scheduler = */ null, /* name = */ null);
+    @VisibleForTesting
+    FakeSettings(Map<String, String> initialValues) {
+        this();
         for (Map.Entry<String, String> kv : initialValues.entrySet()) {
             putString(kv.getKey(), kv.getValue());
         }
     }
 
     @Override
+    @NonNull
     public ContentResolver getContentResolver() {
-        return null;
+        throw new UnsupportedOperationException(
+                "FakeSettings.getContentResolver is not implemented");
     }
 
+    @NonNull
     @Override
     public UserTracker getUserTracker() {
-        return null;
+        return mUserTracker;
     }
 
+    @NonNull
     @Override
     public CoroutineDispatcher getBackgroundDispatcher() {
         return mDispatcher;
     }
 
     @Override
-    public void registerContentObserverForUserSync(Uri uri, boolean notifyDescendants,
-            ContentObserver settingsObserver, int userHandle) {
+    public void registerContentObserverForUserSync(@NonNull Uri uri, boolean notifyDescendants,
+            @NonNull ContentObserver settingsObserver, int userHandle) {
         List<ContentObserver> observers;
         if (userHandle == UserHandle.USER_ALL) {
             mContentObserversAllUsers.putIfAbsent(uri.toString(), new ArrayList<>());
@@ -98,19 +119,18 @@
     }
 
     @Override
-    public void unregisterContentObserverSync(ContentObserver settingsObserver) {
-        for (SettingsKey key : mContentObservers.keySet()) {
-            List<ContentObserver> observers = mContentObservers.get(key);
+    public void unregisterContentObserverSync(@NonNull ContentObserver settingsObserver) {
+        for (List<ContentObserver> observers : mContentObservers.values()) {
             observers.remove(settingsObserver);
         }
-        for (String key : mContentObserversAllUsers.keySet()) {
-            List<ContentObserver> observers = mContentObserversAllUsers.get(key);
+        for (List<ContentObserver> observers : mContentObserversAllUsers.values()) {
             observers.remove(settingsObserver);
         }
     }
 
+    @NonNull
     @Override
-    public Uri getUriFor(String name) {
+    public Uri getUriFor(@NonNull String name) {
         return Uri.withAppendedPath(CONTENT_URI, name);
     }
 
@@ -124,33 +144,34 @@
     }
 
     @Override
-    public String getString(String name) {
+    public String getString(@NonNull String name) {
         return getStringForUser(name, getUserId());
     }
 
     @Override
-    public String getStringForUser(String name, int userHandle) {
+    public String getStringForUser(@NonNull String name, int userHandle) {
         return mValues.get(new SettingsKey(userHandle, getUriFor(name).toString()));
     }
 
     @Override
-    public boolean putString(String name, String value, boolean overrideableByRestore) {
+    public boolean putString(@NonNull String name, @NonNull String value,
+            boolean overrideableByRestore) {
         return putStringForUser(name, value, null, false, getUserId(), overrideableByRestore);
     }
 
     @Override
-    public boolean putString(String name, String value) {
+    public boolean putString(@NonNull String name, @NonNull String value) {
         return putString(name, value, false);
     }
 
     @Override
-    public boolean putStringForUser(String name, String value, int userHandle) {
+    public boolean putStringForUser(@NonNull String name, @NonNull String value, int userHandle) {
         return putStringForUser(name, value, null, false, userHandle, false);
     }
 
     @Override
-    public boolean putStringForUser(String name, String value, String tag, boolean makeDefault,
-            int userHandle, boolean overrideableByRestore) {
+    public boolean putStringForUser(@NonNull String name, @NonNull String value, String tag,
+            boolean makeDefault, int userHandle, boolean overrideableByRestore) {
         SettingsKey key = new SettingsKey(userHandle, getUriFor(name).toString());
         mValues.put(key, value);
 
@@ -166,7 +187,8 @@
     }
 
     @Override
-    public boolean putString(@NonNull String name, String value, String tag, boolean makeDefault) {
+    public boolean putString(@NonNull String name, @NonNull String value, @NonNull String tag,
+            boolean makeDefault) {
         return putString(name, value);
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt
index bcb5848..55044bf 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt
@@ -18,5 +18,7 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.settings.fakeUserTracker
 
-val Kosmos.fakeSettings: FakeSettings by Fixture { FakeSettings() }
+val Kosmos.fakeSettings: FakeSettings by Fixture { FakeSettings(testDispatcher, fakeUserTracker) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioSharingRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioSharingRepository.kt
index d391750..0a617d1 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioSharingRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioSharingRepository.kt
@@ -16,10 +16,7 @@
 
 package com.android.systemui.volume.data.repository
 
-import androidx.annotation.IntRange
 import com.android.settingslib.volume.data.repository.AudioSharingRepository
-import com.android.settingslib.volume.data.repository.AudioSharingRepository.Companion.AUDIO_SHARING_VOLUME_MAX
-import com.android.settingslib.volume.data.repository.AudioSharingRepository.Companion.AUDIO_SHARING_VOLUME_MIN
 import com.android.settingslib.volume.data.repository.GroupIdToVolumes
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -27,11 +24,14 @@
 
 class FakeAudioSharingRepository : AudioSharingRepository {
     private val mutableInAudioSharing: MutableStateFlow<Boolean> = MutableStateFlow(false)
+    private val mutablePrimaryGroupId: MutableStateFlow<Int> =
+        MutableStateFlow(TEST_GROUP_ID_INVALID)
     private val mutableSecondaryGroupId: MutableStateFlow<Int> =
         MutableStateFlow(TEST_GROUP_ID_INVALID)
     private val mutableVolumeMap: MutableStateFlow<GroupIdToVolumes> = MutableStateFlow(emptyMap())
 
     override val inAudioSharing: Flow<Boolean> = mutableInAudioSharing
+    override val primaryGroupId: StateFlow<Int> = mutablePrimaryGroupId
     override val secondaryGroupId: StateFlow<Int> = mutableSecondaryGroupId
     override val volumeMap: StateFlow<GroupIdToVolumes> = mutableVolumeMap
 
@@ -41,6 +41,10 @@
         mutableInAudioSharing.value = state
     }
 
+    fun setPrimaryGroupId(groupId: Int) {
+        mutablePrimaryGroupId.value = groupId
+    }
+
     fun setSecondaryGroupId(groupId: Int) {
         mutableSecondaryGroupId.value = groupId
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorKosmos.kt
index 03981bb..ce8aba5 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/domain/interactor/AudioSharingInteractorKosmos.kt
@@ -18,12 +18,15 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.backgroundCoroutineContext
 import com.android.systemui.volume.data.repository.audioSharingRepository
 
 val Kosmos.audioSharingInteractor by
     Kosmos.Fixture {
         AudioSharingInteractorImpl(
             applicationCoroutineScope,
+            backgroundCoroutineContext,
+            audioVolumeInteractor,
             audioSharingRepository,
         )
     }
diff --git a/ravenwood/Android.bp b/ravenwood/Android.bp
index 8d89cc1..7c8fd42 100644
--- a/ravenwood/Android.bp
+++ b/ravenwood/Android.bp
@@ -68,7 +68,10 @@
     srcs: [
         "runtime-common-ravenwood-src/**/*.java",
     ],
-    visibility: ["//frameworks/base"],
+    visibility: [
+        // Some tests need to access the utilities.
+        ":__subpackages__",
+    ],
 }
 
 java_library {
@@ -318,6 +321,9 @@
 
 android_ravenwood_libgroup {
     name: "ravenwood-runtime",
+    data: [
+        "framework-res",
+    ],
     libs: [
         "100-framework-minus-apex.ravenwood",
         "200-kxml2-android",
@@ -330,6 +336,11 @@
         "services.core.ravenwood-jarjar",
         "services.fakes.ravenwood-jarjar",
 
+        // ICU
+        "core-icu4j-for-host.ravenwood",
+        "icu4j-icudata-jarjar",
+        "icu4j-icutzdata-jarjar",
+
         // Provide runtime versions of utils linked in below
         "junit",
         "truth",
diff --git a/ravenwood/TEST_MAPPING b/ravenwood/TEST_MAPPING
index f6885e1..fbf27fa 100644
--- a/ravenwood/TEST_MAPPING
+++ b/ravenwood/TEST_MAPPING
@@ -12,6 +12,9 @@
     {
       "name": "RavenwoodBivalentTest_device"
     },
+    {
+      "name": "RavenwoodResApkTest"
+    },
     // The sysui tests should match vendor/unbundled_google/packages/SystemUIGoogle/TEST_MAPPING
     {
       "name": "SystemUIGoogleTests",
diff --git a/ravenwood/resapk_test/Android.bp b/ravenwood/resapk_test/Android.bp
new file mode 100644
index 0000000..c145765
--- /dev/null
+++ b/ravenwood/resapk_test/Android.bp
@@ -0,0 +1,30 @@
+package {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_ravenwood_test {
+    name: "RavenwoodResApkTest",
+
+    resource_apk: "RavenwoodResApkTest-apk",
+
+    libs: [
+        // Normally, tests shouldn't directly access it, but we need to access RavenwoodCommonUtils
+        // in this test.
+        "ravenwood-runtime-common-ravenwood",
+    ],
+    static_libs: [
+        "androidx.annotation_annotation",
+        "androidx.test.ext.junit",
+        "androidx.test.rules",
+    ],
+    srcs: [
+        "test/**/*.java",
+    ],
+    sdk_version: "test_current",
+    auto_gen_config: true,
+}
diff --git a/ravenwood/resapk_test/apk/Android.bp b/ravenwood/resapk_test/apk/Android.bp
new file mode 100644
index 0000000..10ed5e2
--- /dev/null
+++ b/ravenwood/resapk_test/apk/Android.bp
@@ -0,0 +1,14 @@
+package {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_app {
+    name: "RavenwoodResApkTest-apk",
+
+    sdk_version: "current",
+}
diff --git a/ravenwood/resapk_test/apk/AndroidManifest.xml b/ravenwood/resapk_test/apk/AndroidManifest.xml
new file mode 100644
index 0000000..f34d8b2
--- /dev/null
+++ b/ravenwood/resapk_test/apk/AndroidManifest.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+        package="com.android.ravenwood.restest_apk">
+</manifest>
diff --git a/ravenwood/resapk_test/apk/res/values/strings.xml b/ravenwood/resapk_test/apk/res/values/strings.xml
new file mode 100644
index 0000000..23d4c0f
--- /dev/null
+++ b/ravenwood/resapk_test/apk/res/values/strings.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+<resources xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <!-- Test string 1 -->
+    <string name="test_string_1" translatable="false" >Test String 1</string>
+</resources>
diff --git a/ravenwood/resapk_test/test/com/android/ravenwood/resapk_test/RavenwoodResApkTest.java b/ravenwood/resapk_test/test/com/android/ravenwood/resapk_test/RavenwoodResApkTest.java
new file mode 100644
index 0000000..1029ed2
--- /dev/null
+++ b/ravenwood/resapk_test/test/com/android/ravenwood/resapk_test/RavenwoodResApkTest.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright (C) 2024 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.ravenwood.resapk_test;
+
+
+import static junit.framework.TestCase.assertTrue;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.ravenwood.common.RavenwoodCommonUtils;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+
+@RunWith(AndroidJUnit4.class)
+public class RavenwoodResApkTest {
+    /**
+     * Ensure the file "ravenwood-res.apk" exists.
+     * TODO Check the content of it, once Ravenwood supports resources. The file should
+     * be a copy of RavenwoodResApkTest-apk.apk
+     */
+    @Test
+    public void testResApkExists() {
+        var file = "ravenwood-res-apks/ravenwood-res.apk";
+
+        assertTrue(new File(file).exists());
+    }
+
+    @Test
+    public void testFrameworkResExists() {
+        var file = "ravenwood-data/framework-res.apk";
+
+        assertTrue(new File(
+                RavenwoodCommonUtils.getRavenwoodRuntimePath() + "/" + file).exists());
+    }
+}
diff --git a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java
index 8a2bc1d..f7a59a4b 100644
--- a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java
+++ b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java
@@ -26,7 +26,6 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.util.Slog;
-import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.AccessibilityNodeInfo;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -60,6 +59,7 @@
     private final Set<AccessibilityHierarchyCheck> mHierarchyChecks;
     private final ATFHierarchyBuilder mATFHierarchyBuilder;
     private final Set<AccessibilityCheckResultReported> mCachedResults = new HashSet<>();
+
     @VisibleForTesting
     final A11yCheckerTimer mTimer = new A11yCheckerTimer();
 
@@ -86,10 +86,8 @@
      */
     @RequiresPermission(allOf = {android.Manifest.permission.INTERACT_ACROSS_USERS_FULL})
     public Set<AccessibilityCheckResultReported> maybeRunA11yChecker(
-            List<AccessibilityNodeInfo> nodes,
-            @Nullable AccessibilityEvent accessibilityEvent,
-            ComponentName sourceComponentName,
-            @UserIdInt int userId) {
+            List<AccessibilityNodeInfo> nodes, @Nullable String sourceEventClassName,
+            ComponentName a11yServiceComponentName, @UserIdInt int userId) {
         if (!shouldRunA11yChecker()) {
             return Set.of();
         }
@@ -108,14 +106,13 @@
                 List<AccessibilityHierarchyCheckResult> checkResults = runChecksOnNode(nodeInfo);
                 Set<AccessibilityCheckResultReported> filteredResults =
                         AccessibilityCheckerUtils.processResults(nodeInfo, checkResults,
-                                accessibilityEvent, mPackageManager, sourceComponentName);
+                                sourceEventClassName, mPackageManager, a11yServiceComponentName);
                 allResults.addAll(filteredResults);
             }
             mCachedResults.addAll(allResults);
         } catch (RuntimeException e) {
             Slog.e(LOG_TAG, "An unknown error occurred while running a11y checker.", e);
         }
-
         return allResults;
     }
 
diff --git a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java
index 4171108..fa0fed2 100644
--- a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java
+++ b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java
@@ -22,7 +22,6 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.util.Slog;
-import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.AccessibilityNodeInfo;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -96,7 +95,7 @@
     static Set<AccessibilityCheckResultReported> processResults(
             AccessibilityNodeInfo nodeInfo,
             List<AccessibilityHierarchyCheckResult> checkResults,
-            @Nullable AccessibilityEvent accessibilityEvent,
+            @Nullable String activityClassName,
             PackageManager packageManager,
             ComponentName a11yServiceComponentName) {
         String appPackageName = nodeInfo.getPackageName().toString();
@@ -110,7 +109,8 @@
                     .setPackageName(appPackageName)
                     .setAppVersionCode(getAppVersionCode(packageManager, appPackageName))
                     .setUiElementPath(nodePath)
-                    .setActivityName(getActivityName(packageManager, accessibilityEvent))
+                    .setActivityName(
+                            getActivityName(packageManager, appPackageName, activityClassName))
                     .setWindowTitle(getWindowTitle(nodeInfo))
                     .setSourceComponentName(a11yServiceComponentName.flattenToString())
                     .setSourceVersionCode(
@@ -140,31 +140,23 @@
     }
 
     /**
-     * Returns the simple class name of the Activity providing the cache update, if available,
+     * Returns the simple class name of the Activity associated with the window, if available,
      * or an empty String if not.
      */
     @VisibleForTesting
     static String getActivityName(
-            PackageManager packageManager, @Nullable AccessibilityEvent accessibilityEvent) {
-        if (accessibilityEvent == null) {
+            PackageManager packageManager, String packageName, @Nullable String activityClassName) {
+        if (activityClassName == null) {
             return "";
         }
-        CharSequence activityName = accessibilityEvent.getClassName();
-        if (accessibilityEvent.getEventType() == AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED
-                && accessibilityEvent.getPackageName() != null
-                && activityName != null) {
-            try {
-                // Check class is for a valid Activity.
-                packageManager
-                        .getActivityInfo(
-                                new ComponentName(accessibilityEvent.getPackageName().toString(),
-                                        activityName.toString()), 0);
-                int qualifierEnd = activityName.toString().lastIndexOf('.');
-                return activityName.toString().substring(qualifierEnd + 1);
-            } catch (PackageManager.NameNotFoundException e) {
-                // No need to spam the logs. This is very frequent when the class doesn't match
-                // an activity.
-            }
+        try {
+            // Check class is for a valid Activity.
+            packageManager.getActivityInfo(new ComponentName(packageName, activityClassName), 0);
+            int qualifierEnd = activityClassName.lastIndexOf('.');
+            return activityClassName.substring(qualifierEnd + 1);
+        } catch (PackageManager.NameNotFoundException e) {
+            // No need to spam the logs. This is very frequent when the class doesn't match
+            // an activity.
         }
         return "";
     }
diff --git a/services/autofill/java/com/android/server/autofill/AutofillManagerService.java b/services/autofill/java/com/android/server/autofill/AutofillManagerService.java
index eae516e..9f7fb57 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillManagerService.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillManagerService.java
@@ -1985,12 +1985,13 @@
         }
 
         @Override
-        public void setAutofillFailure(int sessionId, @NonNull List<AutofillId> ids, int userId) {
+        public void setAutofillFailure(
+                int sessionId, @NonNull List<AutofillId> ids, boolean isRefill, int userId) {
             synchronized (mLock) {
                 final AutofillManagerServiceImpl service =
                         peekServiceForUserWithLocalBinderIdentityLocked(userId);
                 if (service != null) {
-                    service.setAutofillFailureLocked(sessionId, getCallingUid(), ids);
+                    service.setAutofillFailureLocked(sessionId, getCallingUid(), ids, isRefill);
                 } else if (sVerbose) {
                     Slog.v(TAG, "setAutofillFailure(): no service for " + userId);
                 }
@@ -2011,6 +2012,46 @@
         }
 
         @Override
+        public void notifyNotExpiringResponseDuringAuth(int sessionId, int userId) {
+            synchronized (mLock) {
+                final AutofillManagerServiceImpl service =
+                        peekServiceForUserWithLocalBinderIdentityLocked(userId);
+                if (service != null) {
+                    service.notifyNotExpiringResponseDuringAuth(sessionId, getCallingUid());
+                } else if (sVerbose) {
+                    Slog.v(TAG, "notifyNotExpiringResponseDuringAuth(): no service for " + userId);
+                }
+            }
+        }
+
+        @Override
+        public void notifyViewEnteredIgnoredDuringAuthCount(int sessionId, int userId) {
+            synchronized (mLock) {
+                final AutofillManagerServiceImpl service =
+                        peekServiceForUserWithLocalBinderIdentityLocked(userId);
+                if (service != null) {
+                    service.notifyViewEnteredIgnoredDuringAuthCount(sessionId, getCallingUid());
+                } else if (sVerbose) {
+                    Slog.v(TAG, "notifyNotExpiringResponseDuringAuth(): no service for " + userId);
+                }
+            }
+        }
+
+        @Override
+        public void setAutofillIdsAttemptedForRefill(
+                int sessionId, @NonNull List<AutofillId> ids, int userId) {
+            synchronized (mLock) {
+                final AutofillManagerServiceImpl service =
+                        peekServiceForUserWithLocalBinderIdentityLocked(userId);
+                if (service != null) {
+                    service.setAutofillIdsAttemptedForRefill(sessionId, ids, getCallingUid());
+                } else if (sVerbose) {
+                    Slog.v(TAG, "setAutofillIdsAttemptedForRefill(): no service for " + userId);
+                }
+            }
+        }
+
+        @Override
         public void finishSession(int sessionId, int userId,
                 @AutofillCommitReason int commitReason) {
             synchronized (mLock) {
diff --git a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
index 2bf319e..c9f8929 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
@@ -464,7 +464,8 @@
     }
 
     @GuardedBy("mLock")
-    void setAutofillFailureLocked(int sessionId, int uid, @NonNull List<AutofillId> ids) {
+    void setAutofillFailureLocked(
+            int sessionId, int uid, @NonNull List<AutofillId> ids, boolean isRefill) {
         if (!isEnabledLocked()) {
             Slog.wtf(TAG, "Service not enabled");
             return;
@@ -474,7 +475,7 @@
             Slog.v(TAG, "setAutofillFailure(): no session for " + sessionId + "(" + uid + ")");
             return;
         }
-        session.setAutofillFailureLocked(ids);
+        session.setAutofillFailureLocked(ids, isRefill);
     }
 
     @GuardedBy("mLock")
@@ -492,6 +493,52 @@
     }
 
     @GuardedBy("mLock")
+    void notifyNotExpiringResponseDuringAuth(int sessionId, int uid) {
+        if (!isEnabledLocked()) {
+            Slog.wtf(TAG, "Service not enabled");
+            return;
+        }
+        final Session session = mSessions.get(sessionId);
+        if (session == null || uid != session.uid) {
+            Slog.v(TAG, "notifyNotExpiringResponseDuringAuth(): no session for "
+                    + sessionId + "(" + uid + ")");
+            return;
+        }
+        session.setNotifyNotExpiringResponseDuringAuth();
+    }
+
+    @GuardedBy("mLock")
+    void notifyViewEnteredIgnoredDuringAuthCount(int sessionId, int uid) {
+        if (!isEnabledLocked()) {
+            Slog.wtf(TAG, "Service not enabled");
+            return;
+        }
+        final Session session = mSessions.get(sessionId);
+        if (session == null || uid != session.uid) {
+            Slog.v(TAG, "notifyViewEnteredIgnoredDuringAuthCount(): no session for "
+                    + sessionId + "(" + uid + ")");
+            return;
+        }
+        session.setLogViewEnteredIgnoredDuringAuth();
+    }
+
+    @GuardedBy("mLock")
+    public void setAutofillIdsAttemptedForRefill(
+            int sessionId, @NonNull List<AutofillId> ids, int uid) {
+        if (!isEnabledLocked()) {
+            Slog.wtf(TAG, "Service not enabled");
+            return;
+        }
+        final Session session = mSessions.get(sessionId);
+        if (session == null || uid != session.uid) {
+            Slog.v(TAG, "setAutofillIdsAttemptedForRefill(): no session for "
+                    + sessionId + "(" + uid + ")");
+            return;
+        }
+        session.setAutofillIdsAttemptedForRefillLocked(ids);
+    }
+
+    @GuardedBy("mLock")
     void finishSessionLocked(int sessionId, int uid, @AutofillCommitReason int commitReason) {
         if (!isEnabledLocked()) {
             Slog.wtf(TAG, "Service not enabled");
diff --git a/services/autofill/java/com/android/server/autofill/PresentationStatsEventLogger.java b/services/autofill/java/com/android/server/autofill/PresentationStatsEventLogger.java
index 49ca297..930af5e 100644
--- a/services/autofill/java/com/android/server/autofill/PresentationStatsEventLogger.java
+++ b/services/autofill/java/com/android/server/autofill/PresentationStatsEventLogger.java
@@ -58,6 +58,7 @@
 import static com.android.server.autofill.Helper.sVerbose;
 
 import android.annotation.IntDef;
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.ComponentName;
 import android.content.Context;
@@ -685,6 +686,19 @@
     }
 
     /**
+     * Set views_fillable_total_count as long as mEventInternal presents.
+     */
+    public void maybeUpdateViewFillablesForRefillAttempt(List<AutofillId> autofillIds) {
+        mEventInternal.ifPresent(event -> {
+            // These autofill ids would be the ones being re-attempted.
+            event.mAutofillIdsAttemptedAutofill = new ArraySet<>(autofillIds);
+            // These autofill id's are being refilled, so they had failed previously.
+            // Note that these autofillIds correspond to the new autofill ids after relayout.
+            event.mFailedAutofillIds = new ArraySet<>(autofillIds);
+        });
+    }
+
+    /**
      * Set how many views are filtered from fill because they are not in current session
      */
     public void maybeSetFilteredFillableViewsCount(int filteredViewsCount) {
@@ -697,9 +711,16 @@
      * Set views_filled_failure_count using failure count as long as mEventInternal
      * presents.
      */
-    public void maybeSetViewFillFailureCounts(int failureCount) {
+    public void maybeSetViewFillFailureCounts(@NonNull List<AutofillId> ids, boolean isRefill) {
         mEventInternal.ifPresent(event -> {
-            event.mViewFillFailureCount = failureCount;
+            int failureCount = ids.size();
+            if (isRefill) {
+                event.mViewFailedOnRefillCount = failureCount;
+            } else {
+                event.mViewFillFailureCount = failureCount;
+                event.mViewFailedPriorToRefillCount = failureCount;
+                event.mFailedAutofillIds = new ArraySet<>(ids);
+            }
         });
     }
 
@@ -719,7 +740,7 @@
      * Set views_filled_failure_count using failure count as long as mEventInternal
      * presents.
      */
-    public void maybeAddSuccessId(AutofillId autofillId) {
+    public synchronized void maybeAddSuccessId(AutofillId autofillId) {
         mEventInternal.ifPresent(event -> {
             ArraySet<AutofillId> autofillIds = event.mAutofillIdsAttemptedAutofill;
             if (autofillIds == null) {
@@ -727,9 +748,21 @@
                         + " successfully filled");
                 event.mViewFilledButUnexpectedCount++;
             } else if (autofillIds.contains(autofillId)) {
-                if (sVerbose) {
-                    Slog.v(TAG, "Logging autofill for id:" + autofillId);
+                ArraySet<AutofillId> failedIds = event.mFailedAutofillIds;
+                if (failedIds.contains(autofillId)) {
+                    if (sVerbose) {
+                        Slog.v(TAG, "Logging autofill refill of id:" + autofillId);
+                    }
+                    // This indicates the success after refill attempt
+                    event.mViewFilledSuccessfullyOnRefillCount++;
+                    // Remove so if we don't reprocess duplicate requests
+                    failedIds.remove(autofillId);
+                } else {
+                    if (sVerbose) {
+                        Slog.v(TAG, "Logging autofill for id:" + autofillId);
+                    }
                 }
+                // Common actions to take irrespective of being filled by refill attempt or not.
                 event.mViewFillSuccessCount++;
                 autofillIds.remove(autofillId);
                 event.mAlreadyFilledAutofillIds.add(autofillId);
@@ -746,6 +779,23 @@
         });
     }
 
+    /**
+     * Set how many views are filtered from fill because they are not in current session
+     */
+    public void maybeSetNotifyNotExpiringResponseDuringAuth() {
+        mEventInternal.ifPresent(event -> {
+            event.mFixExpireResponseDuringAuthCount++;
+        });
+    }
+    /**
+     * Set how many views are filtered from fill because they are not in current session
+     */
+    public void notifyViewEnteredIgnoredDuringAuthCount() {
+        mEventInternal.ifPresent(event -> {
+            event.mNotifyViewEnteredIgnoredDuringAuthCount++;
+        });
+    }
+
     public void logAndEndEvent() {
         if (!mEventInternal.isPresent()) {
             Slog.w(TAG, "Shouldn't be logging AutofillPresentationEventReported again for same "
@@ -933,6 +983,7 @@
         int mNotifyViewEnteredIgnoredDuringAuthCount = 0;
 
         ArraySet<AutofillId> mAutofillIdsAttemptedAutofill;
+        ArraySet<AutofillId> mFailedAutofillIds = new ArraySet<>();
         ArraySet<AutofillId> mAlreadyFilledAutofillIds = new ArraySet<>();
 
         // Not logged - used for internal logic
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index 21df7a5..b7508b4 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -5360,6 +5360,8 @@
             saveTriggerId = null;
         }
 
+        boolean hasAuthentication = (response.getAuthentication() != null);
+
         // Must also track that are part of datasets, otherwise the FillUI won't be hidden when
         // they go away (if they're not savable).
 
@@ -5379,6 +5381,9 @@
                         }
                     }
                 }
+                if (dataset.getAuthentication() != null) {
+                    hasAuthentication = true;
+                }
             }
         }
 
@@ -5390,7 +5395,7 @@
                         + " hasSaveInfo: " + (saveInfo != null));
             }
             mClient.setTrackedViews(id, toArray(trackedViews), mSaveOnAllViewsInvisible,
-                    saveOnFinish, toArray(fillableIds), saveTriggerId);
+                    saveOnFinish, toArray(fillableIds), saveTriggerId, hasAuthentication);
         } catch (RemoteException e) {
             Slog.w(TAG, "Cannot set tracked ids", e);
         }
@@ -5400,7 +5405,7 @@
      * Sets the state of views that failed to autofill.
      */
     @GuardedBy("mLock")
-    void setAutofillFailureLocked(@NonNull List<AutofillId> ids) {
+    void setAutofillFailureLocked(@NonNull List<AutofillId> ids, boolean isRefill) {
         if (sVerbose && !ids.isEmpty()) {
             Slog.v(TAG, "Total views that failed to populate: " + ids.size());
         }
@@ -5418,7 +5423,7 @@
                 Slog.v(TAG, "Changed state of " + id + " to " + viewState.getStateAsString());
             }
         }
-        mPresentationStatsEventLogger.maybeSetViewFillFailureCounts(ids.size());
+        mPresentationStatsEventLogger.maybeSetViewFillFailureCounts(ids, isRefill);
     }
 
     /**
@@ -5435,6 +5440,23 @@
         mPresentationStatsEventLogger.maybeAddSuccessId(id);
     }
 
+    /**
+     * Sets the state of views that failed to autofill.
+     */
+    void setNotifyNotExpiringResponseDuringAuth() {
+        synchronized (mLock) {
+            mPresentationStatsEventLogger.maybeSetNotifyNotExpiringResponseDuringAuth();
+        }
+    }
+    /**
+     * Sets the state of views that failed to autofill.
+     */
+    void setLogViewEnteredIgnoredDuringAuth() {
+        synchronized (mLock) {
+            mPresentationStatsEventLogger.notifyViewEnteredIgnoredDuringAuthCount();
+        }
+    }
+
     @GuardedBy("mLock")
     private void replaceResponseLocked(@NonNull FillResponse oldResponse,
             @NonNull FillResponse newResponse, @Nullable Bundle newClientState) {
@@ -6665,6 +6687,11 @@
         }
     }
 
+    @GuardedBy("mLock")
+    public void setAutofillIdsAttemptedForRefillLocked(@NonNull List<AutofillId> ids) {
+        mPresentationStatsEventLogger.maybeUpdateViewFillablesForRefillAttempt(ids);
+    }
+
     private AutoFillUI getUiForShowing() {
         synchronized (mLock) {
             mUi.setCallback(this);
diff --git a/services/core/java/com/android/server/MmsServiceBroker.java b/services/core/java/com/android/server/MmsServiceBroker.java
index ced7773..11de258 100644
--- a/services/core/java/com/android/server/MmsServiceBroker.java
+++ b/services/core/java/com/android/server/MmsServiceBroker.java
@@ -130,17 +130,18 @@
         }
 
         @Override
-        public void sendMessage(int subId, String callingPkg, Uri contentUri, String locationUrl,
-                Bundle configOverrides, PendingIntent sentIntent, long messageId,
+        public void sendMessage(int subId, int callingUser, String callingPkg,
+                Uri contentUri, String locationUrl, Bundle configOverrides,
+                PendingIntent sentIntent, long messageId,
                 String attributionTag) throws RemoteException {
             returnPendingIntentWithError(sentIntent);
         }
 
         @Override
-        public void downloadMessage(int subId, String callingPkg, String locationUrl,
-                Uri contentUri, Bundle configOverrides, PendingIntent downloadedIntent,
-                long messageId, String attributionTag)
-                throws RemoteException {
+        public void downloadMessage(int subId, int callingUser, String callingPkg,
+                String locationUrl, Uri contentUri, Bundle configOverrides,
+                PendingIntent downloadedIntent,
+                long messageId, String attributionTag) throws RemoteException {
             returnPendingIntentWithError(downloadedIntent);
         }
 
@@ -151,8 +152,9 @@
         }
 
         @Override
-        public Uri importMultimediaMessage(String callingPkg, Uri contentUri, String messageId,
-                long timestampSecs, boolean seen, boolean read) throws RemoteException {
+        public Uri importMultimediaMessage(int callingUser, String callingPkg,
+                Uri contentUri, String messageId, long timestampSecs,
+                boolean seen, boolean read) throws RemoteException {
             return null;
         }
 
@@ -187,8 +189,8 @@
         }
 
         @Override
-        public Uri addMultimediaMessageDraft(String callingPkg, Uri contentUri)
-                throws RemoteException {
+        public Uri addMultimediaMessageDraft(int callingUser, String callingPkg,
+                Uri contentUri) throws RemoteException {
             return null;
         }
 
@@ -333,9 +335,9 @@
         private static final String PHONE_PACKAGE_NAME = "com.android.phone";
 
         @Override
-        public void sendMessage(int subId, String callingPkg, Uri contentUri,
-                String locationUrl, Bundle configOverrides, PendingIntent sentIntent,
-                long messageId, String attributionTag)
+        public void sendMessage(int subId, int callingUser, String callingPkg,
+                Uri contentUri, String locationUrl, Bundle configOverrides,
+                PendingIntent sentIntent, long messageId, String attributionTag)
                 throws RemoteException {
             Slog.d(TAG, "sendMessage() by " + callingPkg);
             mContext.enforceCallingPermission(Manifest.permission.SEND_SMS, "Send MMS message");
@@ -360,14 +362,15 @@
                     CarrierMessagingService.SERVICE_INTERFACE,
                     Intent.FLAG_GRANT_READ_URI_PERMISSION,
                     subId);
-            getServiceGuarded().sendMessage(subId, callingPkg, contentUri, locationUrl,
-                    configOverrides, sentIntent, messageId, attributionTag);
+            getServiceGuarded().sendMessage(subId, getCallingUserId(), callingPkg, contentUri,
+                    locationUrl, configOverrides, sentIntent, messageId, attributionTag);
         }
 
         @Override
-        public void downloadMessage(int subId, String callingPkg, String locationUrl,
-                Uri contentUri, Bundle configOverrides, PendingIntent downloadedIntent,
-                long messageId, String attributionTag) throws RemoteException {
+        public void downloadMessage(int subId, int callingUser, String callingPkg,
+                String locationUrl, Uri contentUri, Bundle configOverrides,
+                PendingIntent downloadedIntent, long messageId, String attributionTag)
+                throws RemoteException {
             Slog.d(TAG, "downloadMessage() by " + callingPkg);
             mContext.enforceCallingPermission(Manifest.permission.RECEIVE_MMS,
                     "Download MMS message");
@@ -381,8 +384,8 @@
                     Intent.FLAG_GRANT_READ_URI_PERMISSION | Intent.FLAG_GRANT_WRITE_URI_PERMISSION,
                     subId);
 
-            getServiceGuarded().downloadMessage(subId, callingPkg, locationUrl, contentUri,
-                    configOverrides, downloadedIntent, messageId, attributionTag);
+            getServiceGuarded().downloadMessage(subId, getCallingUserId(), callingPkg, locationUrl,
+                    contentUri, configOverrides, downloadedIntent, messageId, attributionTag);
         }
 
         @Override
@@ -399,8 +402,8 @@
         }
 
         @Override
-        public Uri importMultimediaMessage(String callingPkg, Uri contentUri,
-                String messageId, long timestampSecs, boolean seen, boolean read)
+        public Uri importMultimediaMessage(int callingUser, String callingPkg,
+                Uri contentUri, String messageId, long timestampSecs, boolean seen, boolean read)
                 throws RemoteException {
             if (getAppOpsManager().noteOp(AppOpsManager.OP_WRITE_SMS, Binder.getCallingUid(),
                     callingPkg, null, null) != AppOpsManager.MODE_ALLOWED) {
@@ -408,8 +411,8 @@
                 // while writing the TelephonyProvider
                 return FAKE_MMS_SENT_URI;
             }
-            return getServiceGuarded().importMultimediaMessage(
-                    callingPkg, contentUri, messageId, timestampSecs, seen, read);
+            return getServiceGuarded().importMultimediaMessage(getCallingUserId(), callingPkg,
+                    contentUri, messageId, timestampSecs, seen, read);
         }
 
         @Override
@@ -467,15 +470,16 @@
         }
 
         @Override
-        public Uri addMultimediaMessageDraft(String callingPkg, Uri contentUri)
-                throws RemoteException {
+        public Uri addMultimediaMessageDraft(int callingUser, String callingPkg,
+                Uri contentUri) throws RemoteException {
             if (getAppOpsManager().noteOp(AppOpsManager.OP_WRITE_SMS, Binder.getCallingUid(),
                     callingPkg, null, null) != AppOpsManager.MODE_ALLOWED) {
                 // Silently fail AppOps failure due to not being the default SMS app
                 // while writing the TelephonyProvider
                 return FAKE_MMS_DRAFT_URI;
             }
-            return getServiceGuarded().addMultimediaMessageDraft(callingPkg, contentUri);
+            return getServiceGuarded().addMultimediaMessageDraft(getCallingUserId(), callingPkg,
+                    contentUri);
         }
 
         @Override
@@ -572,4 +576,13 @@
         if (info == null) return INVALID_SIM_SLOT_INDEX;
         return info.getSimSlotIndex();
     }
+
+    /**
+     * Retrieves the  calling user id.
+     * @return The id of the calling user.
+     */
+    private int getCallingUserId() {
+        return Binder.getCallingUserHandle().getIdentifier();
+    }
+
 }
diff --git a/services/core/java/com/android/server/TelephonyRegistry.java b/services/core/java/com/android/server/TelephonyRegistry.java
index 3633d0f..33cf842 100644
--- a/services/core/java/com/android/server/TelephonyRegistry.java
+++ b/services/core/java/com/android/server/TelephonyRegistry.java
@@ -55,6 +55,7 @@
 import android.telephony.BarringInfo;
 import android.telephony.CallQuality;
 import android.telephony.CallState;
+import android.telephony.CarrierConfigManager;
 import android.telephony.CellIdentity;
 import android.telephony.CellInfo;
 import android.telephony.CellSignalStrength;
@@ -426,6 +427,7 @@
     private boolean[] mSCBMStarted;
 
     private boolean[] mCarrierRoamingNtnMode = null;
+    private boolean[] mCarrierRoamingNtnEligible = null;
 
     /**
      * Per-phone map of precise data connection state. The key of the map is the pair of transport
@@ -726,6 +728,7 @@
             mSCBMReason = copyOf(mSCBMReason, mNumPhones);
             mSCBMStarted = copyOf(mSCBMStarted, mNumPhones);
             mCarrierRoamingNtnMode = copyOf(mCarrierRoamingNtnMode, mNumPhones);
+            mCarrierRoamingNtnEligible = copyOf(mCarrierRoamingNtnEligible, mNumPhones);
             // ds -> ss switch.
             if (mNumPhones < oldNumPhones) {
                 cutListToSize(mCellInfo, mNumPhones);
@@ -785,6 +788,7 @@
                 mSCBMReason[i] = TelephonyManager.STOP_REASON_UNKNOWN;
                 mSCBMStarted[i] = false;
                 mCarrierRoamingNtnMode[i] = false;
+                mCarrierRoamingNtnEligible[i] = false;
             }
         }
     }
@@ -859,6 +863,7 @@
         mSCBMReason = new int[numPhones];
         mSCBMStarted = new boolean[numPhones];
         mCarrierRoamingNtnMode = new boolean[numPhones];
+        mCarrierRoamingNtnEligible = new boolean[numPhones];
 
         for (int i = 0; i < numPhones; i++) {
             mCallState[i] =  TelephonyManager.CALL_STATE_IDLE;
@@ -903,6 +908,7 @@
             mSCBMReason[i] = TelephonyManager.STOP_REASON_UNKNOWN;
             mSCBMStarted[i] = false;
             mCarrierRoamingNtnMode[i] = false;
+            mCarrierRoamingNtnEligible[i] = false;
         }
 
         mAppOps = mContext.getSystemService(AppOpsManager.class);
@@ -1518,6 +1524,15 @@
                         remove(r.binder);
                     }
                 }
+                if (events.contains(
+                        TelephonyCallback.EVENT_CARRIER_ROAMING_NTN_ELIGIBLE_STATE_CHANGED)) {
+                    try {
+                        r.callback.onCarrierRoamingNtnEligibleStateChanged(
+                                mCarrierRoamingNtnEligible[r.phoneId]);
+                    } catch (RemoteException ex) {
+                        remove(r.binder);
+                    }
+                }
             }
         }
     }
@@ -3536,6 +3551,53 @@
         }
     }
 
+    /**
+     * Notify external listeners that device eligibility to connect to carrier roaming
+     * non-terrestrial network changed.
+     *
+     * @param subId subscription ID.
+     * @param eligible {@code true} when the device is eligible for satellite
+     * communication if all the following conditions are met:
+     * <ul>
+     * <li>Any subscription on the device supports P2P satellite messaging which is defined by
+     * {@link CarrierConfigManager#KEY_SATELLITE_ATTACH_SUPPORTED_BOOL} </li>
+     * <li>{@link CarrierConfigManager#KEY_CARRIER_ROAMING_NTN_CONNECT_TYPE_INT} set to
+     * {@link CarrierConfigManager#CARRIER_ROAMING_NTN_CONNECT_MANUAL} </li>
+     * <li>The device is in {@link ServiceState#STATE_OUT_OF_SERVICE}, not connected to Wi-Fi,
+     * and the hysteresis timer defined by {@link CarrierConfigManager
+     * #KEY_CARRIER_SUPPORTED_SATELLITE_NOTIFICATION_HYSTERESIS_SEC_INT} is expired. </li>
+     * </ul>
+     */
+    public void notifyCarrierRoamingNtnEligibleStateChanged(int subId, boolean eligible) {
+        if (!checkNotifyPermission("notifyCarrierRoamingNtnEligibleStateChanged")) {
+            log("notifyCarrierRoamingNtnEligibleStateChanged: caller does not have required "
+                    + "permissions.");
+            return;
+        }
+
+        if (VDBG) {
+            log("notifyCarrierRoamingNtnEligibleStateChanged: "
+                    + "subId=" + subId + " eligible" + eligible);
+        }
+
+        synchronized (mRecords) {
+            int phoneId = getPhoneIdFromSubId(subId);
+            mCarrierRoamingNtnEligible[phoneId] = eligible;
+            for (Record r : mRecords) {
+                if (r.matchTelephonyCallbackEvent(
+                        TelephonyCallback.EVENT_CARRIER_ROAMING_NTN_ELIGIBLE_STATE_CHANGED)
+                        && idMatch(r, subId, phoneId)) {
+                    try {
+                        r.callback.onCarrierRoamingNtnEligibleStateChanged(eligible);
+                    } catch (RemoteException ex) {
+                        mRemoveList.add(r.binder);
+                    }
+                }
+            }
+            handleRemoveListLocked();
+        }
+    }
+
     @NeverCompile // Avoid size overhead of debugging code.
     @Override
     public void dump(FileDescriptor fd, PrintWriter writer, String[] args) {
@@ -3589,6 +3651,8 @@
                 pw.println("mECBMStarted=" + mECBMStarted[i]);
                 pw.println("mSCBMReason=" + mSCBMReason[i]);
                 pw.println("mSCBMStarted=" + mSCBMStarted[i]);
+                pw.println("mCarrierRoamingNtnMode=" + mCarrierRoamingNtnMode[i]);
+                pw.println("mCarrierRoamingNtnEligible=" + mCarrierRoamingNtnEligible[i]);
 
                 // We need to obfuscate package names, and primitive arrays' native toString is ugly
                 Pair<List<String>, int[]> carrierPrivilegeState = mCarrierPrivilegeStates.get(i);
diff --git a/services/core/java/com/android/server/am/ActiveServices.java b/services/core/java/com/android/server/am/ActiveServices.java
index 0fa5260..2e1416b 100644
--- a/services/core/java/com/android/server/am/ActiveServices.java
+++ b/services/core/java/com/android/server/am/ActiveServices.java
@@ -6653,9 +6653,10 @@
 
             // If unbound while waiting to start and there is no connection left in this service,
             // remove the pending service
-            if (s.getConnections().isEmpty()) {
+            if (s.getConnections().isEmpty() && !s.startRequested) {
                 mPendingServices.remove(s);
                 mPendingBringups.remove(s);
+                if (DEBUG_SERVICE) Slog.v(TAG_SERVICE, "Removed pending service: " + s);
             }
 
             if (c.hasFlag(Context.BIND_AUTO_CREATE)) {
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 8df4e77..03fbfd37 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -124,6 +124,7 @@
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.AggregatedPowerStatsConfig;
+import com.android.server.power.stats.AmbientDisplayPowerStatsProcessor;
 import com.android.server.power.stats.AudioPowerStatsProcessor;
 import com.android.server.power.stats.BatteryExternalStatsWorker;
 import com.android.server.power.stats.BatteryStatsDumpHelperImpl;
@@ -142,6 +143,8 @@
 import com.android.server.power.stats.PowerStatsScheduler;
 import com.android.server.power.stats.PowerStatsStore;
 import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.ScreenPowerStatsProcessor;
+import com.android.server.power.stats.SensorPowerStatsProcessor;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 import com.android.server.power.stats.VideoPowerStatsProcessor;
 import com.android.server.power.stats.WifiPowerStatsProcessor;
@@ -488,6 +491,20 @@
                 .setProcessor(
                         new CpuPowerStatsProcessor(mPowerProfile, mCpuScalingPolicies));
 
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .setProcessor(
+                        new ScreenPowerStatsProcessor(mPowerProfile));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .setProcessor(new AmbientDisplayPowerStatsProcessor());
+
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .trackDeviceStates(
                         AggregatedPowerStatsConfig.STATE_POWER,
@@ -579,6 +596,17 @@
                 .setProcessor(
                         new GnssPowerStatsProcessor(mPowerProfile, mPowerStatsUidResolver));
 
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SENSORS)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessor(new SensorPowerStatsProcessor(
+                        () -> mContext.getSystemService(SensorManager.class)));
+
         config.trackCustomPowerComponents(CustomEnergyConsumerPowerStatsProcessor::new)
                 .trackDeviceStates(
                         AggregatedPowerStatsConfig.STATE_POWER,
@@ -636,6 +664,18 @@
                 BatteryConsumer.POWER_COMPONENT_CPU,
                 Flags.streamlinedBatteryStats());
 
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_SCREEN,
+                Flags.streamlinedMiscBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_SCREEN,
+                Flags.streamlinedMiscBatteryStats());
+
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                Flags.streamlinedMiscBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                Flags.streamlinedMiscBatteryStats());
+
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                 Flags.streamlinedConnectivityBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
@@ -678,6 +718,10 @@
                 BatteryConsumer.POWER_COMPONENT_GNSS,
                 Flags.streamlinedMiscBatteryStats());
 
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_SENSORS,
+                Flags.streamlinedMiscBatteryStats());
+
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CAMERA,
                 Flags.streamlinedMiscBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
diff --git a/services/core/java/com/android/server/am/ContentProviderHelper.java b/services/core/java/com/android/server/am/ContentProviderHelper.java
index 4ff1367..afb7bb4 100644
--- a/services/core/java/com/android/server/am/ContentProviderHelper.java
+++ b/services/core/java/com/android/server/am/ContentProviderHelper.java
@@ -538,6 +538,8 @@
                             if (!pr.hasProvider(cpi.name)) {
                                 checkTime(startTime, "getContentProviderImpl: scheduling install");
                                 pr.installProvider(cpi.name, cpr);
+                                mService.mOomAdjuster.unfreezeTemporarily(proc,
+                                        CachedAppOptimizer.UNFREEZE_REASON_GET_PROVIDER);
                                 try {
                                     thread.scheduleInstallProvider(cpi);
                                 } catch (RemoteException e) {
diff --git a/services/core/java/com/android/server/appop/AttributedOp.java b/services/core/java/com/android/server/appop/AttributedOp.java
index 8cf47d0..430be03 100644
--- a/services/core/java/com/android/server/appop/AttributedOp.java
+++ b/services/core/java/com/android/server/appop/AttributedOp.java
@@ -109,7 +109,7 @@
                 uidState, flags);
 
         mAppOpsService.mHistoricalRegistry.incrementOpAccessedCount(parent.op, parent.uid,
-                parent.packageName, tag, uidState, flags, accessTime,
+                parent.packageName, persistentDeviceId, tag, uidState, flags, accessTime,
                 AppOpsManager.ATTRIBUTION_FLAGS_NONE, AppOpsManager.ATTRIBUTION_CHAIN_ID_NONE);
     }
 
@@ -253,8 +253,8 @@
 
         if (isStarted) {
             mAppOpsService.mHistoricalRegistry.incrementOpAccessedCount(parent.op, parent.uid,
-                    parent.packageName, tag, uidState, flags, startTime, attributionFlags,
-                    attributionChainId);
+                    parent.packageName, persistentDeviceId, tag, uidState, flags, startTime,
+                    attributionFlags, attributionChainId);
         }
     }
 
@@ -333,7 +333,7 @@
                     finishedEvent);
 
             mAppOpsService.mHistoricalRegistry.increaseOpAccessDuration(parent.op, parent.uid,
-                    parent.packageName, tag, event.getUidState(),
+                    parent.packageName, persistentDeviceId, tag, event.getUidState(),
                     event.getFlags(), finishedEvent.getNoteTime(), finishedEvent.getDuration(),
                     event.getAttributionFlags(), event.getAttributionChainId());
 
@@ -441,8 +441,9 @@
             event.setStartElapsedTime(SystemClock.elapsedRealtime());
             event.setStartTime(startTime);
             mAppOpsService.mHistoricalRegistry.incrementOpAccessedCount(parent.op, parent.uid,
-                    parent.packageName, tag, event.getUidState(), event.getFlags(), startTime,
-                    event.getAttributionFlags(), event.getAttributionChainId());
+                    parent.packageName, persistentDeviceId, tag, event.getUidState(),
+                    event.getFlags(), startTime, event.getAttributionFlags(),
+                    event.getAttributionChainId());
             if (shouldSendActive) {
                 mAppOpsService.scheduleOpActiveChangedIfNeededLocked(parent.op, parent.uid,
                         parent.packageName, tag, event.getVirtualDeviceId(), true,
diff --git a/services/core/java/com/android/server/appop/DiscreteRegistry.java b/services/core/java/com/android/server/appop/DiscreteRegistry.java
index 5d83ad6..539dbca 100644
--- a/services/core/java/com/android/server/appop/DiscreteRegistry.java
+++ b/services/core/java/com/android/server/appop/DiscreteRegistry.java
@@ -41,6 +41,7 @@
 import static android.app.AppOpsManager.OP_RESERVED_FOR_TESTING;
 import static android.app.AppOpsManager.flagsToString;
 import static android.app.AppOpsManager.getUidStateName;
+import static android.companion.virtual.VirtualDeviceManager.PERSISTENT_DEVICE_ID_DEFAULT;
 
 import static java.lang.Long.min;
 import static java.lang.Math.max;
@@ -52,6 +53,7 @@
 import android.os.Build;
 import android.os.Environment;
 import android.os.FileUtils;
+import android.permission.flags.Flags;
 import android.provider.DeviceConfig;
 import android.util.ArrayMap;
 import android.util.AtomicFile;
@@ -76,7 +78,7 @@
 import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
+import java.util.Comparator;
 import java.util.Date;
 import java.util.List;
 import java.util.Objects;
@@ -164,6 +166,8 @@
     private static final String TAG_OP = "o";
     private static final String ATTR_OP_ID = "op";
 
+    private static final String ATTR_DEVICE_ID = "di";
+
     private static final String TAG_TAG = "a";
     private static final String ATTR_TAG = "at";
 
@@ -196,11 +200,14 @@
     private boolean mDebugMode = false;
 
     DiscreteRegistry(Object inMemoryLock) {
+        this(inMemoryLock, new File(new File(Environment.getDataSystemDirectory(), "appops"),
+                "discrete"));
+    }
+
+    DiscreteRegistry(Object inMemoryLock, File discreteAccessDir) {
         mInMemoryLock = inMemoryLock;
         synchronized (mOnDiskLock) {
-            mDiscreteAccessDir = new File(
-                    new File(Environment.getDataSystemDirectory(), "appops"),
-                    "discrete");
+            mDiscreteAccessDir = discreteAccessDir;
             createDiscreteAccessDirLocked();
             int largestChainId = readLargestChainIdFromDiskLocked();
             synchronized (mInMemoryLock) {
@@ -245,16 +252,24 @@
                 DEFAULT_DISCRETE_OPS);
     }
 
-    void recordDiscreteAccess(int uid, String packageName, int op, @Nullable String attributionTag,
-            @AppOpsManager.OpFlags int flags, @AppOpsManager.UidState int uidState, long accessTime,
-            long accessDuration, @AppOpsManager.AttributionFlags int attributionFlags,
-            int attributionChainId) {
+    void recordDiscreteAccess(int uid, String packageName, @NonNull String deviceId, int op,
+            @Nullable String attributionTag, @AppOpsManager.OpFlags int flags,
+            @AppOpsManager.UidState int uidState, long accessTime, long accessDuration,
+            @AppOpsManager.AttributionFlags int attributionFlags, int attributionChainId) {
         if (!isDiscreteOp(op, flags)) {
             return;
         }
         synchronized (mInMemoryLock) {
-            mDiscreteOps.addDiscreteAccess(op, uid, packageName, attributionTag, flags, uidState,
-                    accessTime, accessDuration, attributionFlags, attributionChainId);
+            if (Flags.deviceAwareAppOpNewSchemaEnabled()) {
+                mDiscreteOps.addDiscreteAccess(op, uid, packageName, deviceId, attributionTag,
+                        flags, uidState, accessTime, accessDuration, attributionFlags,
+                        attributionChainId);
+            } else {
+                mDiscreteOps.addDiscreteAccess(op, uid, packageName, PERSISTENT_DEVICE_ID_DEFAULT,
+                        attributionTag, flags, uidState, accessTime, accessDuration,
+                        attributionFlags, attributionChainId);
+            }
+
         }
     }
 
@@ -294,7 +309,6 @@
         discreteOps.filter(beginTimeMillis, endTimeMillis, filter, uidFilter, packageNameFilter,
                 opNamesFilter, attributionTagFilter, flagsFilter, attributionChains);
         discreteOps.applyToHistoricalOps(result, attributionChains);
-        return;
     }
 
     private int readLargestChainIdFromDiskLocked() {
@@ -355,28 +369,36 @@
                 String pkg = pkgs.keyAt(pkgNum);
                 int nOps = ops.size();
                 for (int opNum = 0; opNum < nOps; opNum++) {
-                    ArrayMap<String, List<DiscreteOpEvent>> attrOps =
-                            ops.valueAt(opNum).mAttributedOps;
                     int op = ops.keyAt(opNum);
-                    int nAttrOps = attrOps.size();
-                    for (int attrOpNum = 0; attrOpNum < nAttrOps; attrOpNum++) {
-                        List<DiscreteOpEvent> opEvents = attrOps.valueAt(attrOpNum);
-                        String attributionTag = attrOps.keyAt(attrOpNum);
-                        int nOpEvents = opEvents.size();
-                        for (int opEventNum = 0; opEventNum < nOpEvents; opEventNum++) {
-                            DiscreteOpEvent event = opEvents.get(opEventNum);
-                            if (event == null
-                                    || event.mAttributionChainId == ATTRIBUTION_CHAIN_ID_NONE
-                                    || (event.mAttributionFlags & ATTRIBUTION_FLAG_TRUSTED) == 0) {
-                                continue;
-                            }
+                    ArrayMap<String, DiscreteDeviceOp> deviceOps =
+                            ops.valueAt(opNum).mDeviceAttributedOps;
 
-                            if (!chains.containsKey(event.mAttributionChainId)) {
-                                chains.put(event.mAttributionChainId,
-                                        new AttributionChain(attributionExemptPkgs));
+                    int nDeviceOps = deviceOps.size();
+                    for (int deviceNum = 0; deviceNum < nDeviceOps; deviceNum++) {
+                        ArrayMap<String, List<DiscreteOpEvent>> attrOps =
+                                deviceOps.valueAt(deviceNum).mAttributedOps;
+
+                        int nAttrOps = attrOps.size();
+                        for (int attrOpNum = 0; attrOpNum < nAttrOps; attrOpNum++) {
+                            List<DiscreteOpEvent> opEvents = attrOps.valueAt(attrOpNum);
+                            String attributionTag = attrOps.keyAt(attrOpNum);
+                            int nOpEvents = opEvents.size();
+                            for (int opEventNum = 0; opEventNum < nOpEvents; opEventNum++) {
+                                DiscreteOpEvent event = opEvents.get(opEventNum);
+                                if (event == null
+                                        || event.mAttributionChainId == ATTRIBUTION_CHAIN_ID_NONE
+                                        || (event.mAttributionFlags & ATTRIBUTION_FLAG_TRUSTED)
+                                                == 0) {
+                                    continue;
+                                }
+
+                                if (!chains.containsKey(event.mAttributionChainId)) {
+                                    chains.put(event.mAttributionChainId,
+                                            new AttributionChain(attributionExemptPkgs));
+                                }
+                                chains.get(event.mAttributionChainId)
+                                        .addEvent(pkg, uid, attributionTag, op, event);
                             }
-                            chains.get(event.mAttributionChainId)
-                                    .addEvent(pkg, uid, attributionTag, op, event);
                         }
                     }
                 }
@@ -464,7 +486,7 @@
         createDiscreteAccessDir();
     }
 
-    private DiscreteOps getAllDiscreteOps() {
+    DiscreteOps getAllDiscreteOps() {
         DiscreteOps discreteOps = new DiscreteOps(0);
 
         synchronized (mOnDiskLock) {
@@ -608,7 +630,7 @@
         }
     }
 
-    private final class DiscreteOps {
+    static final class DiscreteOps {
         ArrayMap<Integer, DiscreteUidOps> mUids;
         int mChainIdOffset;
         int mLargestChainId;
@@ -634,8 +656,9 @@
         }
 
         void addDiscreteAccess(int op, int uid, @NonNull String packageName,
-                @Nullable String attributionTag, @AppOpsManager.OpFlags int flags,
-                @AppOpsManager.UidState int uidState, long accessTime, long accessDuration,
+                @NonNull String deviceId, @Nullable String attributionTag,
+                @AppOpsManager.OpFlags int flags, @AppOpsManager.UidState int uidState,
+                long accessTime, long accessDuration,
                 @AppOpsManager.AttributionFlags int attributionFlags, int attributionChainId) {
             int offsetChainId = attributionChainId;
             if (attributionChainId != ATTRIBUTION_CHAIN_ID_NONE) {
@@ -649,8 +672,9 @@
                     mChainIdOffset = -1 * attributionChainId;
                 }
             }
-            getOrCreateDiscreteUidOps(uid).addDiscreteAccess(op, packageName, attributionTag, flags,
-                    uidState, accessTime, accessDuration, attributionFlags, offsetChainId);
+            getOrCreateDiscreteUidOps(uid).addDiscreteAccess(op, packageName, deviceId,
+                    attributionTag, flags, uidState, accessTime, accessDuration, attributionFlags,
+                    offsetChainId);
         }
 
         private void filter(long beginTimeMillis, long endTimeMillis,
@@ -837,7 +861,7 @@
         }
     }
 
-    private final class DiscreteUidOps {
+    static final class DiscreteUidOps {
         ArrayMap<String, DiscretePackageOps> mPackages;
 
         DiscreteUidOps() {
@@ -889,12 +913,13 @@
             mPackages.remove(packageName);
         }
 
-        void addDiscreteAccess(int op, @NonNull String packageName, @Nullable String attributionTag,
-                @AppOpsManager.OpFlags int flags, @AppOpsManager.UidState int uidState,
-                long accessTime, long accessDuration,
+        void addDiscreteAccess(int op, @NonNull String packageName, @NonNull String deviceId,
+                @Nullable String attributionTag, @AppOpsManager.OpFlags int flags,
+                @AppOpsManager.UidState int uidState, long accessTime, long accessDuration,
                 @AppOpsManager.AttributionFlags int attributionFlags, int attributionChainId) {
-            getOrCreateDiscretePackageOps(packageName).addDiscreteAccess(op, attributionTag, flags,
-                    uidState, accessTime, accessDuration, attributionFlags, attributionChainId);
+            getOrCreateDiscretePackageOps(packageName).addDiscreteAccess(op, deviceId,
+                    attributionTag, flags, uidState, accessTime, accessDuration,
+                    attributionFlags, attributionChainId);
         }
 
         private DiscretePackageOps getOrCreateDiscretePackageOps(String packageName) {
@@ -948,7 +973,7 @@
         }
     }
 
-    private final class DiscretePackageOps {
+    static final class DiscretePackageOps {
         ArrayMap<Integer, DiscreteOp> mPackageOps;
 
         DiscretePackageOps() {
@@ -959,12 +984,12 @@
             return mPackageOps.isEmpty();
         }
 
-        void addDiscreteAccess(int op, @Nullable String attributionTag,
+        void addDiscreteAccess(int op, @NonNull String deviceId, @Nullable String attributionTag,
                 @AppOpsManager.OpFlags int flags, @AppOpsManager.UidState int uidState,
                 long accessTime, long accessDuration,
                 @AppOpsManager.AttributionFlags int attributionFlags, int attributionChainId) {
-            getOrCreateDiscreteOp(op).addDiscreteAccess(attributionTag, flags, uidState, accessTime,
-                    accessDuration, attributionFlags, attributionChainId);
+            getOrCreateDiscreteOp(op).addDiscreteAccess(deviceId, attributionTag, flags, uidState,
+                    accessTime, accessDuration, attributionFlags, attributionChainId);
         }
 
         void merge(DiscretePackageOps other) {
@@ -1056,10 +1081,148 @@
         }
     }
 
-    private final class DiscreteOp {
-        ArrayMap<String, List<DiscreteOpEvent>> mAttributedOps;
+    static final class DiscreteOp {
+        ArrayMap<String, DiscreteDeviceOp> mDeviceAttributedOps;
 
         DiscreteOp() {
+            mDeviceAttributedOps = new ArrayMap<>();
+        }
+
+        boolean isEmpty() {
+            return mDeviceAttributedOps.isEmpty();
+        }
+
+        void merge(DiscreteOp other) {
+            int nDevices = other.mDeviceAttributedOps.size();
+            for (int i = 0; i < nDevices; i++) {
+                String deviceId = other.mDeviceAttributedOps.keyAt(i);
+                DiscreteDeviceOp otherDeviceOps = other.mDeviceAttributedOps.valueAt(i);
+                getOrCreateDiscreteDeviceOp(deviceId).merge(otherDeviceOps);
+            }
+        }
+
+        // Note: Update this method when we want to filter by device Id.
+        private void filter(long beginTimeMillis, long endTimeMillis,
+                @AppOpsManager.HistoricalOpsRequestFilter int filter,
+                @Nullable String attributionTagFilter, @AppOpsManager.OpFlags int flagsFilter,
+                int currentUid, String currentPkgName, int currentOp,
+                ArrayMap<Integer, AttributionChain> attributionChains) {
+            int nDevices = mDeviceAttributedOps.size();
+            for (int i = nDevices - 1; i >= 0; i--) {
+                mDeviceAttributedOps.valueAt(i).filter(beginTimeMillis, endTimeMillis, filter,
+                        attributionTagFilter, flagsFilter, currentUid, currentPkgName, currentOp,
+                        attributionChains);
+                if (mDeviceAttributedOps.valueAt(i).isEmpty()) {
+                    mDeviceAttributedOps.removeAt(i);
+                }
+            }
+        }
+
+        private void offsetHistory(long offset) {
+            int nDevices = mDeviceAttributedOps.size();
+            for (int i = 0; i < nDevices; i++) {
+                mDeviceAttributedOps.valueAt(i).offsetHistory(offset);
+            }
+        }
+
+        void addDiscreteAccess(@NonNull String deviceId, @Nullable String attributionTag,
+                @AppOpsManager.OpFlags int flags, @AppOpsManager.UidState int uidState,
+                long accessTime, long accessDuration,
+                @AppOpsManager.AttributionFlags int attributionFlags, int attributionChainId) {
+            getOrCreateDiscreteDeviceOp(deviceId).addDiscreteAccess(attributionTag, flags, uidState,
+                    accessTime, accessDuration, attributionFlags, attributionChainId);
+        }
+
+        private DiscreteDeviceOp getOrCreateDiscreteDeviceOp(String deviceId) {
+            return mDeviceAttributedOps.computeIfAbsent(deviceId, k -> new DiscreteDeviceOp());
+        }
+
+        // TODO: b/308716962 Retrieve discrete histories from all devices and integrate them with
+        // HistoricalOps
+        private void applyToHistory(AppOpsManager.HistoricalOps result, int uid,
+                @NonNull String packageName, int op,
+                @NonNull ArrayMap<Integer, AttributionChain> attributionChains) {
+            if (mDeviceAttributedOps.get(PERSISTENT_DEVICE_ID_DEFAULT) != null) {
+                mDeviceAttributedOps.get(PERSISTENT_DEVICE_ID_DEFAULT).applyToHistory(result, uid,
+                        packageName, op, attributionChains);
+            }
+        }
+
+        private void dump(@NonNull PrintWriter pw, @NonNull SimpleDateFormat sdf,
+                @NonNull Date date, @NonNull String prefix, int nDiscreteOps) {
+            int nDevices = mDeviceAttributedOps.size();
+            for (int i = 0; i < nDevices; i++) {
+                pw.print(prefix);
+                pw.print("Device: ");
+                pw.print(mDeviceAttributedOps.keyAt(i));
+                pw.println();
+                mDeviceAttributedOps.valueAt(i).dump(pw, sdf, date, prefix + "  ",
+                        nDiscreteOps);
+            }
+        }
+
+        void serialize(TypedXmlSerializer out) throws Exception {
+            int nDevices = mDeviceAttributedOps.size();
+            for (int i = 0; i < nDevices; i++) {
+                String deviceId = mDeviceAttributedOps.keyAt(i);
+                mDeviceAttributedOps.valueAt(i).serialize(out, deviceId);
+            }
+        }
+
+        void deserialize(TypedXmlPullParser parser, long beginTimeMillis) throws Exception {
+            int outerDepth = parser.getDepth();
+            while (XmlUtils.nextElementWithin(parser, outerDepth)) {
+                if (TAG_TAG.equals(parser.getName())) {
+                    String attributionTag = parser.getAttributeValue(null, ATTR_TAG);
+
+                    int innerDepth = parser.getDepth();
+                    while (XmlUtils.nextElementWithin(parser, innerDepth)) {
+                        if (TAG_ENTRY.equals(parser.getName())) {
+                            long noteTime = parser.getAttributeLong(null, ATTR_NOTE_TIME);
+                            long noteDuration = parser.getAttributeLong(null, ATTR_NOTE_DURATION,
+                                    -1);
+                            int uidState = parser.getAttributeInt(null, ATTR_UID_STATE);
+                            int opFlags = parser.getAttributeInt(null, ATTR_FLAGS);
+                            int attributionFlags = parser.getAttributeInt(null,
+                                    ATTR_ATTRIBUTION_FLAGS, AppOpsManager.ATTRIBUTION_FLAGS_NONE);
+                            int attributionChainId = parser.getAttributeInt(null, ATTR_CHAIN_ID,
+                                    AppOpsManager.ATTRIBUTION_CHAIN_ID_NONE);
+                            String deviceId = parser.getAttributeValue(null, ATTR_DEVICE_ID);
+                            if (deviceId == null) {
+                                deviceId = PERSISTENT_DEVICE_ID_DEFAULT;
+                            }
+                            if (noteTime + noteDuration < beginTimeMillis) {
+                                continue;
+                            }
+
+                            DiscreteDeviceOp deviceOps = getOrCreateDiscreteDeviceOp(deviceId);
+                            List<DiscreteOpEvent> events =
+                                    deviceOps.getOrCreateDiscreteOpEventsList(attributionTag);
+                            DiscreteOpEvent event = new DiscreteOpEvent(noteTime, noteDuration,
+                                    uidState, opFlags, attributionFlags, attributionChainId);
+                            events.add(event);
+                        }
+                    }
+                }
+            }
+
+            int nDeviceOps = mDeviceAttributedOps.size();
+            for (int i = 0; i < nDeviceOps; i++) {
+                DiscreteDeviceOp deviceOp = mDeviceAttributedOps.valueAt(i);
+
+                int nAttrOps = deviceOp.mAttributedOps.size();
+                for (int j = 0; j < nAttrOps; j++) {
+                    List<DiscreteOpEvent> events = deviceOp.mAttributedOps.valueAt(j);
+                    events.sort(Comparator.comparingLong(a -> a.mNoteTime));
+                }
+            }
+        }
+    }
+
+    static final class DiscreteDeviceOp {
+        ArrayMap<String, List<DiscreteOpEvent>> mAttributedOps;
+
+        DiscreteDeviceOp() {
             mAttributedOps = new ArrayMap<>();
         }
 
@@ -1067,7 +1230,7 @@
             return mAttributedOps.isEmpty();
         }
 
-        void merge(DiscreteOp other) {
+        void merge(DiscreteDeviceOp other) {
             int nTags = other.mAttributedOps.size();
             for (int i = 0; i < nTags; i++) {
                 String tag = other.mAttributedOps.keyAt(i);
@@ -1097,7 +1260,7 @@
                         currentUid, currentPkgName, currentOp, mAttributedOps.keyAt(i),
                         attributionChains);
                 mAttributedOps.put(tag, list);
-                if (list.size() == 0) {
+                if (list.isEmpty()) {
                     mAttributedOps.removeAt(i);
                 }
             }
@@ -1125,8 +1288,7 @@
             List<DiscreteOpEvent> attributedOps = getOrCreateDiscreteOpEventsList(
                     attributionTag);
 
-            int nAttributedOps = attributedOps.size();
-            int i = nAttributedOps;
+            int i = attributedOps.size();
             for (; i > 0; i--) {
                 DiscreteOpEvent previousOp = attributedOps.get(i - 1);
                 if (discretizeTimeStamp(previousOp.mNoteTime) < discretizeTimeStamp(accessTime)) {
@@ -1148,12 +1310,8 @@
         }
 
         private List<DiscreteOpEvent> getOrCreateDiscreteOpEventsList(String attributionTag) {
-            List<DiscreteOpEvent> result = mAttributedOps.get(attributionTag);
-            if (result == null) {
-                result = new ArrayList<>();
-                mAttributedOps.put(attributionTag, result);
-            }
-            return result;
+            return mAttributedOps.computeIfAbsent(attributionTag,
+                    k -> new ArrayList<>());
         }
 
         private void applyToHistory(AppOpsManager.HistoricalOps result, int uid,
@@ -1167,8 +1325,7 @@
                 for (int j = 0; j < nEvents; j++) {
                     DiscreteOpEvent event = events.get(j);
                     AppOpsManager.OpEventProxyInfo proxy = null;
-                    if (event.mAttributionChainId != ATTRIBUTION_CHAIN_ID_NONE
-                            && attributionChains != null) {
+                    if (event.mAttributionChainId != ATTRIBUTION_CHAIN_ID_NONE) {
                         AttributionChain chain = attributionChains.get(event.mAttributionChainId);
                         if (chain != null && chain.isComplete()
                                 && chain.isStart(packageName, uid, tag, op, event)
@@ -1198,65 +1355,31 @@
                 int first = nDiscreteOps < 1 ? 0 : max(0, nOps - nDiscreteOps);
                 for (int j = first; j < nOps; j++) {
                     ops.get(j).dump(pw, sdf, date, prefix + "  ");
-
                 }
             }
         }
 
-        void serialize(TypedXmlSerializer out) throws Exception {
+        void serialize(TypedXmlSerializer out, String deviceId) throws Exception {
             int nAttributions = mAttributedOps.size();
             for (int i = 0; i < nAttributions; i++) {
                 out.startTag(null, TAG_TAG);
                 String tag = mAttributedOps.keyAt(i);
                 if (tag != null) {
-                    out.attribute(null, ATTR_TAG, mAttributedOps.keyAt(i));
+                    out.attribute(null, ATTR_TAG, tag);
                 }
                 List<DiscreteOpEvent> ops = mAttributedOps.valueAt(i);
                 int nOps = ops.size();
                 for (int j = 0; j < nOps; j++) {
                     out.startTag(null, TAG_ENTRY);
-                    ops.get(j).serialize(out);
+                    ops.get(j).serialize(out, deviceId);
                     out.endTag(null, TAG_ENTRY);
                 }
                 out.endTag(null, TAG_TAG);
             }
         }
-
-        void deserialize(TypedXmlPullParser parser, long beginTimeMillis) throws Exception {
-            int outerDepth = parser.getDepth();
-            while (XmlUtils.nextElementWithin(parser, outerDepth)) {
-                if (TAG_TAG.equals(parser.getName())) {
-                    String attributionTag = parser.getAttributeValue(null, ATTR_TAG);
-                    List<DiscreteOpEvent> events = getOrCreateDiscreteOpEventsList(
-                            attributionTag);
-                    int innerDepth = parser.getDepth();
-                    while (XmlUtils.nextElementWithin(parser, innerDepth)) {
-                        if (TAG_ENTRY.equals(parser.getName())) {
-                            long noteTime = parser.getAttributeLong(null, ATTR_NOTE_TIME);
-                            long noteDuration = parser.getAttributeLong(null, ATTR_NOTE_DURATION,
-                                    -1);
-                            int uidState = parser.getAttributeInt(null, ATTR_UID_STATE);
-                            int opFlags = parser.getAttributeInt(null, ATTR_FLAGS);
-                            int attributionFlags = parser.getAttributeInt(null,
-                                    ATTR_ATTRIBUTION_FLAGS, AppOpsManager.ATTRIBUTION_FLAGS_NONE);
-                            int attributionChainId = parser.getAttributeInt(null, ATTR_CHAIN_ID,
-                                    AppOpsManager.ATTRIBUTION_CHAIN_ID_NONE);
-                            if (noteTime + noteDuration < beginTimeMillis) {
-                                continue;
-                            }
-                            DiscreteOpEvent event = new DiscreteOpEvent(noteTime, noteDuration,
-                                    uidState, opFlags, attributionFlags, attributionChainId);
-                            events.add(event);
-                        }
-                    }
-                    Collections.sort(events, (a, b) -> a.mNoteTime < b.mNoteTime ? -1
-                            : (a.mNoteTime == b.mNoteTime ? 0 : 1));
-                }
-            }
-        }
     }
 
-    private final class DiscreteOpEvent {
+    static final class DiscreteOpEvent {
         final long mNoteTime;
         final long mNoteDuration;
         final @AppOpsManager.UidState int mUidState;
@@ -1306,7 +1429,7 @@
             pw.println();
         }
 
-        private void serialize(TypedXmlSerializer out) throws Exception {
+        private void serialize(TypedXmlSerializer out, String deviceId) throws Exception {
             out.attributeLong(null, ATTR_NOTE_TIME, mNoteTime);
             if (mNoteDuration != -1) {
                 out.attributeLong(null, ATTR_NOTE_DURATION, mNoteDuration);
@@ -1317,6 +1440,9 @@
             if (mAttributionChainId != AppOpsManager.ATTRIBUTION_CHAIN_ID_NONE) {
                 out.attributeInt(null, ATTR_CHAIN_ID, mAttributionChainId);
             }
+            if (!Objects.equals(deviceId, PERSISTENT_DEVICE_ID_DEFAULT)) {
+                out.attribute(null, ATTR_DEVICE_ID, deviceId);
+            }
             out.attributeInt(null, ATTR_UID_STATE, mUidState);
             out.attributeInt(null, ATTR_FLAGS, mOpFlag);
         }
diff --git a/services/core/java/com/android/server/appop/HistoricalRegistry.java b/services/core/java/com/android/server/appop/HistoricalRegistry.java
index dbd47d0..fffb108 100644
--- a/services/core/java/com/android/server/appop/HistoricalRegistry.java
+++ b/services/core/java/com/android/server/appop/HistoricalRegistry.java
@@ -472,9 +472,9 @@
     }
 
     void incrementOpAccessedCount(int op, int uid, @NonNull String packageName,
-            @Nullable String attributionTag, @UidState int uidState, @OpFlags int flags,
-            long accessTime, @AppOpsManager.AttributionFlags int attributionFlags,
-            int attributionChainId) {
+            @NonNull String deviceId, @Nullable String attributionTag, @UidState int uidState,
+            @OpFlags int flags, long accessTime,
+            @AppOpsManager.AttributionFlags int attributionFlags, int attributionChainId) {
         synchronized (mInMemoryLock) {
             if (mMode == AppOpsManager.HISTORICAL_MODE_ENABLED_ACTIVE) {
                 if (!isPersistenceInitializedMLocked()) {
@@ -485,8 +485,9 @@
                         System.currentTimeMillis()).increaseAccessCount(op, uid, packageName,
                         attributionTag, uidState, flags, 1);
 
-                mDiscreteRegistry.recordDiscreteAccess(uid, packageName, op, attributionTag,
-                        flags, uidState, accessTime, -1, attributionFlags, attributionChainId);
+                mDiscreteRegistry.recordDiscreteAccess(uid, packageName, deviceId, op,
+                        attributionTag, flags, uidState, accessTime, -1, attributionFlags,
+                        attributionChainId);
             }
         }
     }
@@ -507,8 +508,8 @@
     }
 
     void increaseOpAccessDuration(int op, int uid, @NonNull String packageName,
-            @Nullable String attributionTag, @UidState int uidState, @OpFlags int flags,
-            long eventStartTime, long increment,
+            @NonNull String deviceId, @Nullable String attributionTag, @UidState int uidState,
+            @OpFlags int flags, long eventStartTime, long increment,
             @AppOpsManager.AttributionFlags int attributionFlags, int attributionChainId) {
         synchronized (mInMemoryLock) {
             if (mMode == AppOpsManager.HISTORICAL_MODE_ENABLED_ACTIVE) {
@@ -519,9 +520,9 @@
                 getUpdatedPendingHistoricalOpsMLocked(
                         System.currentTimeMillis()).increaseAccessDuration(op, uid, packageName,
                         attributionTag, uidState, flags, increment);
-                mDiscreteRegistry.recordDiscreteAccess(uid, packageName, op, attributionTag,
-                        flags, uidState, eventStartTime, increment, attributionFlags,
-                        attributionChainId);
+                mDiscreteRegistry.recordDiscreteAccess(uid, packageName, deviceId, op,
+                        attributionTag, flags, uidState, eventStartTime, increment,
+                        attributionFlags, attributionChainId);
             }
         }
     }
diff --git a/services/core/java/com/android/server/hdmi/RequestActiveSourceAction.java b/services/core/java/com/android/server/hdmi/RequestActiveSourceAction.java
index 539a00d..a33d70a 100644
--- a/services/core/java/com/android/server/hdmi/RequestActiveSourceAction.java
+++ b/services/core/java/com/android/server/hdmi/RequestActiveSourceAction.java
@@ -19,6 +19,7 @@
 import android.hardware.hdmi.HdmiControlManager;
 import android.hardware.hdmi.IHdmiControlCallback;
 import android.util.Slog;
+import com.android.internal.annotations.VisibleForTesting;
 
 /**
  * Feature action that sends <Request Active Source> message and waits for <Active Source> on TV
@@ -39,6 +40,10 @@
     // Number of retries <Request Active Source> is sent if no device answers this message.
     private static final int MAX_SEND_RETRY_COUNT = 1;
 
+    // Timeout to wait for the LauncherX API call to be completed.
+    @VisibleForTesting
+    protected static final int TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS = 10000;
+
     private int mSendRetryCount = 0;
 
 
@@ -55,7 +60,7 @@
         // We wait for default timeout to allow the message triggered by the LauncherX API call to
         // be sent by the TV and another default timeout in case the message has to be answered
         // (e.g. TV sent a <Set Stream Path> or <Routing Change>).
-        addTimer(mState, HdmiConfig.TIMEOUT_MS * 2);
+        addTimer(mState, TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS);
         return true;
     }
 
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index 36a9c80..a06ad14 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -1254,14 +1254,15 @@
     /**
      * Start drag and drop.
      *
-     * @param fromChannel The input channel that is currently receiving a touch gesture that should
-     *                    be turned into the drag pointer.
-     * @param dragAndDropChannel The input channel associated with the system drag window.
+     * @param fromChannelToken The token of the input channel that is currently receiving a touch
+     *                        gesture that should be turned into the drag pointer.
+     * @param dragAndDropChannelToken The token of the input channel associated with the system drag
+     *                               window.
      * @return true if drag and drop was successfully started, false otherwise.
      */
-    public boolean startDragAndDrop(@NonNull InputChannel fromChannel,
-            @NonNull InputChannel dragAndDropChannel) {
-        return mNative.transferTouchGesture(fromChannel.getToken(), dragAndDropChannel.getToken(),
+    public boolean startDragAndDrop(@NonNull IBinder fromChannelToken,
+            @NonNull IBinder dragAndDropChannelToken) {
+        return mNative.transferTouchGesture(fromChannelToken, dragAndDropChannelToken,
                 true /* isDragDrop */);
     }
 
diff --git a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
index 99f4747..b08f917 100644
--- a/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
+++ b/services/core/java/com/android/server/inputmethod/AdditionalSubtypeMapRepository.java
@@ -38,10 +38,11 @@
 final class AdditionalSubtypeMapRepository {
     private static final String TAG = "AdditionalSubtypeMapRepository";
 
-    // TODO(b/352594784): Should we user other lock primitives?
-    @GuardedBy("sPerUserMap")
+    private static final Object sMutationLock = new Object();
+
     @NonNull
-    private static final SparseArray<AdditionalSubtypeMap> sPerUserMap = new SparseArray<>();
+    private static volatile ImmutableSparseArray<AdditionalSubtypeMap> sPerUserMap =
+            ImmutableSparseArray.empty();
 
     record WriteTask(@UserIdInt int userId, @NonNull AdditionalSubtypeMap subtypeMap,
                      @NonNull InputMethodMap inputMethodMap) {
@@ -198,7 +199,7 @@
     /**
      * Returns {@link AdditionalSubtypeMap} for the given user.
      *
-     * <p>This method is expected be called after {@link #ensureInitializedAndGet(int)}. Otherwise
+     * <p>This method is expected be called after {@link #initializeIfNecessary(int)}. Otherwise
      * {@link AdditionalSubtypeMap#EMPTY_MAP} will be returned.</p>
      *
      * @param userId the user to be queried about
@@ -207,10 +208,7 @@
     @AnyThread
     @NonNull
     static AdditionalSubtypeMap get(@UserIdInt int userId) {
-        final AdditionalSubtypeMap map;
-        synchronized (sPerUserMap) {
-            map = sPerUserMap.get(userId);
-        }
+        final AdditionalSubtypeMap map = sPerUserMap.get(userId);
         if (map == null) {
             Slog.e(TAG, "get(userId=" + userId + ") is called before loadInitialDataAndGet()."
                     + " Returning an empty map");
@@ -220,28 +218,24 @@
     }
 
     /**
-     * Ensures that {@link AdditionalSubtypeMap} is initialized for the given user. Load it from
-     * the persistent storage if {@link #putAndSave(int, AdditionalSubtypeMap, InputMethodMap)} has
-     * not been called yet.
+     * Ensures that {@link AdditionalSubtypeMap} is initialized for the given user.
      *
      * @param userId the user to be initialized
-     * @return {@link AdditionalSubtypeMap} that is associated with the given user. If
-     *         {@link #putAndSave(int, AdditionalSubtypeMap, InputMethodMap)} is already called
-     *         then the given {@link AdditionalSubtypeMap}.
      */
     @AnyThread
     @NonNull
-    static AdditionalSubtypeMap ensureInitializedAndGet(@UserIdInt int userId) {
-        final var map = AdditionalSubtypeUtils.load(userId);
-        synchronized (sPerUserMap) {
-            final AdditionalSubtypeMap previous = sPerUserMap.get(userId);
-            // If putAndSave() has already been called, then use it.
-            if (previous != null) {
-                return previous;
-            }
-            sPerUserMap.put(userId, map);
+    static void initializeIfNecessary(@UserIdInt int userId) {
+        if (sPerUserMap.contains(userId)) {
+            // Fast-pass. If putAndSave() is already called, then do nothing.
+            return;
         }
-        return map;
+        final var map = AdditionalSubtypeUtils.load(userId);
+        synchronized (sMutationLock) {
+            // Check the condition again.
+            if (!sPerUserMap.contains(userId)) {
+                sPerUserMap = sPerUserMap.cloneWithPutOrSelf(userId, map);
+            }
+        }
     }
 
     /**
@@ -255,12 +249,8 @@
     @AnyThread
     static void putAndSave(@UserIdInt int userId, @NonNull AdditionalSubtypeMap map,
             @NonNull InputMethodMap inputMethodMap) {
-        synchronized (sPerUserMap) {
-            final AdditionalSubtypeMap previous = sPerUserMap.get(userId);
-            if (previous == map) {
-                return;
-            }
-            sPerUserMap.put(userId, map);
+        synchronized (sMutationLock) {
+            sPerUserMap = sPerUserMap.cloneWithPutOrSelf(userId, map);
             sWriter.scheduleWriteTask(userId, map, inputMethodMap);
         }
     }
@@ -277,9 +267,9 @@
 
     @AnyThread
     static void remove(@UserIdInt int userId) {
-        synchronized (sPerUserMap) {
+        synchronized (sMutationLock) {
             sWriter.onUserRemoved(userId);
-            sPerUserMap.remove(userId);
+            sPerUserMap = sPerUserMap.cloneWithRemoveOrSelf(userId);
         }
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java b/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
index a7280e6..7f7ae10 100644
--- a/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
+++ b/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
@@ -154,7 +154,7 @@
         void reportPerceptibleAsync(IBinder windowToken, boolean perceptible);
 
         @PermissionVerified(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
-        void removeImeSurface();
+        void removeImeSurface(int displayId);
 
         void removeImeSurfaceFromWindowAsync(IBinder windowToken);
 
@@ -384,10 +384,10 @@
 
     @EnforcePermission(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
     @Override
-    public void removeImeSurface() {
+    public void removeImeSurface(int displayId) {
         super.removeImeSurface_enforcePermission();
 
-        mCallback.removeImeSurface();
+        mCallback.removeImeSurface(displayId);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java b/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java
index 7ebf595..0b3f3f0 100644
--- a/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java
+++ b/services/core/java/com/android/server/inputmethod/ImeVisibilityStateComputer.java
@@ -36,8 +36,10 @@
 import static com.android.server.inputmethod.InputMethodManagerService.computeImeDisplayIdForTarget;
 
 import android.accessibilityservice.AccessibilityService;
+import android.annotation.AnyThread;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.annotation.UserIdInt;
 import android.content.res.Configuration;
 import android.os.Binder;
@@ -83,6 +85,7 @@
      * A map used to track the requested IME target window and its state. The key represents the
      * token of the window and the value is the corresponding IME window state.
      */
+    @GuardedBy("ImfLock.class")
     private final WeakHashMap<IBinder, ImeTargetWindowState> mRequestWindowStateMap =
             new WeakHashMap<>();
 
@@ -93,6 +96,7 @@
      * @see InputMethodManager#HIDE_IMPLICIT_ONLY that system will not hide IME when the value is
      * {@code true}.
      */
+    @GuardedBy("ImfLock.class")
     boolean mRequestedShowExplicitly;
 
     /**
@@ -101,25 +105,39 @@
      * @see InputMethodManager#SHOW_FORCED
      * @see InputMethodManager#HIDE_NOT_ALWAYS
      */
+    @GuardedBy("ImfLock.class")
     boolean mShowForced;
 
     /**
      * Set if we last told the input method to show itself.
      */
+    @GuardedBy("ImfLock.class")
     private boolean mInputShown;
 
     /**
      * Set if we called
      * {@link com.android.server.wm.ImeTargetVisibilityPolicy#showImeScreenshot(IBinder, int)}.
      */
+    @GuardedBy("ImfLock.class")
     private boolean mRequestedImeScreenshot;
 
     /** The window token of the current visible IME layering target overlay. */
+    @GuardedBy("ImfLock.class")
     private IBinder mCurVisibleImeLayeringOverlay;
 
     /** The window token of the current visible IME input target. */
+    @GuardedBy("ImfLock.class")
     private IBinder mCurVisibleImeInputTarget;
 
+    /**
+     * The last window token that we confirmed that IME started talking to.  This is always updated
+     * upon reports from the input method.  If the window state is already changed before the report
+     * is handled, this field just keeps the last value.
+     */
+    @GuardedBy("ImfLock.class")
+    @Nullable
+    private IBinder mLastImeTargetWindow;
+
     /** Represent the invalid IME visibility state */
     public static final int STATE_INVALID = -1;
 
@@ -203,25 +221,31 @@
             public void onImeTargetOverlayVisibilityChanged(IBinder overlayWindowToken,
                     @WindowManager.LayoutParams.WindowType int windowType, boolean visible,
                     boolean removed) {
-                mCurVisibleImeLayeringOverlay =
-                        // Ignoring the starting window since it's ok to cover the IME target
-                        // window in temporary without affecting the IME visibility.
-                        (visible && !removed && windowType != TYPE_APPLICATION_STARTING)
+                // Ignoring the starting window since it's ok to cover the IME target
+                // window in temporary without affecting the IME visibility.
+                final var overlay = (visible && !removed && windowType != TYPE_APPLICATION_STARTING)
                                 ? overlayWindowToken : null;
+                synchronized (ImfLock.class) {
+                    mCurVisibleImeLayeringOverlay = overlay;
+                }
             }
 
             @Override
             public void onImeInputTargetVisibilityChanged(IBinder imeInputTarget,
                     boolean visibleRequested, boolean removed) {
-                if (mCurVisibleImeInputTarget == imeInputTarget && (!visibleRequested || removed)
-                        && mCurVisibleImeLayeringOverlay != null) {
-                    final int reason = SoftInputShowHideReason.HIDE_WHEN_INPUT_TARGET_INVISIBLE;
-                    final var statsToken = ImeTracker.forLogging().onStart(ImeTracker.TYPE_HIDE,
-                            ImeTracker.ORIGIN_SERVER, reason, false /* fromUser */);
-                    mService.onApplyImeVisibilityFromComputer(imeInputTarget, statsToken,
-                            new ImeVisibilityResult(STATE_HIDE_IME_EXPLICIT, reason));
+                synchronized (ImfLock.class) {
+                    if (mCurVisibleImeInputTarget == imeInputTarget && (!visibleRequested
+                            || removed)
+                            && mCurVisibleImeLayeringOverlay != null) {
+                        final int reason = SoftInputShowHideReason.HIDE_WHEN_INPUT_TARGET_INVISIBLE;
+                        final var statsToken = ImeTracker.forLogging().onStart(ImeTracker.TYPE_HIDE,
+                                ImeTracker.ORIGIN_SERVER, reason, false /* fromUser */);
+                        mService.onApplyImeVisibilityFromComputerLocked(imeInputTarget, statsToken,
+                                new ImeVisibilityResult(STATE_HIDE_IME_EXPLICIT, reason));
+                    }
+                    mCurVisibleImeInputTarget =
+                            (visibleRequested && !removed) ? imeInputTarget : null;
                 }
-                mCurVisibleImeInputTarget = (visibleRequested && !removed) ? imeInputTarget : null;
             }
         });
     }
@@ -232,6 +256,7 @@
      * @param statsToken The token tracking the current IME request.
      * @return {@code true} when the show request can proceed.
      */
+    @GuardedBy("ImfLock.class")
     boolean onImeShowFlags(@NonNull ImeTracker.Token statsToken,
             @InputMethodManager.ShowFlags int showFlags) {
         if (mPolicy.mA11yRequestingNoSoftKeyboard || mPolicy.mImeHiddenByDisplayPolicy) {
@@ -258,6 +283,7 @@
      * @param statsToken The token tracking the current IME request.
      * @return {@code true} when the hide request can proceed.
      */
+    @GuardedBy("ImfLock.class")
     boolean canHideIme(@NonNull ImeTracker.Token statsToken,
             @InputMethodManager.HideFlags int hideFlags) {
         if ((hideFlags & InputMethodManager.HIDE_IMPLICIT_ONLY) != 0
@@ -279,6 +305,7 @@
      * Returns the show flags for IME. This translates from {@link InputMethodManager.ShowFlags}
      * to {@link InputMethod.ShowFlags}.
      */
+    @GuardedBy("ImfLock.class")
     @InputMethod.ShowFlags
     int getShowFlagsForInputMethodServiceOnly() {
         int flags = 0;
@@ -294,6 +321,7 @@
      * Returns the show flags for IMM. This translates from {@link InputMethod.ShowFlags}
      * to {@link InputMethodManager.ShowFlags}.
      */
+    @GuardedBy("ImfLock.class")
     @InputMethodManager.ShowFlags
     int getShowFlags() {
         int flags = 0;
@@ -305,12 +333,14 @@
         return flags;
     }
 
+    @GuardedBy("ImfLock.class")
     void clearImeShowFlags() {
         mRequestedShowExplicitly = false;
         mShowForced = false;
         mInputShown = false;
     }
 
+    @GuardedBy("ImfLock.class")
     int computeImeDisplayId(@NonNull ImeTargetWindowState state, int displayId) {
         final int displayToShowIme = computeImeDisplayIdForTarget(displayId, mImeDisplayValidator);
         state.setImeDisplayId(displayToShowIme);
@@ -328,6 +358,7 @@
      *                            visibility state, it could be {@link #STATE_SHOW_IME} or
      *                            {@link #STATE_HIDE_IME}.
      */
+    @GuardedBy("ImfLock.class")
     void requestImeVisibility(IBinder windowToken, boolean showIme) {
         ImeTargetWindowState state = getOrCreateWindowState(windowToken);
         if (!mPolicy.mPendingA11yRequestingHideKeyboard) {
@@ -343,6 +374,7 @@
         setWindowStateInner(windowToken, state);
     }
 
+    @GuardedBy("ImfLock.class")
     ImeTargetWindowState getOrCreateWindowState(IBinder windowToken) {
         ImeTargetWindowState state = mRequestWindowStateMap.get(windowToken);
         if (state == null) {
@@ -351,11 +383,13 @@
         return state;
     }
 
+    @GuardedBy("ImfLock.class")
     ImeTargetWindowState getWindowStateOrNull(IBinder windowToken) {
         ImeTargetWindowState state = mRequestWindowStateMap.get(windowToken);
         return state;
     }
 
+    @GuardedBy("ImfLock.class")
     void setWindowState(IBinder windowToken, @NonNull ImeTargetWindowState newState) {
         final ImeTargetWindowState state = mRequestWindowStateMap.get(windowToken);
         if (state != null && newState.hasEditorFocused()
@@ -367,6 +401,7 @@
         setWindowStateInner(windowToken, newState);
     }
 
+    @GuardedBy("ImfLock.class")
     private void setWindowStateInner(IBinder windowToken, @NonNull ImeTargetWindowState newState) {
         if (DEBUG) Slog.d(TAG, "setWindowStateInner, windowToken=" + windowToken
                 + ", state=" + newState);
@@ -391,6 +426,7 @@
         }
     }
 
+    @GuardedBy("ImfLock.class")
     ImeVisibilityResult computeState(ImeTargetWindowState state, boolean allowVisible) {
         // TODO: Output the request IME visibility state according to the requested window state
         final int softInputVisibility = state.mSoftInputModeState & SOFT_INPUT_MASK_STATE;
@@ -452,8 +488,7 @@
                 break;
             case WindowManager.LayoutParams.SOFT_INPUT_STATE_UNCHANGED:
                 // Do nothing but preserving the last IME requested visibility state.
-                final ImeTargetWindowState lastState =
-                        getWindowStateOrNull(mService.mLastImeTargetWindow);
+                final ImeTargetWindowState lastState = getWindowStateOrNull(mLastImeTargetWindow);
                 if (lastState != null) {
                     state.setRequestedImeVisible(lastState.mRequestedImeVisible);
                 }
@@ -540,7 +575,7 @@
         return null;
     }
 
-    @VisibleForTesting
+    @GuardedBy("ImfLock.class")
     ImeVisibilityResult onInteractiveChanged(IBinder windowToken, boolean interactive) {
         final ImeTargetWindowState state = getWindowStateOrNull(windowToken);
         if (state != null && state.isRequestedImeVisible() && mInputShown && !interactive) {
@@ -568,6 +603,7 @@
         return userData.mImeBindingState.mFocusedWindow;
     }
 
+    @GuardedBy("ImfLock.class")
     IBinder getWindowTokenFrom(ImeTargetWindowState windowState) {
         for (IBinder windowToken : mRequestWindowStateMap.keySet()) {
             final ImeTargetWindowState state = mRequestWindowStateMap.get(windowToken);
@@ -578,6 +614,7 @@
         return null;
     }
 
+    @GuardedBy("ImfLock.class")
     boolean shouldRestoreImeVisibility(@NonNull ImeTargetWindowState state) {
         final int softInputMode = state.getSoftInputModeState();
         switch (softInputMode & WindowManager.LayoutParams.SOFT_INPUT_MASK_STATE) {
@@ -591,14 +628,28 @@
         return mWindowManagerInternal.shouldRestoreImeVisibility(getWindowTokenFrom(state));
     }
 
+    @GuardedBy("ImfLock.class")
     boolean isInputShown() {
         return mInputShown;
     }
 
+    @GuardedBy("ImfLock.class")
     void setInputShown(boolean inputShown) {
         mInputShown = inputShown;
     }
 
+    @GuardedBy("ImfLock.class")
+    @Nullable
+    IBinder getLastImeTargetWindow() {
+        return mLastImeTargetWindow;
+    }
+
+    @GuardedBy("ImfLock.class")
+    void setLastImeTargetWindow(@Nullable IBinder imeTargetWindow) {
+        mLastImeTargetWindow = imeTargetWindow;
+    }
+
+    @GuardedBy("ImfLock.class")
     void dumpDebug(ProtoOutputStream proto, long fieldId) {
         proto.write(SHOW_EXPLICITLY_REQUESTED, mRequestedShowExplicitly);
         proto.write(SHOW_FORCED, mShowForced);
@@ -607,12 +658,14 @@
         proto.write(INPUT_SHOWN, mInputShown);
     }
 
+    @GuardedBy("ImfLock.class")
     void dump(@NonNull PrintWriter pw, @NonNull String prefix) {
         final Printer p = new PrintWriterPrinter(pw);
         p.println(prefix + "mRequestedShowExplicitly=" + mRequestedShowExplicitly
                 + " mShowForced=" + mShowForced);
         p.println(prefix + "mImeHiddenByDisplayPolicy=" + mPolicy.isImeHiddenByDisplayPolicy());
         p.println(prefix + "mInputShown=" + mInputShown);
+        p.println(prefix + "mLastImeTargetWindow=" + mLastImeTargetWindow);
     }
 
     /**
@@ -629,12 +682,14 @@
          *
          * This prevents the IME from showing when it otherwise may have shown.
          */
+        @GuardedBy("ImfLock.class")
         private boolean mImeHiddenByDisplayPolicy;
 
         /**
          * Set when the accessibility service requests to hide IME by
          * {@link AccessibilityService.SoftKeyboardController#setShowMode}
          */
+        @GuardedBy("ImfLock.class")
         private boolean mA11yRequestingNoSoftKeyboard;
 
         /**
@@ -643,16 +698,20 @@
          * {@link android.provider.Settings.Secure#ACCESSIBILITY_SOFT_KEYBOARD_MODE} without
          * changing the requested IME visible state.
          */
+        @GuardedBy("ImfLock.class")
         private boolean mPendingA11yRequestingHideKeyboard;
 
+        @GuardedBy("ImfLock.class")
         void setImeHiddenByDisplayPolicy(boolean hideIme) {
             mImeHiddenByDisplayPolicy = hideIme;
         }
 
+        @GuardedBy("ImfLock.class")
         boolean isImeHiddenByDisplayPolicy() {
             return mImeHiddenByDisplayPolicy;
         }
 
+        @GuardedBy("ImfLock.class")
         void setA11yRequestNoSoftKeyboard(int keyboardShowMode) {
             mA11yRequestingNoSoftKeyboard =
                     (keyboardShowMode & AccessibilityService.SHOW_MODE_MASK) == SHOW_MODE_HIDDEN;
@@ -661,11 +720,13 @@
             }
         }
 
+        @GuardedBy("ImfLock.class")
         boolean isA11yRequestNoSoftKeyboard() {
             return mA11yRequestingNoSoftKeyboard;
         }
     }
 
+    @GuardedBy("ImfLock.class")
     ImeVisibilityPolicy getImePolicy() {
         return mPolicy;
     }
@@ -721,63 +782,78 @@
         /**
          * Set if the client has asked for the input method to be shown.
          */
+        @GuardedBy("ImfLock.class")
         private boolean mRequestedImeVisible;
 
         /**
          * A identifier for knowing the requester of {@link InputMethodManager#showSoftInput} or
          * {@link InputMethodManager#hideSoftInputFromWindow}.
          */
+        @GuardedBy("ImfLock.class")
         private IBinder mRequestImeToken;
 
         /**
          * The IME target display id for which the latest startInput was called.
          */
+        @GuardedBy("ImfLock.class")
         private int mImeDisplayId = DEFAULT_DISPLAY;
 
+        @AnyThread
         boolean hasImeFocusChanged() {
             return mImeFocusChanged;
         }
 
+        @AnyThread
         boolean hasEditorFocused() {
             return mHasFocusedEditor;
         }
 
+        @AnyThread
         boolean isStartInputByGainFocus() {
             return mIsStartInputByGainFocus;
         }
 
+        @AnyThread
         int getSoftInputModeState() {
             return mSoftInputModeState;
         }
 
+        @AnyThread
         int getWindowFlags() {
             return mWindowFlags;
         }
 
+        @AnyThread
         int getToolType() {
             return mToolType;
         }
 
+        @GuardedBy("ImfLock.class")
         private void setImeDisplayId(int imeDisplayId) {
             mImeDisplayId = imeDisplayId;
         }
 
+        @GuardedBy("ImfLock.class")
         int getImeDisplayId() {
             return mImeDisplayId;
         }
 
+        @GuardedBy("ImfLock.class")
         private void setRequestedImeVisible(boolean requestedImeVisible) {
             mRequestedImeVisible = requestedImeVisible;
         }
 
+        @GuardedBy("ImfLock.class")
         boolean isRequestedImeVisible() {
             return mRequestedImeVisible;
         }
 
+        @GuardedBy("ImfLock.class")
         void setRequestImeToken(IBinder token) {
             mRequestImeToken = token;
         }
 
+        @GuardedBy("ImfLock.class")
         IBinder getRequestImeToken() {
             return mRequestImeToken;
         }
diff --git a/services/core/java/com/android/server/inputmethod/ImmutableSparseArray.java b/services/core/java/com/android/server/inputmethod/ImmutableSparseArray.java
new file mode 100644
index 0000000..382aa8a
--- /dev/null
+++ b/services/core/java/com/android/server/inputmethod/ImmutableSparseArray.java
@@ -0,0 +1,183 @@
+/*
+ * Copyright (C) 2024 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.server.inputmethod;
+
+
+import android.annotation.AnyThread;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.util.SparseArray;
+
+import java.util.function.Consumer;
+
+/**
+ * A holder object to expose {@link SparseArray} to multiple threads in a thread-safe manner through
+ * "final Field Semantics" defined in JLS 17.5, with only exposing thread-safe methods such as
+ * {@link SparseArray#get(int)} and {@link SparseArray#size()} from {@link SparseArray}, and with
+ * adding clone-with-update style methods {@link #cloneWithPutOrSelf(int, Object)} and
+ * {@link #cloneWithRemoveOrSelf(int)} instead of exposing mutation methods.
+ *
+ * @param <E> Type of the element
+ */
+final class ImmutableSparseArray<E> {
+    @NonNull
+    private final SparseArray<E> mArray;
+
+    private static final ImmutableSparseArray<Object> EMPTY =
+            new ImmutableSparseArray<>(new SparseArray<>());
+
+    /**
+     * Returns an empty {@link ImmutableSparseArray} instance.
+     *
+     * @return An empty {@link ImmutableSparseArray} instance.
+     * @param <T> Type of the element
+     */
+    @SuppressWarnings("unchecked")
+    @AnyThread
+    @NonNull
+    static <T> ImmutableSparseArray<T> empty() {
+        return (ImmutableSparseArray<T>) EMPTY;
+    }
+
+    private ImmutableSparseArray(@NonNull SparseArray<E> array) {
+        mArray = array;
+    }
+
+    /**
+     * @return the size of this array
+     */
+    @AnyThread
+    int size() {
+        return mArray.size();
+    }
+
+    /**
+     * Returns the key of the specified index.
+     *
+     * @return the key of the specified index
+     * @throws ArrayIndexOutOfBoundsException when the index is out of range
+     */
+    @AnyThread
+    int keyAt(int index) {
+        return mArray.keyAt(index);
+    }
+
+    /**
+     * Returns the value of the specified index.
+     *
+     * @return the value of the specified index
+     * @throws ArrayIndexOutOfBoundsException when the index is out of range
+     */
+    @AnyThread
+    @Nullable
+    public E valueAt(int index) {
+        return mArray.valueAt(index);
+    }
+
+    /**
+     * Returns the index of the specified key.
+     *
+     * @return the index of the specified key if exists. Otherwise {@code -1}
+     */
+    @AnyThread
+    int indexOfKey(int key) {
+        return mArray.indexOfKey(key);
+    }
+
+    /**
+     * Returns {@code true} if the given {@code key} exists.
+     *
+     * @param key the key to be queried
+     * @return    {@code true} if the given {@code key} exists
+     */
+    @AnyThread
+    boolean contains(int key) {
+        return mArray.contains(key);
+    }
+
+    /**
+     * Returns the value associated with the {@code key}.
+     *
+     * @param key the key to be queried
+     * @return    the value associated with the {@code key} if exists. Otherwise {@code null}
+     */
+    @AnyThread
+    @Nullable
+    E get(int key) {
+        return mArray.get(key);
+    }
+
+    /**
+     * Run {@link Consumer} for each value.
+     *
+     * @param consumer {@link Consumer} to be called back
+     */
+    @AnyThread
+    void forEach(@NonNull Consumer<E> consumer) {
+        final int size = mArray.size();
+        for (int i = 0; i < size; ++i) {
+            consumer.accept(mArray.valueAt(i));
+        }
+    }
+
+    /**
+     * Returns an instance of {@link ImmutableSparseArray} that has the given key and value on top
+     * of items cloned from this instance.
+     *
+     * @param key   the key to be added
+     * @param value the value to be added
+     * @return      the same {@link ImmutableSparseArray} instance if there is actually no update.
+     *              Otherwise, a new instance of {@link ImmutableSparseArray}
+     */
+    @AnyThread
+    @NonNull
+    ImmutableSparseArray<E> cloneWithPutOrSelf(int key, @Nullable E value) {
+        final var prevKeyIndex = mArray.indexOfKey(key);
+        if (prevKeyIndex >= 0) {
+            final var prevValue = mArray.valueAt(prevKeyIndex);
+            if (prevValue == value) {
+                return this;
+            }
+        }
+        final var clone = mArray.clone();
+        clone.put(key, value);
+        return new ImmutableSparseArray<>(clone);
+    }
+
+    /**
+     * Returns an instance of {@link ImmutableSparseArray} that does not have the given key on top
+     * of items cloned from this instance.
+     *
+     * @param key the key to be removed
+     * @return    the same {@link ImmutableSparseArray} instance if there is actually no update.
+     *            Otherwise, a new instance of {@link ImmutableSparseArray}
+     */
+    @AnyThread
+    @NonNull
+    ImmutableSparseArray<E> cloneWithRemoveOrSelf(int key) {
+        final int index = indexOfKey(key);
+        if (index < 0) {
+            return this;
+        }
+        if (mArray.size() == 1) {
+            return empty();
+        }
+        final var clone = mArray.clone();
+        clone.remove(key);
+        return new ImmutableSparseArray<>(clone);
+    }
+}
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
index 9837ab1..03cbab5 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodBindingController.java
@@ -463,7 +463,7 @@
                     // should now try to restart the service for us.
                     mLastBindTime = SystemClock.uptimeMillis();
                     clearCurMethodAndSessions();
-                    mService.clearInputShownLocked();
+                    mService.mVisibilityStateComputer.setInputShown(false);
                     mService.unbindCurrentClientLocked(UnbindReason.DISCONNECT_IME, mUserId);
                 }
             }
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 4716e6c..7ff03c2 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -81,6 +81,7 @@
 import android.content.pm.ResolveInfo;
 import android.content.pm.ServiceInfo;
 import android.content.pm.UserInfo;
+import android.content.res.Configuration;
 import android.content.res.Resources;
 import android.hardware.input.InputManager;
 import android.inputmethodservice.InputMethodService;
@@ -253,12 +254,9 @@
     private @interface MultiUserUnawareField {
     }
 
-    private static final int MSG_SHOW_IM_SUBTYPE_PICKER = 1;
-
     private static final int MSG_HIDE_ALL_INPUT_METHODS = 1035;
     private static final int MSG_REMOVE_IME_SURFACE = 1060;
     private static final int MSG_REMOVE_IME_SURFACE_FROM_WINDOW = 1061;
-    private static final int MSG_UPDATE_IME_WINDOW_STATUS = 1070;
 
     private static final int MSG_RESET_HANDWRITING = 1090;
     private static final int MSG_START_HANDWRITING = 1100;
@@ -305,6 +303,28 @@
     private final String[] mNonPreemptibleInputMethods;
 
     /**
+     * Whether the new Input Method Switcher menu is enabled.
+     *
+     * @see #shouldEnableNewInputMethodSwitcherMenu
+     */
+    @SharedByAllUsersField
+    private final boolean mNewInputMethodSwitcherMenuEnabled;
+
+    /**
+     * Returns {@code true} if the new Input Method Switcher menu is enabled. This will be
+     * {@code false} for watches and small screen devices.
+     *
+     * @param context the context to check the device configuration for.
+     */
+    private static boolean shouldEnableNewInputMethodSwitcherMenu(@NonNull Context context) {
+        final boolean isWatch = context.getPackageManager()
+                .hasSystemFeature(PackageManager.FEATURE_WATCH);
+        final boolean isSmallScreen = (context.getResources().getConfiguration().screenLayout
+                & Configuration.SCREENLAYOUT_SIZE_MASK) == Configuration.SCREENLAYOUT_SIZE_SMALL;
+        return Flags.imeSwitcherRevamp() && !isWatch && !isSmallScreen;
+    }
+
+    /**
      * See {@link #shouldEnableConcurrentMultiUserMode(Context)} about when set to be {@code true}.
      */
     @SharedByAllUsersField
@@ -339,6 +359,35 @@
         return mConcurrentMultiUserModeEnabled ? callingProcessUserId : mCurrentUserId;
     }
 
+    /**
+     * Figures out the target IME user ID associated with the given {@code displayId}.
+     *
+     * @param displayId the display ID to be queried about
+     * @return User ID to be used for this {@code displayId}.
+     */
+    @GuardedBy("ImfLock.class")
+    @UserIdInt
+    private int resolveImeUserIdFromDisplayIdLocked(int displayId) {
+        return mConcurrentMultiUserModeEnabled
+                ? mUserManagerInternal.getUserAssignedToDisplay(displayId) : mCurrentUserId;
+    }
+
+    /**
+     * Figures out the target IME user ID associated with the given {@code windowToken}.
+     *
+     * @param windowToken the Window token to be queried about
+     * @return User ID to be used for this {@code displayId}.
+     */
+    @GuardedBy("ImfLock.class")
+    @UserIdInt
+    private int resolveImeUserIdFromWindowLocked(@NonNull IBinder windowToken) {
+        if (mConcurrentMultiUserModeEnabled) {
+            final int displayId = mWindowManagerInternal.getDisplayIdForWindow(windowToken);
+            return mUserManagerInternal.getUserAssignedToDisplay(displayId);
+        }
+        return mCurrentUserId;
+    }
+
     final Context mContext;
     final Resources mRes;
     private final Handler mHandler;
@@ -372,7 +421,7 @@
     @GuardedBy("ImfLock.class")
     @MultiUserUnawareField
     @NonNull
-    private final ImeVisibilityStateComputer mVisibilityStateComputer;
+    final ImeVisibilityStateComputer mVisibilityStateComputer;
 
     @GuardedBy("ImfLock.class")
     @SharedByAllUsersField
@@ -423,6 +472,9 @@
         IInputMethodSession mSession;
         InputChannel mChannel;
 
+        @UserIdInt
+        final int mUserId;
+
         @Override
         public String toString() {
             return "SessionState{uid=" + mClient.mUid + " pid=" + mClient.mPid
@@ -431,15 +483,17 @@
                     + " session=" + Integer.toHexString(
                     System.identityHashCode(mSession))
                     + " channel=" + mChannel
+                    + " userId=" + mUserId
                     + "}";
         }
 
         SessionState(ClientState client, IInputMethodInvoker method,
-                IInputMethodSession session, InputChannel channel) {
+                IInputMethodSession session, InputChannel channel, @UserIdInt int userId) {
             mClient = client;
             mMethod = method;
             mSession = session;
             mChannel = channel;
+            mUserId = userId;
         }
     }
 
@@ -495,14 +549,6 @@
     }
 
     /**
-     * The last window token that we confirmed that IME started talking to.  This is always updated
-     * upon reports from the input method.  If the window state is already changed before the report
-     * is handled, this field just keeps the last value.
-     */
-    @MultiUserUnawareField
-    IBinder mLastImeTargetWindow;
-
-    /**
      * Map of window perceptible states indexed by their associated window tokens.
      *
      * The value {@code true} indicates that IME has not been mostly hidden via
@@ -571,7 +617,7 @@
     private void onSecureSettingsChangedLocked(@NonNull String key, @UserIdInt int userId) {
         switch (key) {
             case Settings.Secure.SHOW_IME_WITH_HARD_KEYBOARD: {
-                if (!Flags.imeSwitcherRevamp()) {
+                if (!mNewInputMethodSwitcherMenuEnabled) {
                     if (userId == mCurrentUserId) {
                         mMenuController.updateKeyboardFromSettingsLocked();
                     }
@@ -640,7 +686,7 @@
                         }
                     }
                 }
-                if (Flags.imeSwitcherRevamp()) {
+                if (mNewInputMethodSwitcherMenuEnabled) {
                     synchronized (ImfLock.class) {
                         final var bindingController = getInputMethodBindingController(senderUserId);
                         mMenuControllerNew.hide(bindingController.getCurTokenDisplayId(),
@@ -944,7 +990,7 @@
             // TODO(b/196206770): Disallow I/O on this thread. Currently it's needed for loading
             // additional subtypes in switchUserOnHandlerLocked().
             final ServiceThread thread = new ServiceThread(HANDLER_THREAD_NAME,
-                    Process.THREAD_PRIORITY_FOREGROUND, true /* allowIo */);
+                    Process.THREAD_PRIORITY_FOREGROUND, false /* allowIo */);
             thread.start();
 
             final ServiceThread ioThread = new ServiceThread(PACKAGE_MONITOR_THREAD_NAME,
@@ -1046,10 +1092,8 @@
             SecureSettingsWrapper.onUserStarting(userId);
             mService.mIoHandler.post(() -> {
                 synchronized (ImfLock.class) {
-                    if (mService.mConcurrentMultiUserModeEnabled) {
-                        if (mService.mCurrentUserId != userId && mService.mSystemReady) {
-                            mService.initializeVisibleBackgroundUserLocked(userId);
-                        }
+                    if (mService.mSystemReady) {
+                        mService.onUserReadyLocked(userId);
                     }
                 }
             });
@@ -1065,8 +1109,8 @@
 
                 for (int userId : userIds) {
                     Slog.d(TAG, "Start initialization for user=" + userId);
-                    final var additionalSubtypeMap =
-                            AdditionalSubtypeMapRepository.ensureInitializedAndGet(userId);
+                    AdditionalSubtypeMapRepository.initializeIfNecessary(userId);
+                    final var additionalSubtypeMap = AdditionalSubtypeMapRepository.get(userId);
                     final var settings = InputMethodManagerService.queryInputMethodServicesInternal(
                             context, userId, additionalSubtypeMap,
                             DirectBootAwareness.AUTO).getMethodMap();
@@ -1095,7 +1139,7 @@
                 mUserSwitchHandlerTask.mClientToBeReset = clientToBeReset;
                 return;
             }
-            mHandler.removeCallbacks(mUserSwitchHandlerTask);
+            mIoHandler.removeCallbacks(mUserSwitchHandlerTask);
         }
         // Hide soft input before user switch task since switch task may block main handler a while
         // and delayed the hideCurrentInputLocked().
@@ -1105,7 +1149,7 @@
         final UserSwitchHandlerTask task = new UserSwitchHandlerTask(this, userId,
                 clientToBeReset);
         mUserSwitchHandlerTask = task;
-        mHandler.post(task);
+        mIoHandler.post(task);
     }
 
     @VisibleForTesting
@@ -1133,6 +1177,7 @@
             mUserManagerInternal = LocalServices.getService(UserManagerInternal.class);
 
             mSlotIme = mContext.getString(com.android.internal.R.string.status_bar_ime);
+            mNewInputMethodSwitcherMenuEnabled = shouldEnableNewInputMethodSwitcherMenu(mContext);
 
             mShowOngoingImeSwitcherForPhones = false;
 
@@ -1145,7 +1190,7 @@
                             : bindingControllerFactory);
 
             mMenuController = new InputMethodMenuController(this);
-            mMenuControllerNew = Flags.imeSwitcherRevamp()
+            mMenuControllerNew = mNewInputMethodSwitcherMenuEnabled
                     ? new InputMethodMenuControllerNew() : null;
             mVisibilityStateComputer = new ImeVisibilityStateComputer(this);
             mVisibilityApplier = new DefaultImeVisibilityApplier(this);
@@ -1386,30 +1431,32 @@
                         UserHandle.ALL, broadcastFilterForAllUsers, null, null,
                         Context.RECEIVER_EXPORTED);
 
-                final String defaultImiId = SecureSettingsWrapper.getString(
-                        Settings.Secure.DEFAULT_INPUT_METHOD, null, currentUserId);
-                final boolean imeSelectedOnBoot = !TextUtils.isEmpty(defaultImiId);
-                final var settings = InputMethodSettingsRepository.get(currentUserId);
-                postInputMethodSettingUpdatedLocked(
-                        !imeSelectedOnBoot /* resetDefaultEnabledIme */, currentUserId);
-                updateFromSettingsLocked(true, currentUserId);
-                InputMethodUtils.setNonSelectedSystemImesDisabledUntilUsed(
-                        getPackageManagerForUser(mContext, currentUserId),
-                        settings.getEnabledInputMethodList());
-
                 AdditionalSubtypeMapRepository.startWriterThread();
 
-                if (mConcurrentMultiUserModeEnabled) {
-                    for (int userId : mUserManagerInternal.getUserIds()) {
-                        if (userId != mCurrentUserId) {
-                            initializeVisibleBackgroundUserLocked(userId);
-                        }
-                    }
+                for (int userId : mUserManagerInternal.getUserIds()) {
+                    onUserReadyLocked(userId);
                 }
             }
         }
     }
 
+    @GuardedBy("ImfLock.class")
+    void onUserReadyLocked(@UserIdInt int userId) {
+        if (!mUserManagerInternal.isUserRunning(userId)) {
+            return;
+        }
+
+        final String defaultImiId = SecureSettingsWrapper.getString(
+                Settings.Secure.DEFAULT_INPUT_METHOD, null, userId);
+        final boolean imeSelectedOnBoot = !TextUtils.isEmpty(defaultImiId);
+        final var settings = InputMethodSettingsRepository.get(userId);
+        postInputMethodSettingUpdatedLocked(!imeSelectedOnBoot /* resetDefaultEnabledIme */,
+                userId);
+        updateFromSettingsLocked(true, userId);
+        InputMethodUtils.setNonSelectedSystemImesDisabledUntilUsed(
+                getPackageManagerForUser(mContext, userId), settings.getEnabledInputMethodList());
+    }
+
     void registerImeRequestedChangedListener() {
         mWindowManagerInternal.setOnImeRequestedChangedListener(
                 (windowToken, imeVisible, statsToken) -> {
@@ -1764,7 +1811,7 @@
                     ImeTracker.PHASE_SERVER_WAIT_IME);
             userData.mCurStatsToken = null;
             // TODO: Make mMenuController multi-user aware
-            if (Flags.imeSwitcherRevamp()) {
+            if (mNewInputMethodSwitcherMenuEnabled) {
                 mMenuControllerNew.hide(bindingController.getCurTokenDisplayId(), userId);
             } else {
                 mMenuController.hideInputMethodMenuLocked();
@@ -1807,17 +1854,6 @@
     }
 
     @GuardedBy("ImfLock.class")
-    void clearInputShownLocked() {
-        mVisibilityStateComputer.setInputShown(false);
-    }
-
-    @GuardedBy("ImfLock.class")
-    @Override
-    public boolean isInputShownLocked() {
-        return mVisibilityStateComputer.isInputShown();
-    }
-
-    @GuardedBy("ImfLock.class")
     private boolean isShowRequestedForCurrentWindow(@UserIdInt int userId) {
         final var userData = getUserData(userId);
         // TODO(b/349904272): Make mVisibilityStateComputer multi-user aware
@@ -2316,7 +2352,7 @@
                     if (userData.mCurClient != null) {
                         clearClientSessionLocked(userData.mCurClient);
                         userData.mCurClient.mCurSession = new SessionState(
-                                userData.mCurClient, method, session, channel);
+                                userData.mCurClient, method, session, channel, userId);
                         InputBindResult res = attachNewInputLocked(
                                 StartInputReason.SESSION_CREATED_BY_IME, true, userId);
                         attachNewAccessibilityLocked(StartInputReason.SESSION_CREATED_BY_IME, true,
@@ -2457,9 +2493,7 @@
                     sessionState.mSession.finishSession();
                 } catch (RemoteException e) {
                     Slog.w(TAG, "Session failed to close due to remote exception", e);
-                    // TODO(b/350386877): Propagate userId from the caller or infer it from
-                    //  sessionState
-                    final int userId = mCurrentUserId;
+                    final int userId = sessionState.mUserId;
                     final var bindingController = getInputMethodBindingController(userId);
                     updateSystemUiLocked(0 /* vis */, bindingController.getBackDisposition(),
                             userId);
@@ -2590,7 +2624,7 @@
         if (!mShowOngoingImeSwitcherForPhones) return false;
         // When the IME switcher dialog is shown, the IME switcher button should be hidden.
         // TODO(b/305849394): Make mMenuController multi-user aware.
-        final boolean switcherMenuShowing = Flags.imeSwitcherRevamp()
+        final boolean switcherMenuShowing = mNewInputMethodSwitcherMenuEnabled
                 ? mMenuControllerNew.isShowing()
                 : mMenuController.getSwitchingDialogLocked() != null;
         if (switcherMenuShowing) {
@@ -2610,7 +2644,8 @@
                 || (visibility & InputMethodService.IME_INVISIBLE) != 0) {
             return false;
         }
-        if (mWindowManagerInternal.isHardKeyboardAvailable() && !Flags.imeSwitcherRevamp()) {
+        if (mWindowManagerInternal.isHardKeyboardAvailable()
+                && !mNewInputMethodSwitcherMenuEnabled) {
             // When physical keyboard is attached, we show the ime switcher (or notification if
             // NavBar is not available) because SHOW_IME_WITH_HARD_KEYBOARD settings currently
             // exists in the IME switcher dialog.  Might be OK to remove this condition once
@@ -2621,7 +2656,7 @@
         }
 
         final InputMethodSettings settings = InputMethodSettingsRepository.get(userId);
-        if (Flags.imeSwitcherRevamp()) {
+        if (mNewInputMethodSwitcherMenuEnabled) {
             // The IME switcher button should be shown when the current IME specified a
             // language settings activity.
             final var curImi = settings.getMethodMap().get(settings.getSelectedInputMethod());
@@ -2740,20 +2775,18 @@
             if (targetWindow != null) {
                 mWindowManagerInternal.updateInputMethodTargetWindow(token, targetWindow);
             }
-            mLastImeTargetWindow = targetWindow;
+            mVisibilityStateComputer.setLastImeTargetWindow(targetWindow);
         }
     }
 
-    private void updateImeWindowStatus(boolean disableImeIcon) {
-        synchronized (ImfLock.class) {
-            // TODO(b/350386877): Propagate userId from the caller.
-            final int userId = mCurrentUserId;
-            if (disableImeIcon) {
-                final var bindingController = getInputMethodBindingController(userId);
-                updateSystemUiLocked(0, bindingController.getBackDisposition(), userId);
-            } else {
-                updateSystemUiLocked(userId);
-            }
+    @GuardedBy("ImfLock.class")
+    private void updateImeWindowStatusLocked(boolean disableImeIcon, int displayId) {
+        final int userId = resolveImeUserIdFromDisplayIdLocked(displayId);
+        if (disableImeIcon) {
+            final var bindingController = getInputMethodBindingController(userId);
+            updateSystemUiLocked(0, bindingController.getBackDisposition(), userId);
+        } else {
+            updateSystemUiLocked(userId);
         }
     }
 
@@ -2799,7 +2832,7 @@
             }
             final var curId = bindingController.getCurId();
             // TODO(b/305849394): Make mMenuController multi-user aware.
-            final boolean switcherMenuShowing = Flags.imeSwitcherRevamp()
+            final boolean switcherMenuShowing = mNewInputMethodSwitcherMenuEnabled
                     ? mMenuControllerNew.isShowing()
                     : mMenuController.getSwitchingDialogLocked() != null;
             if (switcherMenuShowing
@@ -2821,65 +2854,11 @@
     @GuardedBy("ImfLock.class")
     void updateFromSettingsLocked(boolean enabledMayChange, @UserIdInt int userId) {
         updateInputMethodsFromSettingsLocked(enabledMayChange, userId);
-        if (!Flags.imeSwitcherRevamp()) {
+        if (!mNewInputMethodSwitcherMenuEnabled) {
             mMenuController.updateKeyboardFromSettingsLocked();
         }
     }
 
-    /**
-     * This initialization logic is used when and only when {@link #mConcurrentMultiUserModeEnabled}
-     * is set to {@code true}.
-     *
-     * <p>There remain several yet-to-be-implemented features. For the canonical and desired
-     * behaviors always refer to single-user code paths such as
-     * {@link #updateInputMethodsFromSettingsLocked(boolean, int)}.</p>
-     *
-     * <p>Here are examples of missing features.</p>
-     * <ul>
-     *     <li>Profiles are not supported.</li>
-     *     <li>
-     *         {@link PackageManager#COMPONENT_ENABLED_STATE_DISABLED_UNTIL_USED} is not updated.
-     *     </li>
-     *     <li>{@link InputMethodBindingController#getDeviceIdToShowIme()} is ignored.</li>
-     *     <li>and so on.</li>
-     * </ul>
-     */
-    @GuardedBy("ImfLock.class")
-    void initializeVisibleBackgroundUserLocked(@UserIdInt int userId) {
-        final var settings = InputMethodSettingsRepository.get(userId);
-
-        // Until we figure out what makes most sense, we enable all the pre-installed IMEs in
-        // concurrent multi-user IME mode.
-        String enabledImeIdsStr = settings.getEnabledInputMethodsStr();
-        for (var imi : settings.getMethodList()) {
-            if (!imi.isSystem()) {
-                continue;
-            }
-            enabledImeIdsStr = InputMethodUtils.concatEnabledImeIds(enabledImeIdsStr, imi.getId());
-        }
-        if (!TextUtils.equals(settings.getEnabledInputMethodsStr(), enabledImeIdsStr)) {
-            settings.putEnabledInputMethodsStr(enabledImeIdsStr);
-        }
-
-        // Also update the currently-selected IME.
-        String id = settings.getSelectedInputMethod();
-        if (TextUtils.isEmpty(id)) {
-            final InputMethodInfo imi = InputMethodInfoUtils.getMostApplicableDefaultIME(
-                    settings.getEnabledInputMethodList());
-            if (imi != null) {
-                id = imi.getId();
-                settings.putSelectedInputMethod(id);
-            }
-        }
-        final var userData = getUserData(userId);
-        final var bindingController = userData.mBindingController;
-        bindingController.setSelectedMethodId(id);
-
-        // Also re-initialize controllers.
-        userData.mSwitchingController.resetCircularListLocked(mContext, settings);
-        userData.mHardwareKeyboardShortcutController.update(settings);
-    }
-
     @GuardedBy("ImfLock.class")
     void updateInputMethodsFromSettingsLocked(boolean enabledMayChange, @UserIdInt int userId) {
         final InputMethodSettings settings = InputMethodSettingsRepository.get(userId);
@@ -3058,52 +3037,75 @@
         }
     }
 
+    @GuardedBy("ImfLock.class")
+    private void sendResultReceiverFailureLocked(@Nullable ResultReceiver resultReceiver) {
+        final boolean isInputShown = mVisibilityStateComputer.isInputShown();
+        resultReceiver.send(isInputShown
+                ? InputMethodManager.RESULT_UNCHANGED_SHOWN
+                : InputMethodManager.RESULT_UNCHANGED_HIDDEN, null);
+    }
+
     @Override
     public boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
             int lastClickToolType, ResultReceiver resultReceiver,
             @SoftInputShowHideReason int reason) {
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "IMMS.showSoftInput");
-        final int uid = Binder.getCallingUid();
-        final int callingUserId = UserHandle.getUserId(uid);
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#showSoftInput", mDumper);
         synchronized (ImfLock.class) {
-            final int userId = resolveImeUserIdLocked(callingUserId);
-            if (!canInteractWithImeLocked(uid, client, "showSoftInput", statsToken,
-                    userId)) {
-                ImeTracker.forLogging().onFailed(
-                        statsToken, ImeTracker.PHASE_SERVER_CLIENT_FOCUSED);
-                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
-                return false;
+            final boolean result = showSoftInputLocked(client, windowToken, statsToken, flags,
+                    lastClickToolType, resultReceiver, reason);
+            // When ZeroJankProxy is enabled, the app has already received "true" as the return
+            // value, and expect "resultReceiver" to be notified later. See b/327751155.
+            if (!result && Flags.useZeroJankProxy()) {
+                sendResultReceiverFailureLocked(resultReceiver);
             }
-            final long ident = Binder.clearCallingIdentity();
-            final var userData = getUserData(userId);
-            try {
-                if (DEBUG) Slog.v(TAG, "Client requesting input be shown");
-                if (Flags.refactorInsetsController()) {
-                    boolean wasVisible = isInputShownLocked();
-                    if (userData.mImeBindingState != null
-                            && userData.mImeBindingState.mFocusedWindowClient != null
-                            && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
-                        userData.mImeBindingState.mFocusedWindowClient.mClient
-                                .setImeVisibility(true, statsToken);
-                        if (resultReceiver != null) {
-                            resultReceiver.send(
-                                    wasVisible ? InputMethodManager.RESULT_UNCHANGED_SHOWN
-                                            : InputMethodManager.RESULT_SHOWN, null);
-                        }
-                        return true;
+            return result;  // ignored when ZeroJankProxy is enabled.
+        }
+    }
+
+    @GuardedBy("ImfLock.class")
+    private boolean showSoftInputLocked(IInputMethodClient client, IBinder windowToken,
+            @NonNull ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
+            int lastClickToolType, ResultReceiver resultReceiver,
+            @SoftInputShowHideReason int reason) {
+        final int uid = Binder.getCallingUid();
+        final int callingUserId = UserHandle.getUserId(uid);
+        final int userId = resolveImeUserIdLocked(callingUserId);
+        if (!canInteractWithImeLocked(uid, client, "showSoftInput", statsToken,
+                userId)) {
+            ImeTracker.forLogging().onFailed(
+                    statsToken, ImeTracker.PHASE_SERVER_CLIENT_FOCUSED);
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+            return false;
+        }
+        final var userData = getUserData(userId);
+        final long ident = Binder.clearCallingIdentity();
+        try {
+            if (DEBUG) Slog.v(TAG, "Client requesting input be shown");
+            if (Flags.refactorInsetsController()) {
+                boolean wasVisible = mVisibilityStateComputer.isInputShown();
+                if (userData.mImeBindingState != null
+                        && userData.mImeBindingState.mFocusedWindowClient != null
+                        && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
+                    userData.mImeBindingState.mFocusedWindowClient.mClient
+                            .setImeVisibility(true, statsToken);
+                    if (resultReceiver != null) {
+                        resultReceiver.send(
+                                wasVisible ? InputMethodManager.RESULT_UNCHANGED_SHOWN
+                                        : InputMethodManager.RESULT_SHOWN, null);
                     }
-                    return false;
-                } else {
-                    return showCurrentInputLocked(windowToken, statsToken, flags, lastClickToolType,
-                            resultReceiver, reason, userId);
+                    return true;
                 }
-            } finally {
-                Binder.restoreCallingIdentity(ident);
-                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+                return false;
+            } else {
+                return showCurrentInputLocked(windowToken, statsToken, flags, lastClickToolType,
+                        resultReceiver, reason, userId);
             }
+        } finally {
+            Binder.restoreCallingIdentity(ident);
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         }
     }
 
@@ -3113,8 +3115,7 @@
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#showSoftInput", mDumper);
         synchronized (ImfLock.class) {
-            // TODO(b/305849394): Infer userId from windowToken
-            final int userId = mCurrentUserId;
+            final int userId = resolveImeUserIdFromWindowLocked(windowToken);
             final long ident = Binder.clearCallingIdentity();
             try {
                 if (DEBUG) Slog.v(TAG, "Client requesting input be shown");
@@ -3134,8 +3135,7 @@
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#hideSoftInput", mDumper);
         synchronized (ImfLock.class) {
-            // TODO(b/305849394): Infer userId from windowToken
-            final int userId = mCurrentUserId;
+            final int userId = resolveImeUserIdFromWindowLocked(windowToken);
             final long ident = Binder.clearCallingIdentity();
             try {
                 if (DEBUG) Slog.v(TAG, "Client requesting input be hidden");
@@ -3397,14 +3397,14 @@
             Objects.requireNonNull(windowToken, "windowToken must not be null");
             synchronized (ImfLock.class) {
                 Boolean windowPerceptible = mFocusedWindowPerceptible.get(windowToken);
-                final int userId = mCurrentUserId;
+                final int userId = resolveImeUserIdFromWindowLocked(windowToken);
                 final var userData = getUserData(userId);
                 if (userData.mImeBindingState.mFocusedWindow != windowToken
                         || (windowPerceptible != null && windowPerceptible == perceptible)) {
                     return;
                 }
                 mFocusedWindowPerceptible.put(windowToken, windowPerceptible);
-                updateSystemUiLocked(mCurrentUserId);
+                updateSystemUiLocked(userId);
             }
         });
     }
@@ -3498,50 +3498,64 @@
     public boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
             ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
-        final int uid = Binder.getCallingUid();
-        final int callingUserId = UserHandle.getUserId(uid);
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#hideSoftInput", mDumper);
         synchronized (ImfLock.class) {
-            final int userId = resolveImeUserIdLocked(callingUserId);
-            if (!canInteractWithImeLocked(uid, client, "hideSoftInput", statsToken, userId)) {
-                if (isInputShownLocked()) {
-                    ImeTracker.forLogging().onFailed(
-                            statsToken, ImeTracker.PHASE_SERVER_CLIENT_FOCUSED);
-                } else {
-                    ImeTracker.forLogging().onCancelled(statsToken,
-                            ImeTracker.PHASE_SERVER_CLIENT_FOCUSED);
+            final boolean result = hideSoftInputLocked(client, windowToken, statsToken, flags,
+                    resultReceiver, reason);
+            // When ZeroJankProxy is enabled, the app has already received "true" as the return
+            // value, and expect "resultReceiver" to be notified later. See b/327751155.
+            if (!result && Flags.useZeroJankProxy()) {
+                sendResultReceiverFailureLocked(resultReceiver);
+            }
+            return result;  // ignored when ZeroJankProxy is enabled.
+        }
+    }
+
+    @GuardedBy("ImfLock.class")
+    private boolean hideSoftInputLocked(IInputMethodClient client, IBinder windowToken,
+            @NonNull ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
+            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
+        final int uid = Binder.getCallingUid();
+        final int callingUserId = UserHandle.getUserId(uid);
+        final int userId = resolveImeUserIdLocked(callingUserId);
+        if (!canInteractWithImeLocked(uid, client, "hideSoftInput", statsToken, userId)) {
+            if (mVisibilityStateComputer.isInputShown()) {
+                ImeTracker.forLogging().onFailed(
+                        statsToken, ImeTracker.PHASE_SERVER_CLIENT_FOCUSED);
+            } else {
+                ImeTracker.forLogging().onCancelled(statsToken,
+                        ImeTracker.PHASE_SERVER_CLIENT_FOCUSED);
+            }
+            return false;
+        }
+        final var userData = getUserData(userId);
+        final long ident = Binder.clearCallingIdentity();
+        try {
+            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "IMMS.hideSoftInput");
+            if (DEBUG) Slog.v(TAG, "Client requesting input be hidden");
+            if (Flags.refactorInsetsController()) {
+                if (userData.mImeBindingState != null
+                        && userData.mImeBindingState.mFocusedWindowClient != null
+                        && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
+                    boolean wasVisible = mVisibilityStateComputer.isInputShown();
+                    // TODO add windowToken to interface
+                    userData.mImeBindingState.mFocusedWindowClient.mClient
+                            .setImeVisibility(false, statsToken);
+                    if (resultReceiver != null) {
+                        resultReceiver.send(wasVisible ? InputMethodManager.RESULT_HIDDEN
+                                : InputMethodManager.RESULT_UNCHANGED_HIDDEN, null);
+                    }
+                    return true;
                 }
                 return false;
+            } else {
+                return InputMethodManagerService.this.hideCurrentInputLocked(
+                        windowToken, statsToken, flags, resultReceiver, reason, userId);
             }
-            final long ident = Binder.clearCallingIdentity();
-            final var userData = getUserData(userId);
-            try {
-                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "IMMS.hideSoftInput");
-                if (DEBUG) Slog.v(TAG, "Client requesting input be hidden");
-                if (Flags.refactorInsetsController()) {
-                    if (userData.mImeBindingState != null
-                            && userData.mImeBindingState.mFocusedWindowClient != null
-                            && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
-                        boolean wasVisible = isInputShownLocked();
-                        // TODO add windowToken to interface
-                        userData.mImeBindingState.mFocusedWindowClient.mClient
-                                .setImeVisibility(false, statsToken);
-                        if (resultReceiver != null) {
-                            resultReceiver.send(wasVisible ? InputMethodManager.RESULT_HIDDEN
-                                    : InputMethodManager.RESULT_UNCHANGED_HIDDEN, null);
-                        }
-                        return true;
-                    }
-                    return false;
-                } else {
-                    return InputMethodManagerService.this.hideCurrentInputLocked(windowToken,
-                            statsToken, flags, resultReceiver, reason, userId);
-                }
-            } finally {
-                Binder.restoreCallingIdentity(ident);
-                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
-            }
+        } finally {
+            Binder.restoreCallingIdentity(ident);
+            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
         }
     }
 
@@ -3586,7 +3600,7 @@
         // TODO(b/246309664): Clean up IMMS#mImeWindowVis
         IInputMethodInvoker curMethod = bindingController.getCurMethod();
         final boolean shouldHideSoftInput = curMethod != null
-                && (isInputShownLocked()
+                && (mVisibilityStateComputer.isInputShown()
                 || (bindingController.getImeWindowVis() & InputMethodService.IME_ACTIVE) != 0);
 
         mVisibilityStateComputer.requestImeVisibility(windowToken, false);
@@ -3934,10 +3948,9 @@
     }
 
     @GuardedBy("ImfLock.class")
-    private boolean canShowInputMethodPickerLocked(IInputMethodClient client) {
+    private boolean canShowInputMethodPickerLocked(IInputMethodClient client,
+            @UserIdInt int userId) {
         final int uid = Binder.getCallingUid();
-        // TODO(b/305849394): Get userId from callers.
-        final int userId = mCurrentUserId;
         final var userData = getUserData(userId);
         if (userData.mImeBindingState.mFocusedWindowClient != null && client != null
                 && userData.mImeBindingState.mFocusedWindowClient.mClient.asBinder()
@@ -3964,19 +3977,22 @@
         }
         final int callingUserId = UserHandle.getCallingUserId();
         synchronized (ImfLock.class) {
-            if (!canShowInputMethodPickerLocked(client)) {
+            final int userId = resolveImeUserIdLocked(callingUserId);
+            if (!canShowInputMethodPickerLocked(client, userId)) {
                 Slog.w(TAG, "Ignoring showInputMethodPickerFromClient of uid "
                         + Binder.getCallingUid() + ": " + client);
                 return;
             }
-            final int userId = resolveImeUserIdLocked(callingUserId);
             final var userData = getUserData(userId);
             // Always call subtype picker, because subtype picker is a superset of input method
             // picker.
             final int displayId = (userData.mCurClient != null)
                     ? userData.mCurClient.mSelfReportedDisplayId : DEFAULT_DISPLAY;
-            mHandler.obtainMessage(MSG_SHOW_IM_SUBTYPE_PICKER, auxiliarySubtypeMode, displayId)
-                    .sendToTarget();
+            mHandler.post(() -> {
+                synchronized (ImfLock.class) {
+                    showInputMethodPickerLocked(auxiliarySubtypeMode, displayId, userId);
+                }
+            });
         }
     }
 
@@ -3985,8 +4001,12 @@
     public void showInputMethodPickerFromSystem(int auxiliarySubtypeMode, int displayId) {
         // Always call subtype picker, because subtype picker is a superset of input method
         // picker.
-        mHandler.obtainMessage(MSG_SHOW_IM_SUBTYPE_PICKER, auxiliarySubtypeMode, displayId)
-                .sendToTarget();
+        mHandler.post(() -> {
+            synchronized (ImfLock.class) {
+                final int userId = resolveImeUserIdFromDisplayIdLocked(displayId);
+                showInputMethodPickerLocked(auxiliarySubtypeMode, displayId, userId);
+            }
+        });
     }
 
     /**
@@ -3995,7 +4015,7 @@
     @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
     public boolean isInputMethodPickerShownForTest() {
         synchronized (ImfLock.class) {
-            return Flags.imeSwitcherRevamp()
+            return mNewInputMethodSwitcherMenuEnabled
                     ? mMenuControllerNew.isShowing()
                     : mMenuController.isisInputMethodPickerShownForTestLocked();
         }
@@ -4074,7 +4094,7 @@
     @Override
     public void onImeSwitchButtonClickFromSystem(int displayId) {
         synchronized (ImfLock.class) {
-            final int userId = mCurrentUserId;
+            final int userId = resolveImeUserIdFromDisplayIdLocked(displayId);
             final var userData = getUserData(userId);
             final var bindingController = userData.mBindingController;
             final var curToken = bindingController.getCurToken();
@@ -4414,7 +4434,7 @@
 
     @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
     @Override
-    public void removeImeSurface() {
+    public void removeImeSurface(int displayId) {
         mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE).sendToTarget();
     }
 
@@ -4683,8 +4703,8 @@
             proto.write(CUR_SEQ, bindingController.getSequenceNumber());
             proto.write(CUR_CLIENT, Objects.toString(userData.mCurClient));
             userData.mImeBindingState.dumpDebug(proto, mWindowManagerInternal);
-            proto.write(LAST_IME_TARGET_WINDOW_NAME,
-                    mWindowManagerInternal.getWindowName(mLastImeTargetWindow));
+            proto.write(LAST_IME_TARGET_WINDOW_NAME, mWindowManagerInternal.getWindowName(
+                    mVisibilityStateComputer.getLastImeTargetWindow()));
             proto.write(CUR_FOCUSED_WINDOW_SOFT_INPUT_MODE, InputMethodDebug.softInputModeToString(
                     userData.mImeBindingState.mFocusedWindowSoftInputMode));
             if (userData.mCurEditorInfo != null) {
@@ -4701,7 +4721,7 @@
             proto.write(IS_INTERACTIVE, mIsInteractive);
             proto.write(BACK_DISPOSITION, bindingController.getBackDisposition());
             proto.write(IME_WINDOW_VISIBILITY, bindingController.getImeWindowVis());
-            if (!Flags.imeSwitcherRevamp()) {
+            if (!mNewInputMethodSwitcherMenuEnabled) {
                 proto.write(SHOW_IME_WITH_HARD_KEYBOARD,
                         mMenuController.getShowImeWithHardKeyboard());
             }
@@ -4858,8 +4878,8 @@
                                     .setImeVisibility(false, statsToken);
                         }
                     } else {
-                        hideCurrentInputLocked(mLastImeTargetWindow, statsToken, flags,
-                                null /* resultReceiver */, reason, userId);
+                        hideCurrentInputLocked(mVisibilityStateComputer.getLastImeTargetWindow(),
+                                statsToken, flags, null /* resultReceiver */, reason, userId);
                     }
                 } finally {
                     Binder.restoreCallingIdentity(ident);
@@ -4897,9 +4917,9 @@
                                     .setImeVisibility(true, statsToken);
                         }
                     } else {
-                        showCurrentInputLocked(mLastImeTargetWindow, statsToken, flags,
-                                MotionEvent.TOOL_TYPE_UNKNOWN, null /* resultReceiver */, reason,
-                                userId);
+                        showCurrentInputLocked(mVisibilityStateComputer.getLastImeTargetWindow(),
+                                statsToken, flags, MotionEvent.TOOL_TYPE_UNKNOWN,
+                                null /* resultReceiver */, reason, userId);
                     }
                 } finally {
                     Binder.restoreCallingIdentity(ident);
@@ -4916,14 +4936,12 @@
         return mVisibilityApplier;
     }
 
-    void onApplyImeVisibilityFromComputer(IBinder windowToken, @NonNull ImeTracker.Token statsToken,
-            @NonNull ImeVisibilityResult result) {
-        synchronized (ImfLock.class) {
-            // TODO(b/305849394): Infer userId from windowToken
-            final int userId = mCurrentUserId;
-            mVisibilityApplier.applyImeVisibility(windowToken, statsToken, result.getState(),
-                    result.getReason(), userId);
-        }
+    @GuardedBy("ImfLock.class")
+    void onApplyImeVisibilityFromComputerLocked(IBinder windowToken,
+            @NonNull ImeTracker.Token statsToken, @NonNull ImeVisibilityResult result) {
+        final int userId = resolveImeUserIdFromWindowLocked(windowToken);
+        mVisibilityApplier.applyImeVisibility(windowToken, statsToken, result.getState(),
+                result.getReason(), userId);
     }
 
     @GuardedBy("ImfLock.class")
@@ -4979,89 +4997,80 @@
         userData.mEnabledAccessibilitySessions = accessibilitySessions;
     }
 
+    @GuardedBy("ImfLock.class")
+    private void showInputMethodPickerLocked(int auxiliarySubtypeMode, int displayId,
+            @UserIdInt int userId) {
+        final boolean showAuxSubtypes;
+        switch (auxiliarySubtypeMode) {
+            // This is undocumented so far, but IMM#showInputMethodPicker() has been
+            // implemented so that auxiliary subtypes will be excluded when the soft
+            // keyboard is invisible.
+            case InputMethodManager.SHOW_IM_PICKER_MODE_AUTO ->
+                    showAuxSubtypes = mVisibilityStateComputer.isInputShown();
+            case InputMethodManager.SHOW_IM_PICKER_MODE_INCLUDE_AUXILIARY_SUBTYPES ->
+                    showAuxSubtypes = true;
+            case InputMethodManager.SHOW_IM_PICKER_MODE_EXCLUDE_AUXILIARY_SUBTYPES ->
+                    showAuxSubtypes = false;
+            default -> {
+                Slog.e(TAG, "Unknown subtype picker mode=" + auxiliarySubtypeMode);
+                return;
+            }
+        }
+        final InputMethodSettings settings = InputMethodSettingsRepository.get(userId);
+        final boolean isScreenLocked = mWindowManagerInternal.isKeyguardLocked()
+                && mWindowManagerInternal.isKeyguardSecure(userId);
+        final String lastInputMethodId = settings.getSelectedInputMethod();
+        int lastInputMethodSubtypeId = settings.getSelectedInputMethodSubtypeId(lastInputMethodId);
+
+        final List<ImeSubtypeListItem> imList = InputMethodSubtypeSwitchingController
+                .getSortedInputMethodAndSubtypeList(
+                        showAuxSubtypes, isScreenLocked, true /* forImeMenu */,
+                        mContext, settings);
+        if (imList.isEmpty()) {
+            Slog.w(TAG, "Show switching menu failed, imList is empty,"
+                    + " showAuxSubtypes: " + showAuxSubtypes
+                    + " isScreenLocked: " + isScreenLocked
+                    + " userId: " + userId);
+            return;
+        }
+
+        if (Flags.imeSwitcherRevamp()) {
+            if (DEBUG) {
+                Slog.v(TAG, "Show IME switcher menu,"
+                        + " showAuxSubtypes=" + showAuxSubtypes
+                        + " displayId=" + displayId
+                        + " preferredInputMethodId=" + lastInputMethodId
+                        + " preferredInputMethodSubtypeId=" + lastInputMethodSubtypeId);
+            }
+
+            final var itemsAndIndex = getInputMethodPickerItems(imList,
+                    lastInputMethodId, lastInputMethodSubtypeId, userId);
+            final var menuItems = itemsAndIndex.first;
+            final int selectedIndex = itemsAndIndex.second;
+
+            if (selectedIndex == -1) {
+                Slog.w(TAG, "Switching menu shown with no item selected"
+                        + ", IME id: " + lastInputMethodId
+                        + ", subtype index: " + lastInputMethodSubtypeId);
+            }
+
+            mMenuControllerNew.show(menuItems, selectedIndex, displayId, userId);
+        } else {
+            mMenuController.showInputMethodMenuLocked(showAuxSubtypes, displayId,
+                    lastInputMethodId, lastInputMethodSubtypeId, imList);
+        }
+    }
+
     @SuppressWarnings("unchecked")
     @UiThread
     @Override
     public boolean handleMessage(Message msg) {
         switch (msg.what) {
-            case MSG_SHOW_IM_SUBTYPE_PICKER:
-                final boolean showAuxSubtypes;
-                final int displayId = msg.arg2;
-                switch (msg.arg1) {
-                    case InputMethodManager.SHOW_IM_PICKER_MODE_AUTO:
-                        // This is undocumented so far, but IMM#showInputMethodPicker() has been
-                        // implemented so that auxiliary subtypes will be excluded when the soft
-                        // keyboard is invisible.
-                        synchronized (ImfLock.class) {
-                            showAuxSubtypes = isInputShownLocked();
-                        }
-                        break;
-                    case InputMethodManager.SHOW_IM_PICKER_MODE_INCLUDE_AUXILIARY_SUBTYPES:
-                        showAuxSubtypes = true;
-                        break;
-                    case InputMethodManager.SHOW_IM_PICKER_MODE_EXCLUDE_AUXILIARY_SUBTYPES:
-                        showAuxSubtypes = false;
-                        break;
-                    default:
-                        Slog.e(TAG, "Unknown subtype picker mode = " + msg.arg1);
-                        return false;
-                }
+            case MSG_HIDE_ALL_INPUT_METHODS: {
+                @SoftInputShowHideReason final int reason = msg.arg1;
+                final int originatingDisplayId = msg.arg2;
                 synchronized (ImfLock.class) {
-                    final InputMethodSettings settings =
-                            InputMethodSettingsRepository.get(mCurrentUserId);
-                    final int userId = settings.getUserId();
-                    final boolean isScreenLocked = mWindowManagerInternal.isKeyguardLocked()
-                            && mWindowManagerInternal.isKeyguardSecure(userId);
-                    final String lastInputMethodId = settings.getSelectedInputMethod();
-                    int lastInputMethodSubtypeId =
-                            settings.getSelectedInputMethodSubtypeId(lastInputMethodId);
-
-                    final List<ImeSubtypeListItem> imList = InputMethodSubtypeSwitchingController
-                            .getSortedInputMethodAndSubtypeList(
-                                    showAuxSubtypes, isScreenLocked, true /* forImeMenu */,
-                                    mContext, settings);
-                    if (imList.isEmpty()) {
-                        Slog.w(TAG, "Show switching menu failed, imList is empty,"
-                                + " showAuxSubtypes: " + showAuxSubtypes
-                                + " isScreenLocked: " + isScreenLocked
-                                + " userId: " + userId);
-                        return false;
-                    }
-
-                    if (Flags.imeSwitcherRevamp()) {
-                        if (DEBUG) {
-                            Slog.v(TAG, "Show IME switcher menu,"
-                                    + " showAuxSubtypes=" + showAuxSubtypes
-                                    + " displayId=" + displayId
-                                    + " preferredInputMethodId=" + lastInputMethodId
-                                    + " preferredInputMethodSubtypeId=" + lastInputMethodSubtypeId);
-                        }
-
-                        final var itemsAndIndex = getInputMethodPickerItems(imList,
-                                lastInputMethodId, lastInputMethodSubtypeId, userId);
-                        final var menuItems = itemsAndIndex.first;
-                        final int selectedIndex = itemsAndIndex.second;
-
-                        if (selectedIndex == -1) {
-                            Slog.w(TAG, "Switching menu shown with no item selected"
-                                    + ", IME id: " + lastInputMethodId
-                                    + ", subtype index: " + lastInputMethodSubtypeId);
-                        }
-
-                        mMenuControllerNew.show(menuItems, selectedIndex, displayId, userId);
-                    } else {
-                        mMenuController.showInputMethodMenuLocked(showAuxSubtypes, displayId,
-                                lastInputMethodId, lastInputMethodSubtypeId, imList);
-                    }
-                }
-                return true;
-
-            // ---------------------------------------------------------
-
-            case MSG_HIDE_ALL_INPUT_METHODS:
-                synchronized (ImfLock.class) {
-                    // TODO(b/305849394): Needs to figure out what to do where for background users.
-                    final int userId = mCurrentUserId;
+                    final int userId = resolveImeUserIdFromDisplayIdLocked(originatingDisplayId);
                     final var userData = getUserData(userId);
                     if (Flags.refactorInsetsController()) {
                         if (userData.mImeBindingState != null
@@ -5069,15 +5078,16 @@
                                 && userData.mImeBindingState.mFocusedWindowClient.mClient != null) {
                             userData.mImeBindingState.mFocusedWindowClient.mClient
                                     .setImeVisibility(false,
-                                    null /* TODO(b329229469) check statsToken */);
+                                            null /* TODO(b329229469) check statsToken */);
                         }
                     } else {
-                        @SoftInputShowHideReason final int reason = (int) msg.obj;
+
                         hideCurrentInputLocked(userData.mImeBindingState.mFocusedWindow,
                                 0 /* flags */, reason, userId);
                     }
                 }
                 return true;
+            }
             case MSG_REMOVE_IME_SURFACE: {
                 synchronized (ImfLock.class) {
                     // TODO(b/305849394): Needs to figure out what to do where for background users.
@@ -5097,8 +5107,7 @@
             case MSG_REMOVE_IME_SURFACE_FROM_WINDOW: {
                 IBinder windowToken = (IBinder) msg.obj;
                 synchronized (ImfLock.class) {
-                    // TODO(b/305849394): Infer userId from windowToken.
-                    final int userId = mCurrentUserId;
+                    final int userId = resolveImeUserIdFromWindowLocked(windowToken);
                     final var userData = getUserData(userId);
                     try {
                         if (windowToken == userData.mImeBindingState.mFocusedWindow
@@ -5111,10 +5120,6 @@
                 }
                 return true;
             }
-            case MSG_UPDATE_IME_WINDOW_STATUS: {
-                updateImeWindowStatus(msg.arg1 == 1);
-                return true;
-            }
 
             // ---------------------------------------------------------
 
@@ -5124,7 +5129,7 @@
 
             // --------------------------------------------------------------
             case MSG_HARD_KEYBOARD_SWITCH_CHANGED:
-                if (!Flags.imeSwitcherRevamp()) {
+                if (!mNewInputMethodSwitcherMenuEnabled) {
                     mMenuController.handleHardKeyboardStatusChange(msg.arg1 == 1);
                 }
                 synchronized (ImfLock.class) {
@@ -5800,7 +5805,8 @@
         public void hideAllInputMethods(@SoftInputShowHideReason int reason,
                 int originatingDisplayId) {
             mHandler.removeMessages(MSG_HIDE_ALL_INPUT_METHODS);
-            mHandler.obtainMessage(MSG_HIDE_ALL_INPUT_METHODS, reason).sendToTarget();
+            mHandler.obtainMessage(MSG_HIDE_ALL_INPUT_METHODS, reason, originatingDisplayId)
+                    .sendToTarget();
         }
 
         @ImfLockFree
@@ -5902,8 +5908,7 @@
         @Override
         public void reportImeControl(@Nullable IBinder windowToken) {
             synchronized (ImfLock.class) {
-                // TODO(b/305849394): Need to infer userId or get userId from callers.
-                final int userId = mCurrentUserId;
+                final int userId = resolveImeUserIdFromWindowLocked(windowToken);
                 final var userData = getUserData(userId);
                 if (userData.mImeBindingState.mFocusedWindow != windowToken) {
                     // A perceptible value was set for the focused window, but it is no longer in
@@ -5918,14 +5923,14 @@
         @Override
         public void onImeParentChanged(int displayId) {
             synchronized (ImfLock.class) {
-                // TODO(b/305849394): Need to infer userId or get userId from callers.
-                final int userId = mCurrentUserId;
+                final int userId = resolveImeUserIdFromDisplayIdLocked(displayId);
                 final var userData = getUserData(userId);
                 // Hide the IME method menu only when the IME surface parent is changed by the
                 // input target changed, in case seeing the dialog dismiss flickering during
                 // the next focused window starting the input connection.
-                if (mLastImeTargetWindow != userData.mImeBindingState.mFocusedWindow) {
-                    if (Flags.imeSwitcherRevamp()) {
+                if (mVisibilityStateComputer.getLastImeTargetWindow()
+                        != userData.mImeBindingState.mFocusedWindow) {
+                    if (mNewInputMethodSwitcherMenuEnabled) {
                         final var bindingController = getInputMethodBindingController(userId);
                         mMenuControllerNew.hide(bindingController.getCurTokenDisplayId(), userId);
                     } else {
@@ -5944,8 +5949,11 @@
         @ImfLockFree
         @Override
         public void updateImeWindowStatus(boolean disableImeIcon, int displayId) {
-            mHandler.obtainMessage(MSG_UPDATE_IME_WINDOW_STATUS, disableImeIcon ? 1 : 0, 0)
-                    .sendToTarget();
+            mHandler.post(() -> {
+                synchronized (ImfLock.class) {
+                    updateImeWindowStatusLocked(disableImeIcon, displayId);
+                }
+            });
         }
 
         @Override
@@ -6043,8 +6051,8 @@
         public void onSwitchKeyboardLayoutShortcut(int direction, int displayId,
                 IBinder targetWindowToken) {
             synchronized (ImfLock.class) {
-                // TODO(b/305849394): Infer userId from displayId
-                switchKeyboardLayoutLocked(direction, getUserData(mCurrentUserId));
+                final int userId = resolveImeUserIdFromDisplayIdLocked(displayId);
+                switchKeyboardLayoutLocked(direction, getUserData(userId));
             }
         }
     }
@@ -6278,7 +6286,7 @@
                     };
             mUserDataRepository.forAllUserData(userDataDump);
 
-            if (Flags.imeSwitcherRevamp()) {
+            if (mNewInputMethodSwitcherMenuEnabled) {
                 p.println("  menuControllerNew:");
                 mMenuControllerNew.dump(p, "  ");
             } else {
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java b/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
index 1b84036..4f5af63 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodSettingsRepository.java
@@ -19,15 +19,13 @@
 import android.annotation.AnyThread;
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
-import android.util.SparseArray;
-
-import com.android.internal.annotations.GuardedBy;
 
 final class InputMethodSettingsRepository {
-    // TODO(b/352594784): Should we user other lock primitives?
-    @GuardedBy("sPerUserMap")
+    private static final Object sMutationLock = new Object();
+
     @NonNull
-    private static final SparseArray<InputMethodSettings> sPerUserMap = new SparseArray<>();
+    private static volatile ImmutableSparseArray<InputMethodSettings> sPerUserMap =
+            ImmutableSparseArray.empty();
 
     /**
      * Not intended to be instantiated.
@@ -38,10 +36,7 @@
     @NonNull
     @AnyThread
     static InputMethodSettings get(@UserIdInt int userId) {
-        final InputMethodSettings obj;
-        synchronized (sPerUserMap) {
-            obj = sPerUserMap.get(userId);
-        }
+        final InputMethodSettings obj = sPerUserMap.get(userId);
         if (obj != null) {
             return obj;
         }
@@ -50,15 +45,15 @@
 
     @AnyThread
     static void put(@UserIdInt int userId, @NonNull InputMethodSettings obj) {
-        synchronized (sPerUserMap) {
-            sPerUserMap.put(userId, obj);
+        synchronized (sMutationLock) {
+            sPerUserMap = sPerUserMap.cloneWithPutOrSelf(userId, obj);
         }
     }
 
     @AnyThread
     static void remove(@UserIdInt int userId) {
-        synchronized (sPerUserMap) {
-            sPerUserMap.remove(userId);
+        synchronized (sMutationLock) {
+            sPerUserMap = sPerUserMap.cloneWithRemoveOrSelf(userId);
         }
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java b/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java
index e7cff20..b3500be 100644
--- a/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java
+++ b/services/core/java/com/android/server/inputmethod/SecureSettingsWrapper.java
@@ -24,7 +24,6 @@
 import android.provider.Settings;
 import android.util.ArrayMap;
 import android.util.ArraySet;
-import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.LocalServices;
@@ -38,6 +37,13 @@
  * to the persistent value when the user storage is unlocked.</p>
  */
 final class SecureSettingsWrapper {
+
+    private static final Object sMutationLock = new Object();
+
+    @NonNull
+    private static volatile ImmutableSparseArray<ReaderWriter> sUserMap =
+            ImmutableSparseArray.empty();
+
     @Nullable
     private static volatile ContentResolver sContentResolver = null;
 
@@ -61,8 +67,8 @@
      */
     @AnyThread
     static void endTestMode() {
-        synchronized (sUserMap) {
-            sUserMap.clear();
+        synchronized (sMutationLock) {
+            sUserMap = ImmutableSparseArray.empty();
         }
         sTestMode = false;
     }
@@ -243,10 +249,6 @@
         }
     }
 
-    @GuardedBy("sUserMap")
-    @NonNull
-    private static final SparseArray<ReaderWriter> sUserMap = new SparseArray<>();
-
     private static final ReaderWriter NOOP = new ReaderWriter() {
         @Override
         public void putString(String key, String str) {
@@ -282,15 +284,15 @@
     private static ReaderWriter putOrGet(@UserIdInt int userId,
             @NonNull ReaderWriter readerWriter) {
         final boolean isUnlockedUserImpl = readerWriter instanceof UnlockedUserImpl;
-        synchronized (sUserMap) {
+        synchronized (sMutationLock) {
             final ReaderWriter current = sUserMap.get(userId);
             if (current == null) {
-                sUserMap.put(userId, readerWriter);
+                sUserMap = sUserMap.cloneWithPutOrSelf(userId, readerWriter);
                 return readerWriter;
             }
             // Upgrading from CopyOnWriteImpl to DirectImpl is allowed.
             if (current instanceof LockedUserImpl && isUnlockedUserImpl) {
-                sUserMap.put(userId, readerWriter);
+                sUserMap = sUserMap.cloneWithPutOrSelf(userId, readerWriter);
                 return readerWriter;
             }
             return current;
@@ -300,11 +302,9 @@
     @NonNull
     @AnyThread
     private static ReaderWriter get(@UserIdInt int userId) {
-        synchronized (sUserMap) {
-            final ReaderWriter readerWriter = sUserMap.get(userId);
-            if (readerWriter != null) {
-                return readerWriter;
-            }
+        final ReaderWriter readerWriter = sUserMap.get(userId);
+        if (readerWriter != null) {
+            return readerWriter;
         }
         if (sTestMode) {
             return putOrGet(userId, new FakeReaderWriterImpl());
@@ -363,8 +363,8 @@
      */
     @AnyThread
     static void onUserRemoved(@UserIdInt int userId) {
-        synchronized (sUserMap) {
-            sUserMap.remove(userId);
+        synchronized (sMutationLock) {
+            sUserMap = sUserMap.cloneWithRemoveOrSelf(userId);
         }
     }
 
diff --git a/services/core/java/com/android/server/inputmethod/UserDataRepository.java b/services/core/java/com/android/server/inputmethod/UserDataRepository.java
index 6f831cc..e3524b1 100644
--- a/services/core/java/com/android/server/inputmethod/UserDataRepository.java
+++ b/services/core/java/com/android/server/inputmethod/UserDataRepository.java
@@ -19,51 +19,39 @@
 import android.annotation.AnyThread;
 import android.annotation.NonNull;
 import android.annotation.UserIdInt;
-import android.util.SparseArray;
 
-import com.android.internal.annotations.GuardedBy;
-
-import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Consumer;
 import java.util.function.IntFunction;
 
 final class UserDataRepository {
 
-    private final ReentrantReadWriteLock mUserDataLock = new ReentrantReadWriteLock();
+    private final Object mMutationLock = new Object();
 
-    @GuardedBy("mUserDataLock")
-    private final SparseArray<UserData> mUserData = new SparseArray<>();
+    @NonNull
+    private volatile ImmutableSparseArray<UserData> mUserData = ImmutableSparseArray.empty();
 
     private final IntFunction<InputMethodBindingController> mBindingControllerFactory;
 
     @AnyThread
     @NonNull
     UserData getOrCreate(@UserIdInt int userId) {
-        mUserDataLock.writeLock().lock();
-        try {
-            UserData userData = mUserData.get(userId);
-            if (userData == null) {
-                userData = new UserData(userId, mBindingControllerFactory.apply(userId));
-                mUserData.put(userId, userData);
-            }
+        // Do optimistic read first for optimization.
+        final var userData = mUserData.get(userId);
+        if (userData != null) {
             return userData;
-        } finally {
-            mUserDataLock.writeLock().unlock();
+        }
+        // Note that the below line can be called concurrently. Here we assume that
+        // instantiating UserData for the same user multiple times would have no side effect.
+        final var newUserData = new UserData(userId, mBindingControllerFactory.apply(userId));
+        synchronized (mMutationLock) {
+            mUserData = mUserData.cloneWithPutOrSelf(userId, newUserData);
+            return newUserData;
         }
     }
 
     @AnyThread
     void forAllUserData(Consumer<UserData> consumer) {
-        final SparseArray<UserData> copiedArray;
-        mUserDataLock.readLock().lock();
-        try {
-            copiedArray = mUserData.clone();
-        } finally {
-            mUserDataLock.readLock().unlock();
-        }
-        for (int i = 0; i < copiedArray.size(); i++) {
-            consumer.accept(copiedArray.valueAt(i));
-        }
+        mUserData.forEach(consumer);
     }
 
     UserDataRepository(
@@ -73,11 +61,8 @@
 
     @AnyThread
     void remove(@UserIdInt int userId) {
-        mUserDataLock.writeLock().lock();
-        try {
-            mUserData.remove(userId);
-        } finally {
-            mUserDataLock.writeLock().unlock();
+        synchronized (mMutationLock) {
+            mUserData = mUserData.cloneWithRemoveOrSelf(userId);
         }
     }
 }
diff --git a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
index 770e12d..f603ff3 100644
--- a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
+++ b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
@@ -86,8 +86,6 @@
     interface Callback extends IInputMethodManagerImpl.Callback {
         @GuardedBy("ImfLock.class")
         ClientState getClientStateLocked(IInputMethodClient client);
-        @GuardedBy("ImfLock.class")
-        boolean isInputShownLocked();
     }
 
     private final Callback mInner;
@@ -178,19 +176,8 @@
             @Nullable ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
             @MotionEvent.ToolType int lastClickToolType, ResultReceiver resultReceiver,
             @SoftInputShowHideReason int reason) {
-        offload(
-                () -> {
-                    if (!mInner.showSoftInput(
-                            client,
-                            windowToken,
-                            statsToken,
-                            flags,
-                            lastClickToolType,
-                            resultReceiver,
-                            reason)) {
-                        sendResultReceiverFailure(resultReceiver);
-                    }
-                });
+        offload(() -> mInner.showSoftInput(
+                client, windowToken, statsToken, flags, lastClickToolType, resultReceiver, reason));
         return true;
     }
 
@@ -198,30 +185,11 @@
     public boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
             @Nullable ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
             ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
-        offload(
-                () -> {
-                    if (!mInner.hideSoftInput(
-                            client, windowToken, statsToken, flags, resultReceiver, reason)) {
-                        sendResultReceiverFailure(resultReceiver);
-                    }
-                });
+        offload(() -> mInner.hideSoftInput(
+                client, windowToken, statsToken, flags, resultReceiver, reason));
         return true;
     }
 
-    private void sendResultReceiverFailure(@Nullable ResultReceiver resultReceiver) {
-        if (resultReceiver == null) {
-            return;
-        }
-        final boolean isInputShown;
-        synchronized (ImfLock.class) {
-            isInputShown = mInner.isInputShownLocked();
-        }
-        resultReceiver.send(isInputShown
-                        ? InputMethodManager.RESULT_UNCHANGED_SHOWN
-                        : InputMethodManager.RESULT_UNCHANGED_HIDDEN,
-                null);
-    }
-
     @Override
     @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
     public void hideSoftInputFromServerForTest() {
@@ -332,8 +300,8 @@
 
     @IInputMethodManagerImpl.PermissionVerified(Manifest.permission.INTERNAL_SYSTEM_WINDOW)
     @Override
-    public void removeImeSurface() {
-        mInner.removeImeSurface();
+    public void removeImeSurface(int displayId) {
+        mInner.removeImeSurface(displayId);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 016abff..4179edd 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -1315,10 +1315,24 @@
                         nv.rank, nv.count);
 
                 StatusBarNotification sbn = r.getSbn();
-                cancelNotification(callingUid, callingPid, sbn.getPackageName(), sbn.getTag(),
-                        sbn.getId(), FLAG_AUTO_CANCEL,
-                        FLAG_FOREGROUND_SERVICE | FLAG_USER_INITIATED_JOB | FLAG_BUBBLE,
-                        false, r.getUserId(), REASON_CLICK, nv.rank, nv.count, null);
+                // Notifications should be cancelled on click if they have been lifetime extended,
+                // regardless of presence or absence of FLAG_AUTO_CANCEL.
+                if (lifetimeExtensionRefactor()
+                        && (sbn.getNotification().flags
+                        & FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY) != 0) {
+                    cancelNotification(callingUid, callingPid, sbn.getPackageName(), sbn.getTag(),
+                            sbn.getId(), FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY,
+                            FLAG_FOREGROUND_SERVICE | FLAG_USER_INITIATED_JOB
+                                    | FLAG_BUBBLE,
+                            false, r.getUserId(), REASON_CLICK, nv.rank, nv.count, null);
+
+                } else {
+                    // Otherwise, only FLAG_AUTO_CANCEL notifications should be canceled on click.
+                    cancelNotification(callingUid, callingPid, sbn.getPackageName(), sbn.getTag(),
+                            sbn.getId(), FLAG_AUTO_CANCEL,
+                            FLAG_FOREGROUND_SERVICE | FLAG_USER_INITIATED_JOB | FLAG_BUBBLE,
+                            false, r.getUserId(), REASON_CLICK, nv.rank, nv.count, null);
+                }
                 nv.recycle();
                 reportUserInteraction(r);
                 mAssistants.notifyAssistantNotificationClicked(r);
diff --git a/services/core/java/com/android/server/notification/ZenModeHelper.java b/services/core/java/com/android/server/notification/ZenModeHelper.java
index b12a917..95d8bb9 100644
--- a/services/core/java/com/android/server/notification/ZenModeHelper.java
+++ b/services/core/java/com/android/server/notification/ZenModeHelper.java
@@ -2811,8 +2811,9 @@
     private final class H extends Handler {
         private static final int MSG_DISPATCH = 1;
         private static final int MSG_METRICS = 2;
-        private static final int MSG_RINGER_AUDIO = 5;
         private static final int MSG_APPLY_EFFECTS = 6;
+        private static final int MSG_AUDIO_APPLIED_TO_RINGER = 7;
+        private static final int MSG_AUDIO_NOT_APPLIED_TO_RINGER = 8;
 
         private static final long METRICS_PERIOD_MS = 6 * 60 * 60 * 1000;
 
@@ -2831,8 +2832,13 @@
         }
 
         private void postUpdateRingerAndAudio(boolean shouldApplyToRinger) {
-            removeMessages(MSG_RINGER_AUDIO);
-            sendMessage(obtainMessage(MSG_RINGER_AUDIO, shouldApplyToRinger));
+            if (shouldApplyToRinger) {
+                removeMessages(MSG_AUDIO_APPLIED_TO_RINGER);
+                sendEmptyMessage(MSG_AUDIO_APPLIED_TO_RINGER);
+            } else {
+                removeMessages(MSG_AUDIO_NOT_APPLIED_TO_RINGER);
+                sendEmptyMessage(MSG_AUDIO_NOT_APPLIED_TO_RINGER);
+            }
         }
 
         private void postApplyDeviceEffects(@ConfigChangeOrigin int origin) {
@@ -2849,9 +2855,11 @@
                 case MSG_METRICS:
                     mMetrics.emit();
                     break;
-                case MSG_RINGER_AUDIO:
-                    boolean shouldApplyToRinger = (boolean) msg.obj;
-                    updateRingerAndAudio(shouldApplyToRinger);
+                case MSG_AUDIO_APPLIED_TO_RINGER:
+                    updateRingerAndAudio(/* shouldApplyToRinger= */ true);
+                    break;
+                case MSG_AUDIO_NOT_APPLIED_TO_RINGER:
+                    updateRingerAndAudio(/* shouldApplyToRinger= */ false);
                     break;
                 case MSG_APPLY_EFFECTS:
                     @ConfigChangeOrigin int origin = msg.arg1;
diff --git a/services/core/java/com/android/server/pm/InstallPackageHelper.java b/services/core/java/com/android/server/pm/InstallPackageHelper.java
index 8d3f07e..4425079 100644
--- a/services/core/java/com/android/server/pm/InstallPackageHelper.java
+++ b/services/core/java/com/android/server/pm/InstallPackageHelper.java
@@ -3264,6 +3264,7 @@
     /**
      * Tries to restore the disabled system package after an update has been deleted.
      */
+    @GuardedBy("mPm.mInstallLock")
     public void restoreDisabledSystemPackageLIF(DeletePackageAction action,
             @NonNull int[] allUserHandles, boolean writeSettings) throws SystemDeleteException {
         final PackageSetting deletedPs = action.mDeletingPs;
@@ -3282,10 +3283,21 @@
         }
         // Install the system package
         if (DEBUG_REMOVE) Slog.d(TAG, "Re-installing system package: " + disabledPs);
-        try (PackageManagerTracedLock installLock = mPm.mInstallLock.acquireLock()) {
+        try {
             final int[] origUsers = outInfo == null ? null : outInfo.mOrigUsers;
-            installPackageFromSystemLIF(disabledPs.getPathString(), allUserHandles,
-                    origUsers, writeSettings);
+            try (PackageManagerTracedLock installLock = mPm.mInstallLock.acquireLock()) {
+                installPackageFromSystemLIF(disabledPs.getPathString(), allUserHandles,
+                        origUsers, writeSettings);
+            }
+            if (origUsers != null) {
+                mPm.commitPackageStateMutation(null, mutator -> {
+                    for (int userId : origUsers) {
+                        mutator.forPackage(disabledPs.getPackageName())
+                                .userState(userId)
+                                .setOverlayPaths(deletedPs.getOverlayPaths(userId));
+                    }
+                });
+            }
         } catch (PackageManagerException e) {
             Slog.w(TAG, "Failed to restore system package:" + deletedPs.getPackageName() + ": "
                     + e.getMessage());
@@ -3813,13 +3825,13 @@
                 // This also has the (beneficial) side effect where if a package disappears from an
                 // APEX, leaving only a /data copy, it will lose its apexModuleName.
                 //
-                // This must be done before scanSystemPackageLI as that will throw in the case of a
+                // This must be done before scanPackageForInitLI as that will throw in the case of a
                 // system -> data package.
                 disabledPkgSetting.setApexModuleName(activeApexInfo.apexModuleName);
             }
         }
 
-        final Pair<ScanResult, Boolean> scanResultPair = scanSystemPackageLI(
+        final Pair<ScanResult, Boolean> scanResultPair = scanPackageForInitLI(
                 parsedPackage, parseFlags, scanFlags, user);
         final ScanResult scanResult = scanResultPair.first;
         boolean shouldHideSystemApp = scanResultPair.second;
@@ -4054,7 +4066,7 @@
         }
     }
 
-    private Pair<ScanResult, Boolean> scanSystemPackageLI(ParsedPackage parsedPackage,
+    private Pair<ScanResult, Boolean> scanPackageForInitLI(ParsedPackage parsedPackage,
             @ParsingPackageUtils.ParseFlags int parseFlags,
             @PackageManagerService.ScanFlags int scanFlags,
             @Nullable UserHandle user) throws PackageManagerException {
@@ -4167,7 +4179,7 @@
                         ParsingPackageUtils.getSigningDetails(input, parsedPackage,
                                 false /*skipVerify*/);
                 if (result.isError()) {
-                    throw new PrepareFailure("Failed collect during scanSystemPackageLI",
+                    throw new PrepareFailure("Failed collect during scanPackageForInitLI",
                             result.getException());
                 }
                 disabledPkgSetting.setSigningDetails(result.getResult());
diff --git a/services/core/java/com/android/server/pm/PackageManagerServiceUtils.java b/services/core/java/com/android/server/pm/PackageManagerServiceUtils.java
index a1dffc6..9757582 100644
--- a/services/core/java/com/android/server/pm/PackageManagerServiceUtils.java
+++ b/services/core/java/com/android/server/pm/PackageManagerServiceUtils.java
@@ -1452,6 +1452,13 @@
             }
 
             if (!ArrayUtils.isEmpty(after.splitNames)) {
+                if (beforeSplitNames.length != beforeSplitRevisionCodes.length) {
+                    throw new PackageManagerException(INSTALL_FAILED_VERSION_DOWNGRADE,
+                            "Current split names and the split revision codes are not 1:1 mapping."
+                                    + "This indicates that the package info data has been"
+                                    + " corrupted.");
+                }
+
                 for (int i = 0; i < after.splitNames.length; i++) {
                     final String splitName = after.splitNames[i];
                     final int j = ArrayUtils.indexOf(beforeSplitNames, splitName);
diff --git a/services/core/java/com/android/server/pm/Settings.java b/services/core/java/com/android/server/pm/Settings.java
index 9177e2b..b7dfd8d 100644
--- a/services/core/java/com/android/server/pm/Settings.java
+++ b/services/core/java/com/android/server/pm/Settings.java
@@ -4489,10 +4489,24 @@
         String splitName = parser.getAttributeValue(null, ATTR_NAME);
         int splitRevision = parser.getAttributeInt(null, ATTR_VERSION, -1);
         if (splitName != null && splitRevision >= 0) {
+            final int beforeSplitNamesLength = outPs.getSplitNames().length;
+            // If the split name already exists in the outPs#getSplitNames, don't add it
+            // into the array and update its revision code below
             outPs.setSplitNames(ArrayUtils.appendElement(String.class,
                     outPs.getSplitNames(), splitName));
-            outPs.setSplitRevisionCodes(ArrayUtils.appendInt(
-                    outPs.getSplitRevisionCodes(), splitRevision));
+
+            // If the same split name has already been added before, update the latest
+            // revision code
+            final int afterSplitNamesLength = outPs.getSplitNames().length;
+            if (beforeSplitNamesLength == afterSplitNamesLength) {
+                final int index = ArrayUtils.indexOf(outPs.getSplitNames(), splitName);
+                final int[] splitRevisionCodes = outPs.getSplitRevisionCodes();
+                splitRevisionCodes[index] = splitRevision;
+                outPs.setSplitRevisionCodes(splitRevisionCodes);
+            } else {
+                outPs.setSplitRevisionCodes(ArrayUtils.appendInt(
+                        outPs.getSplitRevisionCodes(), splitRevision, /* allowDuplicates= */ true));
+            }
         }
 
         XmlUtils.skipCurrentTag(parser);
diff --git a/services/core/java/com/android/server/power/hint/TEST_MAPPING b/services/core/java/com/android/server/power/hint/TEST_MAPPING
index 34c25c6..9988786 100644
--- a/services/core/java/com/android/server/power/hint/TEST_MAPPING
+++ b/services/core/java/com/android/server/power/hint/TEST_MAPPING
@@ -1,27 +1,18 @@
 {
-  "presubmit": [
-    {
-      "name": "FrameworksServicesTests",
-      "options": [
-        {
-          "include-filter": "com.android.server.power.hint"
-        },
-        {
-          "exclude-annotation": "androidx.test.filters.FlakyTest"
-        }
-      ]
-    }
-  ],
   "postsubmit": [
     {
+      "name": "PerformanceHintTests",
+      "options": [
+        {"exclude-annotation": "androidx.test.filters.FlakyTest"},
+        {"exclude-annotation": "org.junit.Ignore"}
+      ]
+    },
+    {
       "name": "CtsStatsdAtomHostTestCases",
       "options": [
+        {"include-filter": "android.cts.statsdatom.performancehintmanager"},
         {"exclude-annotation": "androidx.test.filters.FlakyTest"},
-        {"exclude-annotation": "org.junit.Ignore"},
-        {"include-filter": "android.cts.statsdatom.performancehintmanager"}
-      ],
-      "file_patterns": [
-        "(/|^)HintManagerService.java"
+        {"exclude-annotation": "org.junit.Ignore"}
       ]
     }
   ]
diff --git a/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java
new file mode 100644
index 0000000..a42929f
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.os.BatteryConsumer;
+import android.os.PersistableBundle;
+
+import com.android.internal.os.PowerStats;
+
+public class AmbientDisplayPowerStatsProcessor extends PowerStatsProcessor {
+    private final PowerStatsLayout mStatsLayout;
+    private final PowerStats.Descriptor mDescriptor;
+    private final long[] mTmpDeviceStats;
+    private PowerStats.Descriptor mScreenPowerStatsDescriptor;
+    private ScreenPowerStatsLayout mScreenPowerStatsLayout;
+    private long[] mTmpScreenStats;
+
+    public AmbientDisplayPowerStatsProcessor() {
+        mStatsLayout = new PowerStatsLayout();
+        mStatsLayout.addDeviceSectionPowerEstimate();
+        PersistableBundle extras = new PersistableBundle();
+        mStatsLayout.toExtras(extras);
+        mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                mStatsLayout.getDeviceStatsArrayLength(), null, 0, 0, extras);
+        mTmpDeviceStats = new long[mDescriptor.statsArrayLength];
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats, long timestampMs) {
+        stats.setPowerStatsDescriptor(mDescriptor);
+
+        PowerComponentAggregatedPowerStats screenStats =
+                stats.getAggregatedPowerStats().getPowerComponentStats(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN);
+        if (screenStats == null) {
+            return;
+        }
+
+        if (mScreenPowerStatsDescriptor == null) {
+            mScreenPowerStatsDescriptor = screenStats.getPowerStatsDescriptor();
+            if (mScreenPowerStatsDescriptor == null) {
+                return;
+            }
+
+            mScreenPowerStatsLayout = new ScreenPowerStatsLayout(mScreenPowerStatsDescriptor);
+            mTmpScreenStats = new long[mScreenPowerStatsDescriptor.statsArrayLength];
+        }
+
+        MultiStateStats.States[] deviceStateConfig = screenStats.getConfig().getDeviceStateConfig();
+
+        // Ambient display power estimates have already been calculated by the screen power stats
+        // processor. All that remains to be done is copy the estimates over.
+        MultiStateStats.States.forEachTrackedStateCombination(deviceStateConfig,
+                states -> {
+                    screenStats.getDeviceStats(mTmpScreenStats, states);
+                    double power =
+                            mScreenPowerStatsLayout.getScreenDozePowerEstimate(mTmpScreenStats);
+                    mStatsLayout.setDevicePowerEstimate(mTmpDeviceStats, power);
+                    stats.setDeviceStats(states, mTmpDeviceStats);
+                });
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index 4052a64..143b3ff 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -296,6 +296,7 @@
     private final LongSparseArray<SamplingTimer> mKernelMemoryStats = new LongSparseArray<>();
     private int[] mCpuPowerBracketMap;
     private final CpuPowerStatsCollector mCpuPowerStatsCollector;
+    private final ScreenPowerStatsCollector mScreenPowerStatsCollector;
     private final MobileRadioPowerStatsCollector mMobileRadioPowerStatsCollector;
     private final WifiPowerStatsCollector mWifiPowerStatsCollector;
     private final BluetoothPowerStatsCollector mBluetoothPowerStatsCollector;
@@ -303,6 +304,54 @@
     private final GnssPowerStatsCollector mGnssPowerStatsCollector;
     private final CustomEnergyConsumerPowerStatsCollector mCustomEnergyConsumerPowerStatsCollector;
     private final SparseBooleanArray mPowerStatsCollectorEnabled = new SparseBooleanArray();
+    private ScreenPowerStatsCollector.ScreenUsageTimeRetriever mScreenUsageTimeRetriever =
+            new ScreenPowerStatsCollector.ScreenUsageTimeRetriever() {
+
+                @Override
+                public long getScreenOnTimeMs(int display) {
+                    synchronized (BatteryStatsImpl.this) {
+                        return getDisplayScreenOnTime(display,
+                                mClock.elapsedRealtime() * 1000) / 1000;
+                    }
+                }
+
+                @Override
+                public long getBrightnessLevelTimeMs(int display, int brightnessLevel) {
+                    synchronized (BatteryStatsImpl.this) {
+                        return getDisplayScreenBrightnessTime(display, brightnessLevel,
+                                mClock.elapsedRealtime() * 1000) / 1000;
+                    }
+                }
+
+                @Override
+                public long getScreenDozeTimeMs(int display) {
+                    synchronized (BatteryStatsImpl.this) {
+                        return getDisplayScreenDozeTime(display,
+                                mClock.elapsedRealtime() * 1000) / 1000;
+                    }
+                }
+
+                @Override
+                public void retrieveTopActivityTimes(Callback callback) {
+                    synchronized (BatteryStatsImpl.this) {
+                        long elapsedTimeUs = mClock.elapsedRealtime() * 1000;
+                        for (int i = mUidStats.size() - 1; i >= 0; i--) {
+                            Uid uid = mUidStats.valueAt(i);
+                            long topStateTime = uid.getProcessStateTime(Uid.PROCESS_STATE_TOP,
+                                    elapsedTimeUs, STATS_SINCE_CHARGED) / 1000;
+                            Timer timer = uid.getForegroundActivityTimer();
+                            if (timer == null) {
+                                callback.onUidTopActivityTime(uid.mUid, topStateTime);
+                            } else {
+                                long topActivityTime = timer.getTotalTimeLocked(elapsedTimeUs,
+                                        STATS_SINCE_CHARGED) / 1000;
+                                callback.onUidTopActivityTime(uid.mUid, Math.min(topStateTime,
+                                        topActivityTime));
+                            }
+                        }
+                    }
+                }
+            };
     private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
             new WifiPowerStatsCollector.WifiStatsRetriever() {
                 @Override
@@ -1966,8 +2015,9 @@
     }
 
     private class PowerStatsCollectorInjector implements CpuPowerStatsCollector.Injector,
-            MobileRadioPowerStatsCollector.Injector, WifiPowerStatsCollector.Injector,
-            BluetoothPowerStatsCollector.Injector, EnergyConsumerPowerStatsCollector.Injector {
+            ScreenPowerStatsCollector.Injector, MobileRadioPowerStatsCollector.Injector,
+            WifiPowerStatsCollector.Injector, BluetoothPowerStatsCollector.Injector,
+            EnergyConsumerPowerStatsCollector.Injector {
         private PackageManager mPackageManager;
         private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
         private NetworkStatsManager mNetworkStatsManager;
@@ -2039,6 +2089,16 @@
         }
 
         @Override
+        public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
+            return mScreenUsageTimeRetriever;
+        }
+
+        @Override
+        public int getDisplayCount() {
+            return BatteryStatsImpl.this.getDisplayCount();
+        }
+
+        @Override
         public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
             return () -> readMobileNetworkStatsLocked(mNetworkStatsManager);
         }
@@ -5371,8 +5431,6 @@
         }
     }
 
-    int mSensorNesting;
-
     @GuardedBy("this")
     public void noteStartSensorLocked(int uid, int sensor) {
         noteStartSensorLocked(uid, sensor, mClock.elapsedRealtime(), mClock.uptimeMillis());
@@ -5381,11 +5439,8 @@
     @GuardedBy("this")
     public void noteStartSensorLocked(int uid, int sensor, long elapsedRealtimeMs, long uptimeMs) {
         uid = mapUid(uid);
-        if (mSensorNesting == 0) {
-            mHistory.recordStateStartEvent(elapsedRealtimeMs, uptimeMs,
-                    HistoryItem.STATE_SENSOR_ON_FLAG);
-        }
-        mSensorNesting++;
+        mHistory.recordStateStartEvent(elapsedRealtimeMs, uptimeMs,
+                HistoryItem.STATE_SENSOR_ON_FLAG, uid, "sensor:0x" + Integer.toHexString(sensor));
         getUidStatsLocked(uid, elapsedRealtimeMs, uptimeMs)
                 .noteStartSensor(sensor, elapsedRealtimeMs);
     }
@@ -5398,11 +5453,8 @@
     @GuardedBy("this")
     public void noteStopSensorLocked(int uid, int sensor, long elapsedRealtimeMs, long uptimeMs) {
         uid = mapUid(uid);
-        mSensorNesting--;
-        if (mSensorNesting == 0) {
-            mHistory.recordStateStopEvent(elapsedRealtimeMs, uptimeMs,
-                    HistoryItem.STATE_SENSOR_ON_FLAG);
-        }
+        mHistory.recordStateStopEvent(elapsedRealtimeMs, uptimeMs,
+                HistoryItem.STATE_SENSOR_ON_FLAG, uid, "sensor:0x" + Integer.toHexString(sensor));
         getUidStatsLocked(uid, elapsedRealtimeMs, uptimeMs)
                 .noteStopSensor(sensor, elapsedRealtimeMs);
     }
@@ -5736,13 +5788,17 @@
         maybeUpdateOverallScreenBrightness(overallBin, elapsedRealtimeMs, uptimeMs);
 
         if (shouldScheduleSync) {
-            final int numDisplays = mPerDisplayBatteryStats.length;
-            final int[] displayStates = new int[numDisplays];
-            for (int i = 0; i < numDisplays; i++) {
-                displayStates[i] = mPerDisplayBatteryStats[i].screenState;
+            if (mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_SCREEN)) {
+                mScreenPowerStatsCollector.schedule();
+            } else {
+                final int numDisplays = mPerDisplayBatteryStats.length;
+                final int[] displayStates = new int[numDisplays];
+                for (int i = 0; i < numDisplays; i++) {
+                    displayStates[i] = mPerDisplayBatteryStats[i].screenState;
+                }
+                mExternalSync.scheduleSyncDueToScreenStateChange(externalUpdateFlag,
+                        batteryRunning, batteryScreenOffRunning, state, displayStates);
             }
-            mExternalSync.scheduleSyncDueToScreenStateChange(externalUpdateFlag,
-                    batteryRunning, batteryScreenOffRunning, state, displayStates);
         }
     }
 
@@ -11290,6 +11346,9 @@
         mCpuPowerStatsCollector = new CpuPowerStatsCollector(mPowerStatsCollectorInjector);
         mCpuPowerStatsCollector.addConsumer(this::recordPowerStats);
 
+        mScreenPowerStatsCollector = new ScreenPowerStatsCollector(mPowerStatsCollectorInjector);
+        mScreenPowerStatsCollector.addConsumer(this::recordPowerStats);
+
         mMobileRadioPowerStatsCollector = new MobileRadioPowerStatsCollector(
                 mPowerStatsCollectorInjector, this::onMobileRadioPowerStatsRetrieved);
         mMobileRadioPowerStatsCollector.addConsumer(this::recordPowerStats);
@@ -14750,6 +14809,10 @@
                 mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU));
         mCpuPowerStatsCollector.schedule();
 
+        mScreenPowerStatsCollector.setEnabled(
+                mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_SCREEN));
+        mScreenPowerStatsCollector.schedule();
+
         mMobileRadioPowerStatsCollector.setEnabled(
                 mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO));
         mMobileRadioPowerStatsCollector.schedule();
@@ -14786,6 +14849,8 @@
         switch (powerComponent) {
             case BatteryConsumer.POWER_COMPONENT_CPU:
                 return mCpuPowerStatsCollector;
+            case BatteryConsumer.POWER_COMPONENT_SCREEN:
+                return mScreenPowerStatsCollector;
             case BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO:
                 return mMobileRadioPowerStatsCollector;
             case BatteryConsumer.POWER_COMPONENT_WIFI:
@@ -16329,6 +16394,7 @@
      */
     public void schedulePowerStatsSampleCollection() {
         mCpuPowerStatsCollector.forceSchedule();
+        mScreenPowerStatsCollector.forceSchedule();
         mMobileRadioPowerStatsCollector.forceSchedule();
         mWifiPowerStatsCollector.forceSchedule();
         mBluetoothPowerStatsCollector.forceSchedule();
@@ -16351,6 +16417,7 @@
      */
     public void dumpStatsSample(PrintWriter pw) {
         mCpuPowerStatsCollector.collectAndDump(pw);
+        mScreenPowerStatsCollector.collectAndDump(pw);
         mMobileRadioPowerStatsCollector.collectAndDump(pw);
         mWifiPowerStatsCollector.collectAndDump(pw);
         mBluetoothPowerStatsCollector.collectAndDump(pw);
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index ac68966..b308f38 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -93,8 +93,10 @@
                 if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_BLUETOOTH)) {
                     mPowerCalculators.add(new BluetoothPowerCalculator(mPowerProfile));
                 }
-                mPowerCalculators.add(new SensorPowerCalculator(
-                        mContext.getSystemService(SensorManager.class)));
+                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_SENSORS)) {
+                    mPowerCalculators.add(new SensorPowerCalculator(
+                            mContext.getSystemService(SensorManager.class)));
+                }
                 if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_GNSS)) {
                     mPowerCalculators.add(new GnssPowerCalculator(mPowerProfile));
                 }
@@ -110,8 +112,13 @@
                 if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_VIDEO)) {
                     mPowerCalculators.add(new VideoPowerCalculator(mPowerProfile));
                 }
-                mPowerCalculators.add(new ScreenPowerCalculator(mPowerProfile));
-                mPowerCalculators.add(new AmbientDisplayPowerCalculator(mPowerProfile));
+                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_SCREEN)) {
+                    mPowerCalculators.add(new ScreenPowerCalculator(mPowerProfile));
+                }
+                if (!mPowerStatsExporterEnabled.get(
+                        BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY)) {
+                    mPowerCalculators.add(new AmbientDisplayPowerCalculator(mPowerProfile));
+                }
                 mPowerCalculators.add(new IdlePowerCalculator(mPowerProfile));
                 if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_ANY)) {
                     mPowerCalculators.add(new CustomEnergyConsumerPowerCalculator(mPowerProfile));
diff --git a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
index 8384b2b..6820197 100644
--- a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
@@ -163,15 +163,14 @@
         }
     }
 
-    void setDeviceStats(@AggregatedPowerStatsConfig.TrackedState int[] states, long[] values) {
+    void setDeviceStats(int[] states, long[] values) {
         if (mDeviceStats == null) {
             createDeviceStats(0);
         }
         mDeviceStats.setStats(states, values);
     }
 
-    void setUidStats(int uid, @AggregatedPowerStatsConfig.TrackedState int[] states,
-            long[] values) {
+    void setUidStats(int uid, int[] states, long[] values) {
         UidStats uidStats = getUidStats(uid);
         uidStats.stats.setStats(states, values);
     }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
index dfc8daa..c81c7ff 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
@@ -220,8 +220,7 @@
         }
 
         @Nullable
-        public DeviceStateEstimation getDeviceStateEstimate(
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        public DeviceStateEstimation getDeviceStateEstimate(int[] stateValues) {
             String label = concatLabels(mConfig.getDeviceStateConfig(), stateValues);
             for (int i = 0; i < deviceStateEstimations.size(); i++) {
                 DeviceStateEstimation deviceStateEstimation = this.deviceStateEstimations.get(i);
@@ -233,8 +232,7 @@
         }
 
         public CombinedDeviceStateEstimate getCombinedDeviceStateEstimate(
-                MultiStateStats.States[] deviceStates,
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+                MultiStateStats.States[] deviceStates, int[] stateValues) {
             String label = concatLabels(deviceStates, stateValues);
             for (int i = 0; i < combinedDeviceStateEstimations.size(); i++) {
                 CombinedDeviceStateEstimate cdse = combinedDeviceStateEstimations.get(i);
@@ -259,8 +257,8 @@
             for (int i = deviceStateEstimations.size() - 1; i >= 0; i--) {
                 deviceStateEstimations.get(i).intermediates = null;
             }
-            for (int i = deviceStateEstimations.size() - 1; i >= 0; i--) {
-                deviceStateEstimations.get(i).intermediates = null;
+            for (int i = combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
+                combinedDeviceStateEstimations.get(i).intermediates = null;
             }
             for (int i = uidStateEstimates.size() - 1; i >= 0; i--) {
                 UidStateEstimate uidStateEstimate = uidStateEstimates.get(i);
@@ -275,12 +273,10 @@
 
     protected static class DeviceStateEstimation {
         public final String id;
-        @AggregatedPowerStatsConfig.TrackedState
         public final int[] stateValues;
         public Object intermediates;
 
-        public DeviceStateEstimation(MultiStateStats.States[] config,
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        public DeviceStateEstimation(MultiStateStats.States[] config, int[] stateValues) {
             id = concatLabels(config, stateValues);
             this.stateValues = stateValues;
         }
@@ -288,11 +284,12 @@
 
     protected static class CombinedDeviceStateEstimate {
         public final String id;
+        public final int[] stateValues;
         public List<DeviceStateEstimation> deviceStateEstimations = new ArrayList<>();
         public Object intermediates;
 
-        public CombinedDeviceStateEstimate(MultiStateStats.States[] config,
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        public CombinedDeviceStateEstimate(MultiStateStats.States[] config, int[] stateValues) {
+            this.stateValues = Arrays.copyOf(stateValues, stateValues.length);
             id = concatLabels(config, stateValues);
         }
     }
@@ -310,19 +307,16 @@
     }
 
     protected static class UidStateProportionalEstimate {
-        @AggregatedPowerStatsConfig.TrackedState
         public final int[] stateValues;
         public Object intermediates;
 
-        protected UidStateProportionalEstimate(
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        protected UidStateProportionalEstimate(int[] stateValues) {
             this.stateValues = stateValues;
         }
     }
 
     @NonNull
-    private static String concatLabels(MultiStateStats.States[] config,
-            @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+    private static String concatLabels(MultiStateStats.States[] config, int[] stateValues) {
         List<String> labels = new ArrayList<>();
         for (int state = 0; state < config.length; state++) {
             if (config[state] != null && config[state].isTracked()) {
@@ -334,7 +328,6 @@
         return labels.toString();
     }
 
-    @AggregatedPowerStatsConfig.TrackedState
     private static int[][] getAllTrackedStateCombinations(MultiStateStats.States[] states) {
         List<int[]> combinations = new ArrayList<>();
         MultiStateStats.States.forEachTrackedStateCombination(states, stateValues -> {
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
new file mode 100644
index 0000000..291f289
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
@@ -0,0 +1,227 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.hardware.power.stats.EnergyConsumerType;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.Handler;
+import android.os.PersistableBundle;
+import android.util.Slog;
+import android.util.SparseLongArray;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+import java.util.Arrays;
+import java.util.function.IntSupplier;
+
+public class ScreenPowerStatsCollector extends PowerStatsCollector {
+    private static final String TAG = "ScreenPowerStatsCollector";
+
+    interface ScreenUsageTimeRetriever {
+        interface Callback {
+            void onUidTopActivityTime(int uid, long topActivityTimeMs);
+        }
+
+        void retrieveTopActivityTimes(Callback callback);
+
+        long getScreenOnTimeMs(int display);
+        long getBrightnessLevelTimeMs(int display, int brightnessLevel);
+        long getScreenDozeTimeMs(int display);
+    }
+
+    interface Injector {
+        Handler getHandler();
+        Clock getClock();
+        PowerStatsUidResolver getUidResolver();
+        long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
+        ConsumedEnergyRetriever getConsumedEnergyRetriever();
+        IntSupplier getVoltageSupplier();
+        ScreenUsageTimeRetriever getScreenUsageTimeRetriever();
+        int getDisplayCount();
+    }
+
+    private static final long ENERGY_UNSPECIFIED = -1;
+
+    private final Injector mInjector;
+    private boolean mIsInitialized;
+    private ScreenPowerStatsLayout mLayout;
+    private int mDisplayCount;
+    private PowerStats mPowerStats;
+    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    private IntSupplier mVoltageSupplier;
+    private ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
+    private int[] mEnergyConsumerIds = new int[0];
+    private long[] mLastConsumedEnergyUws;
+    private int mLastVoltageMv;
+    private boolean mFirstSample = true;
+    private long[] mLastScreenOnTime;
+    private long[][] mLastBrightnessLevelTime;
+    private long[] mLastDozeTime;
+    private final SparseLongArray mLastTopActivityTime = new SparseLongArray();
+    private long mLastCollectionTime;
+
+    ScreenPowerStatsCollector(Injector injector) {
+        super(injector.getHandler(),
+                injector.getPowerStatsCollectionThrottlePeriod(
+                        BatteryConsumer.powerComponentIdToString(
+                                BatteryConsumer.POWER_COMPONENT_SCREEN)),
+                injector.getUidResolver(), injector.getClock());
+        mInjector = injector;
+    }
+
+    private boolean ensureInitialized() {
+        if (mIsInitialized) {
+            return true;
+        }
+
+        if (!isEnabled()) {
+            return false;
+        }
+
+        mDisplayCount = mInjector.getDisplayCount();
+        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
+        mVoltageSupplier = mInjector.getVoltageSupplier();
+        mScreenUsageTimeRetriever = mInjector.getScreenUsageTimeRetriever();
+        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
+                EnergyConsumerType.DISPLAY);
+        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
+        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+
+        mLayout = new ScreenPowerStatsLayout();
+        mLayout.addDeviceScreenUsageDurationSection(mInjector.getDisplayCount());
+        mLayout.addDeviceSectionEnergyConsumers(mEnergyConsumerIds.length);
+        mLayout.addDeviceSectionUsageDuration();
+        mLayout.addDeviceSectionPowerEstimate();
+        mLayout.addUidTopActivitiyDuration();
+        mLayout.addUidSectionPowerEstimate();
+
+        PersistableBundle extras = new PersistableBundle();
+        mLayout.toExtras(extras);
+        PowerStats.Descriptor powerStatsDescriptor = new PowerStats.Descriptor(
+                BatteryConsumer.POWER_COMPONENT_SCREEN, mLayout.getDeviceStatsArrayLength(),
+                null, 0, mLayout.getUidStatsArrayLength(),
+                extras);
+
+        mLastScreenOnTime = new long[mDisplayCount];
+        mLastBrightnessLevelTime = new long[mDisplayCount][BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS];
+        mLastDozeTime = new long[mDisplayCount];
+
+        mPowerStats = new PowerStats(powerStatsDescriptor);
+
+        mIsInitialized = true;
+        return true;
+    }
+
+    @Override
+    protected PowerStats collectStats() {
+        if (!ensureInitialized()) {
+            return null;
+        }
+
+        if (mEnergyConsumerIds.length != 0) {
+            collectEnergyConsumers();
+        }
+
+        for (int display = 0; display < mDisplayCount; display++) {
+            long screenOnTimeMs = mScreenUsageTimeRetriever.getScreenOnTimeMs(display);
+            if (!mFirstSample) {
+                mLayout.setScreenOnDuration(mPowerStats.stats, display,
+                        screenOnTimeMs - mLastScreenOnTime[display]);
+            }
+            mLastScreenOnTime[display] = screenOnTimeMs;
+
+            for (int level = 0; level < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; level++) {
+                long brightnessLevelTimeMs =
+                        mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(display, level);
+                if (!mFirstSample) {
+                    mLayout.setBrightnessLevelDuration(mPowerStats.stats, display, level,
+                            brightnessLevelTimeMs - mLastBrightnessLevelTime[display][level]);
+                }
+                mLastBrightnessLevelTime[display][level] = brightnessLevelTimeMs;
+            }
+            long screenDozeTimeMs = mScreenUsageTimeRetriever.getScreenDozeTimeMs(display);
+            if (!mFirstSample) {
+                mLayout.setScreenDozeDuration(mPowerStats.stats, display,
+                        screenDozeTimeMs - mLastDozeTime[display]);
+            }
+            mLastDozeTime[display] = screenDozeTimeMs;
+        }
+
+        mPowerStats.uidStats.clear();
+
+        mScreenUsageTimeRetriever.retrieveTopActivityTimes((uid, topActivityTimeMs) -> {
+            long topActivityDuration = topActivityTimeMs - mLastTopActivityTime.get(uid);
+            if (topActivityDuration == 0) {
+                return;
+            }
+            mLastTopActivityTime.put(uid, topActivityTimeMs);
+
+            int mappedUid = mUidResolver.mapUid(uid);
+            long[] uidStats = mPowerStats.uidStats.get(mappedUid);
+            if (uidStats == null) {
+                uidStats = new long[mLayout.getUidStatsArrayLength()];
+                mPowerStats.uidStats.put(mappedUid, uidStats);
+            }
+
+            mLayout.setUidTopActivityDuration(uidStats,
+                    mLayout.getUidTopActivityDuration(uidStats) + topActivityDuration);
+        });
+
+        long elapsedRealtime = mClock.elapsedRealtime();
+        mPowerStats.durationMs = elapsedRealtime - mLastCollectionTime;
+        mLastCollectionTime = elapsedRealtime;
+
+        mFirstSample = false;
+
+        return mPowerStats;
+    }
+
+    private void collectEnergyConsumers() {
+        int voltageMv = mVoltageSupplier.getAsInt();
+        if (voltageMv <= 0) {
+            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
+                    + " mV) when querying energy consumers");
+            return;
+        }
+
+        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
+        mLastVoltageMv = voltageMv;
+
+        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mEnergyConsumerIds);
+        if (energyUws == null) {
+            return;
+        }
+
+        for (int i = energyUws.length - 1; i >= 0; i--) {
+            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
+                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
+            if (energyDelta < 0) {
+                // Likely, restart of powerstats HAL
+                energyDelta = 0;
+            }
+            mLayout.setConsumedEnergy(mPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
+            mLastConsumedEnergyUws[i] = energyUws[i];
+        }
+    }
+
+    @Override
+    protected void onUidRemoved(int uid) {
+        mLastTopActivityTime.delete(uid);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java
new file mode 100644
index 0000000..f134aa8
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java
@@ -0,0 +1,171 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.annotation.NonNull;
+import android.os.BatteryStats;
+import android.os.PersistableBundle;
+
+import com.android.internal.os.PowerStats;
+
+/**
+ * Captures the positions and lengths of sections of the stats array, such as time-in-state,
+ * power usage estimates etc.
+ */
+public class ScreenPowerStatsLayout extends PowerStatsLayout {
+    private static final String EXTRA_DEVICE_SCREEN_COUNT = "dsc";
+    private static final String EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION = "dsd";
+    private static final String EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS = "dbd";
+    private static final String EXTRA_DEVICE_DOZE_DURATION_POSITION = "ddd";
+    private static final String EXTRA_DEVICE_DOZE_POWER_POSITION = "ddp";
+    private static final String EXTRA_UID_FOREGROUND_DURATION = "uf";
+
+    private int mDisplayCount;
+    private int mDeviceScreenOnDurationPosition;
+    private int[] mDeviceBrightnessDurationPositions;
+    private int mDeviceScreenDozeDurationPosition;
+    private int mDeviceScreenDozePowerPosition;
+    private int mUidTopActivityTimePosition;
+
+    ScreenPowerStatsLayout() {
+    }
+
+    ScreenPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+        super(descriptor);
+    }
+
+    void addDeviceScreenUsageDurationSection(int displayCount) {
+        mDisplayCount = displayCount;
+        mDeviceScreenOnDurationPosition = addDeviceSection(displayCount, "on");
+        mDeviceBrightnessDurationPositions = new int[BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS];
+        for (int level = 0; level < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; level++) {
+            mDeviceBrightnessDurationPositions[level] =
+                    addDeviceSection(displayCount, BatteryStats.SCREEN_BRIGHTNESS_NAMES[level]);
+        }
+        mDeviceScreenDozeDurationPosition = addDeviceSection(displayCount, "doze");
+    }
+
+    @Override
+    public void addDeviceSectionPowerEstimate() {
+        super.addDeviceSectionPowerEstimate();
+        // Used by AmbientDisplayPowerStatsProcessor
+        mDeviceScreenDozePowerPosition = addDeviceSection(1, "doze-power", FLAG_HIDDEN);
+    }
+
+    public int getDisplayCount() {
+        return mDisplayCount;
+    }
+
+    /**
+     * Stores screen-on time for the specified display.
+     */
+    public void setScreenOnDuration(long[] stats, int display, long durationMs) {
+        stats[mDeviceScreenOnDurationPosition + display] = durationMs;
+    }
+
+    /**
+     * Returns screen-on time for the specified display.
+     */
+    public long getScreenOnDuration(long[] stats, int display) {
+        return stats[mDeviceScreenOnDurationPosition + display];
+    }
+
+    /**
+     * Stores time at the specified brightness level for the specified display.
+     */
+    public void setBrightnessLevelDuration(long[] stats, int display, int brightnessLevel,
+            long durationMs) {
+        stats[mDeviceBrightnessDurationPositions[brightnessLevel] + display] = durationMs;
+    }
+
+    /**
+     * Returns time at the specified brightness level for the specified display.
+     */
+    public long getBrightnessLevelDuration(long[] stats, int display, int brightnessLevel) {
+        return stats[mDeviceBrightnessDurationPositions[brightnessLevel] + display];
+    }
+
+    /**
+     * Stores time in the doze (ambient) state for the specified display.
+     */
+    public void setScreenDozeDuration(long[] stats, int display, long durationMs) {
+        stats[mDeviceScreenDozeDurationPosition + display] = durationMs;
+    }
+
+    /**
+     * Retrieves time in the doze (ambient) state for the specified display.
+     */
+    public long getScreenDozeDuration(long[] stats, int display) {
+        return stats[mDeviceScreenDozeDurationPosition + display];
+    }
+
+    /**
+     * Stores estimated power in the doze (ambient) state.
+     */
+    public void setScreenDozePowerEstimate(long[] stats, double power) {
+        stats[mDeviceScreenDozePowerPosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
+    }
+
+    /**
+     * Retrieves estimated power in the doze (ambient) state.
+     */
+    public double getScreenDozePowerEstimate(long[] stats) {
+        return stats[mDeviceScreenDozePowerPosition] / MILLI_TO_NANO_MULTIPLIER;
+    }
+
+    void addUidTopActivitiyDuration() {
+        mUidTopActivityTimePosition = addUidSection(1, "top");
+    }
+
+    /**
+     * Stores time the UID spent in the TOP state.
+     */
+    public void setUidTopActivityDuration(long[] stats, long durationMs) {
+        stats[mUidTopActivityTimePosition] = durationMs;
+    }
+
+    /**
+     * Returns time the UID spent in the TOP state.
+     */
+    public long getUidTopActivityDuration(long[] stats) {
+        return stats[mUidTopActivityTimePosition];
+    }
+
+    @Override
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putInt(EXTRA_DEVICE_SCREEN_COUNT, mDisplayCount);
+        extras.putInt(EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION, mDeviceScreenOnDurationPosition);
+        extras.putIntArray(EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS,
+                mDeviceBrightnessDurationPositions);
+        extras.putInt(EXTRA_DEVICE_DOZE_DURATION_POSITION, mDeviceScreenDozeDurationPosition);
+        extras.putInt(EXTRA_DEVICE_DOZE_POWER_POSITION, mDeviceScreenDozePowerPosition);
+        extras.putInt(EXTRA_UID_FOREGROUND_DURATION, mUidTopActivityTimePosition);
+    }
+
+    @Override
+    public void fromExtras(PersistableBundle extras) {
+        super.fromExtras(extras);
+        mDisplayCount = extras.getInt(EXTRA_DEVICE_SCREEN_COUNT, 1);
+        mDeviceScreenOnDurationPosition = extras.getInt(EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION);
+        mDeviceBrightnessDurationPositions = extras.getIntArray(
+                EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS);
+        mDeviceScreenDozeDurationPosition = extras.getInt(EXTRA_DEVICE_DOZE_DURATION_POSITION);
+        mDeviceScreenDozePowerPosition = extras.getInt(EXTRA_DEVICE_DOZE_POWER_POSITION);
+        mUidTopActivityTimePosition = extras.getInt(EXTRA_UID_FOREGROUND_DURATION);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java
new file mode 100644
index 0000000..908c751
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java
@@ -0,0 +1,239 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import static android.os.BatteryConsumer.PROCESS_STATE_ANY;
+
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_AMBIENT;
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_FULL;
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import android.os.BatteryStats;
+import android.util.Slog;
+
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class ScreenPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "ScreenPowerStatsProcessor";
+    private final int mDisplayCount;
+    private final UsageBasedPowerEstimator[] mScreenOnPowerEstimators;
+    private final UsageBasedPowerEstimator[] mScreenDozePowerEstimators;
+    private final UsageBasedPowerEstimator[][] mScreenBrightnessLevelPowerEstimators;
+    private PowerStats.Descriptor mLastUsedDescriptor;
+    private ScreenPowerStatsLayout mStatsLayout;
+    private PowerEstimationPlan mPlan;
+    private long[] mTmpDeviceStatsArray;
+    private long[] mTmpUidStatsArray;
+
+    private static class Intermediates {
+        public double power;
+    }
+
+    public ScreenPowerStatsProcessor(PowerProfile powerProfile) {
+        mDisplayCount = powerProfile.getNumDisplays();
+        mScreenOnPowerEstimators = new UsageBasedPowerEstimator[mDisplayCount];
+        mScreenDozePowerEstimators = new UsageBasedPowerEstimator[mDisplayCount];
+        mScreenBrightnessLevelPowerEstimators = new UsageBasedPowerEstimator[mDisplayCount][];
+        for (int display = 0; display < mDisplayCount; display++) {
+            mScreenOnPowerEstimators[display] = new UsageBasedPowerEstimator(
+                    powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, display));
+
+            double averagePowerFullBrightness = powerProfile.getAveragePowerForOrdinal(
+                    POWER_GROUP_DISPLAY_SCREEN_FULL, display);
+            mScreenBrightnessLevelPowerEstimators[display] =
+                    new UsageBasedPowerEstimator[BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS];
+            for (int bin = 0; bin < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; bin++) {
+                // For example, if the number of bins is 3, the corresponding averages
+                // are calculated as 0.5 * full, 1.5 * full, 2.5 * full
+                final double binPowerMah = averagePowerFullBrightness * (bin + 0.5)
+                        / BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS;
+                mScreenBrightnessLevelPowerEstimators[display][bin] =
+                        new UsageBasedPowerEstimator(binPowerMah);
+            }
+
+            mScreenDozePowerEstimators[display] = new UsageBasedPowerEstimator(
+                    powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, display));
+        }
+    }
+
+    private boolean unpackPowerStatsDescriptor(PowerStats.Descriptor descriptor) {
+        if (descriptor == null) {
+            return false;
+        }
+
+        if (descriptor.equals(mLastUsedDescriptor)) {
+            return true;
+        }
+
+        mLastUsedDescriptor = descriptor;
+        mStatsLayout = new ScreenPowerStatsLayout(descriptor);
+        if (mStatsLayout.getDisplayCount() != mDisplayCount) {
+            Slog.e(TAG, "Incompatible number of displays: " + mStatsLayout.getDisplayCount()
+                    + ", expected: " + mDisplayCount);
+            return false;
+        }
+
+        mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
+        mTmpUidStatsArray = new long[descriptor.uidStatsArrayLength];
+        return true;
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats, long timestampMs) {
+        if (!unpackPowerStatsDescriptor(stats.getPowerStatsDescriptor())) {
+            return;
+        }
+
+        if (mPlan == null) {
+            mPlan = new PowerEstimationPlan(stats.getConfig());
+        }
+
+        computeDevicePowerEstimates(stats);
+        combineDeviceStateEstimates();
+
+        List<Integer> uids = new ArrayList<>();
+        stats.collectUids(uids);
+
+        if (!uids.isEmpty()) {
+            computeUidPowerEstimates(stats, uids);
+        }
+        mPlan.resetIntermediates();
+    }
+
+    private void computeDevicePowerEstimates(PowerComponentAggregatedPowerStats stats) {
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+            if (!stats.getDeviceStats(mTmpDeviceStatsArray, estimation.stateValues)) {
+                continue;
+            }
+
+            if (estimation.stateValues[STATE_SCREEN] == SCREEN_STATE_ON) {
+                double power;
+                if (mStatsLayout.getEnergyConsumerCount() > 0) {
+                    power = uCtoMah(mStatsLayout.getConsumedEnergy(mTmpDeviceStatsArray, 0));
+                } else {
+                    power = 0;
+                    for (int display = 0; display < mStatsLayout.getDisplayCount(); display++) {
+                        power += computeDisplayPower(mTmpDeviceStatsArray, display);
+                    }
+                }
+                mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray, power);
+                Intermediates intermediates = new Intermediates();
+                intermediates.power = power;
+                estimation.intermediates = intermediates;
+            } else {
+                double power = 0;
+                if (mStatsLayout.getEnergyConsumerCount() > 0) {
+                    power = uCtoMah(mStatsLayout.getConsumedEnergy(mTmpDeviceStatsArray, 0));
+                } else {
+                    for (int display = 0; display < mStatsLayout.getDisplayCount(); display++) {
+                        power += mScreenDozePowerEstimators[display].calculatePower(
+                                mStatsLayout.getScreenDozeDuration(mTmpDeviceStatsArray, display));
+                    }
+                }
+                mStatsLayout.setScreenDozePowerEstimate(mTmpDeviceStatsArray, power);
+            }
+
+            stats.setDeviceStats(estimation.stateValues, mTmpDeviceStatsArray);
+        }
+    }
+
+    private double computeDisplayPower(long[] stats, int display) {
+        double power = mScreenOnPowerEstimators[display]
+                .calculatePower(mStatsLayout.getScreenOnDuration(stats, display));
+        for (int bin = 0; bin < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; bin++) {
+            power += mScreenBrightnessLevelPowerEstimators[display][bin]
+                    .calculatePower(mStatsLayout.getBrightnessLevelDuration(stats, display, bin));
+        }
+        return power;
+    }
+
+    /**
+     * Combine power estimates before distributing them proportionally to UIDs.
+     */
+    private void combineDeviceStateEstimates() {
+        for (int i = mPlan.combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
+            CombinedDeviceStateEstimate cdse = mPlan.combinedDeviceStateEstimations.get(i);
+            List<DeviceStateEstimation> deviceStateEstimations = cdse.deviceStateEstimations;
+            double power = 0;
+            for (int j = deviceStateEstimations.size() - 1; j >= 0; j--) {
+                DeviceStateEstimation dse = deviceStateEstimations.get(j);
+                Intermediates intermediates = (Intermediates) dse.intermediates;
+                if (intermediates != null) {
+                    power += intermediates.power;
+                }
+            }
+            if (power != 0) {
+                Intermediates cdseIntermediates = new Intermediates();
+                cdseIntermediates.power = power;
+                cdse.intermediates = cdseIntermediates;
+            }
+        }
+    }
+
+    private void computeUidPowerEstimates(PowerComponentAggregatedPowerStats stats,
+            List<Integer> uids) {
+        int[] uidStateValues = new int[stats.getConfig().getUidStateConfig().length];
+        uidStateValues[STATE_SCREEN] = SCREEN_STATE_ON;
+        uidStateValues[STATE_PROCESS_STATE] = PROCESS_STATE_ANY;
+
+        for (int i = mPlan.uidStateEstimates.size() - 1; i >= 0; i--) {
+            UidStateEstimate uidStateEstimate = mPlan.uidStateEstimates.get(i);
+            Intermediates intermediates =
+                    (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+            int[] deviceStateValues = uidStateEstimate.combinedDeviceStateEstimate
+                    .stateValues;
+            if (deviceStateValues[STATE_SCREEN] != SCREEN_STATE_ON
+                    || intermediates == null) {
+                continue;
+            }
+
+            uidStateValues[STATE_POWER] = deviceStateValues[STATE_POWER];
+
+            long totalTopActivityDuration = 0;
+            for (int j = uids.size() - 1; j >= 0; j--) {
+                int uid = uids.get(j);
+                if (stats.getUidStats(mTmpUidStatsArray, uid, uidStateValues)) {
+                    totalTopActivityDuration +=
+                            mStatsLayout.getUidTopActivityDuration(mTmpUidStatsArray);
+                }
+            }
+
+            if (totalTopActivityDuration == 0) {
+                return;
+            }
+
+            for (int j = uids.size() - 1; j >= 0; j--) {
+                int uid = uids.get(j);
+                if (stats.getUidStats(mTmpUidStatsArray, uid, uidStateValues)) {
+                    long duration = mStatsLayout.getUidTopActivityDuration(mTmpUidStatsArray);
+                    double power = intermediates.power * duration / totalTopActivityDuration;
+                    mStatsLayout.setUidPowerEstimate(mTmpUidStatsArray, power);
+                    stats.setUidStats(uid, uidStateValues, mTmpUidStatsArray);
+                }
+            }
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/SensorPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/SensorPowerStatsLayout.java
new file mode 100644
index 0000000..e66cd39
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/SensorPowerStatsLayout.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.os.PersistableBundle;
+import android.util.Slog;
+import android.util.SparseIntArray;
+
+public class SensorPowerStatsLayout extends PowerStatsLayout {
+    private static final String TAG = "SensorPowerStatsLayout";
+    private static final String EXTRA_DEVICE_SENSOR_HANDLES = "dsh";
+    private static final String EXTRA_UID_SENSOR_POSITIONS = "usp";
+
+    private final SparseIntArray mSensorPositions = new SparseIntArray();
+
+    void addUidSensorSection(int handle, String label) {
+        mSensorPositions.put(handle, addUidSection(1, label, FLAG_OPTIONAL));
+    }
+
+    /**
+     * Returns the position in the uid stats array of the duration element corresponding
+     * to the specified sensor identified by its handle.
+     */
+    public int getUidSensorDurationPosition(int handle) {
+        return mSensorPositions.get(handle, UNSUPPORTED);
+    }
+
+    /**
+     * Adds the specified duration to the accumulated timer for the specified sensor.
+     */
+    public void addUidSensorDuration(long[] stats, int handle, long durationMs) {
+        int position = mSensorPositions.get(handle, UNSUPPORTED);
+        if (position == UNSUPPORTED) {
+            Slog.e(TAG, "Unknown sensor: " + handle);
+            return;
+        }
+        stats[position] += durationMs;
+    }
+
+    @Override
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+
+        int[] handlers = new int[mSensorPositions.size()];
+        int[] uidDurationPositions = new int[mSensorPositions.size()];
+
+        for (int i = 0; i < mSensorPositions.size(); i++) {
+            handlers[i] = mSensorPositions.keyAt(i);
+            uidDurationPositions[i] = mSensorPositions.valueAt(i);
+        }
+
+        extras.putIntArray(EXTRA_DEVICE_SENSOR_HANDLES, handlers);
+        extras.putIntArray(EXTRA_UID_SENSOR_POSITIONS, uidDurationPositions);
+    }
+
+    @Override
+    public void fromExtras(PersistableBundle extras) {
+        super.fromExtras(extras);
+
+        int[] handlers = extras.getIntArray(EXTRA_DEVICE_SENSOR_HANDLES);
+        int[] uidDurationPositions = extras.getIntArray(EXTRA_UID_SENSOR_POSITIONS);
+
+        for (int i = 0; i < handlers.length; i++) {
+            mSensorPositions.put(handlers[i], uidDurationPositions[i]);
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/SensorPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/SensorPowerStatsProcessor.java
new file mode 100644
index 0000000..5bd3288
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/SensorPowerStatsProcessor.java
@@ -0,0 +1,312 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import android.hardware.Sensor;
+import android.hardware.SensorManager;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.PersistableBundle;
+import android.util.Slog;
+import android.util.SparseArray;
+
+import com.android.internal.os.PowerStats;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
+import java.util.function.Supplier;
+
+public class SensorPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "SensorPowerStatsProcessor";
+    private static final String ANDROID_SENSOR_TYPE_PREFIX = "android.sensor.";
+
+    private static final double MILLIS_IN_HOUR = 1000.0 * 60 * 60;
+    private static final String SENSOR_EVENT_TAG_PREFIX = "sensor:0x";
+    private final Supplier<SensorManager> mSensorManagerSupplier;
+
+    private static final long INITIAL_TIMESTAMP = -1;
+    private SensorManager mSensorManager;
+    private SensorPowerStatsLayout mStatsLayout;
+    private PowerStats mPowerStats;
+    private boolean mIsInitialized;
+    private PowerStats.Descriptor mDescriptor;
+    private long mLastUpdateTimestamp;
+    private PowerEstimationPlan mPlan;
+
+    private static class SensorState {
+        public int sensorHandle;
+        public boolean stateOn;
+        public int uid;
+        public long startTime = INITIAL_TIMESTAMP;
+    }
+
+    private static class Intermediates {
+        public double power;
+    }
+
+    private final SparseArray<SensorState> mSensorStates = new SparseArray<>();
+    private long[] mTmpDeviceStatsArray;
+    private long[] mTmpUidStatsArray;
+
+    public SensorPowerStatsProcessor(Supplier<SensorManager> sensorManagerSupplier) {
+        mSensorManagerSupplier = sensorManagerSupplier;
+    }
+
+    private boolean ensureInitialized() {
+        if (mIsInitialized) {
+            return true;
+        }
+
+        mSensorManager = mSensorManagerSupplier.get();
+        if (mSensorManager == null) {
+            return false;
+        }
+
+        mStatsLayout = new SensorPowerStatsLayout();
+        List<Sensor> sensorList = new ArrayList<>(mSensorManager.getSensorList(Sensor.TYPE_ALL));
+        sensorList.sort(Comparator.comparingInt(Sensor::getId));
+        for (int i = 0; i < sensorList.size(); i++) {
+            Sensor sensor = sensorList.get(i);
+            String label = makeLabel(sensor, sensorList);
+            mStatsLayout.addUidSensorSection(sensor.getHandle(), label);
+        }
+        mStatsLayout.addUidSectionPowerEstimate();
+        mStatsLayout.addDeviceSectionPowerEstimate();
+
+        PersistableBundle extras = new PersistableBundle();
+        mStatsLayout.toExtras(extras);
+        mDescriptor = new PowerStats.Descriptor(
+                BatteryConsumer.POWER_COMPONENT_SENSORS, mStatsLayout.getDeviceStatsArrayLength(),
+                null, 0, mStatsLayout.getUidStatsArrayLength(),
+                extras);
+
+        mPowerStats = new PowerStats(mDescriptor);
+        mTmpUidStatsArray = new long[mDescriptor.uidStatsArrayLength];
+        mTmpDeviceStatsArray = new long[mDescriptor.statsArrayLength];
+
+        mIsInitialized = true;
+        return true;
+    }
+
+    private String makeLabel(Sensor sensor, List<Sensor> sensorList) {
+        int type = sensor.getType();
+        String label = sensor.getStringType();
+
+        boolean isSingleton = true;
+        for (int i = sensorList.size() - 1; i >= 0; i--) {
+            Sensor s = sensorList.get(i);
+            if (s == sensor) {
+                continue;
+            }
+            if (s.getType() == type) {
+                isSingleton = false;
+                break;
+            }
+        }
+        if (!isSingleton) {
+            StringBuilder sb = new StringBuilder(label).append('.');
+            if (sensor.getId() > 0) { // 0 and -1 are reserved
+                sb.append(sensor.getId());
+            } else {
+                sb.append(sensor.getName());
+            }
+            label = sb.toString();
+        }
+        if (label.startsWith(ANDROID_SENSOR_TYPE_PREFIX)) {
+            label = label.substring(ANDROID_SENSOR_TYPE_PREFIX.length());
+        }
+        return label.replace(' ', '_');
+    }
+
+    @Override
+    void start(PowerComponentAggregatedPowerStats stats, long timestampMs) {
+        if (!ensureInitialized()) {
+            return;
+        }
+
+        // Establish a baseline at the beginning of an accumulation pass
+        mLastUpdateTimestamp = timestampMs;
+        flushPowerStats(stats, timestampMs);
+    }
+
+    @Override
+    void noteStateChange(PowerComponentAggregatedPowerStats stats, BatteryStats.HistoryItem item) {
+        if (!mIsInitialized) {
+            return;
+        }
+
+        if (item.eventTag == null || item.eventTag.string == null
+                || !item.eventTag.string.startsWith(SENSOR_EVENT_TAG_PREFIX)) {
+            return;
+        }
+
+        int sensorHandle;
+        try {
+            sensorHandle = Integer.parseInt(item.eventTag.string, SENSOR_EVENT_TAG_PREFIX.length(),
+                    item.eventTag.string.length(), 16);
+        } catch (NumberFormatException e) {
+            Slog.wtf(TAG, "Bad format of event tag: " + item.eventTag.string);
+            return;
+        }
+
+        SensorState sensor = mSensorStates.get(sensorHandle);
+        if (sensor == null) {
+            sensor = new SensorState();
+            sensor.sensorHandle = sensorHandle;
+            mSensorStates.put(sensorHandle, sensor);
+        }
+
+        int uid = item.eventTag.uid;
+        boolean sensorOn = (item.states & BatteryStats.HistoryItem.STATE_SENSOR_ON_FLAG) != 0;
+        if (sensorOn) {
+            if (!sensor.stateOn) {
+                sensor.stateOn = true;
+                sensor.uid = uid;
+                sensor.startTime = item.time;
+            } else if (sensor.uid != uid) {
+                recordUsageDuration(sensor, item.time);
+                sensor.uid = uid;
+            }
+        } else {
+            if (sensor.stateOn) {
+                recordUsageDuration(sensor, item.time);
+                sensor.stateOn = false;
+            }
+        }
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats, long timestampMs) {
+        if (!mIsInitialized) {
+            return;
+        }
+
+        for (int i = mSensorStates.size() - 1; i >= 0; i--) {
+            SensorState sensor = mSensorStates.valueAt(i);
+            if (sensor.stateOn) {
+                recordUsageDuration(sensor, timestampMs);
+            }
+        }
+        flushPowerStats(stats, timestampMs);
+
+        if (mPlan == null) {
+            mPlan = new PowerEstimationPlan(stats.getConfig());
+        }
+
+        List<Integer> uids = new ArrayList<>();
+        stats.collectUids(uids);
+
+        computeUidPowerEstimates(stats, uids);
+        computeDevicePowerEstimates(stats);
+
+        mPlan.resetIntermediates();
+    }
+
+    protected void recordUsageDuration(SensorState sensorState, long time) {
+        long durationMs = Math.max(0, time - sensorState.startTime);
+        if (durationMs > 0) {
+            long[] uidStats = mPowerStats.uidStats.get(sensorState.uid);
+            if (uidStats == null) {
+                uidStats = new long[mDescriptor.uidStatsArrayLength];
+                mPowerStats.uidStats.put(sensorState.uid, uidStats);
+            }
+            mStatsLayout.addUidSensorDuration(uidStats, sensorState.sensorHandle, durationMs);
+        }
+        sensorState.startTime = time;
+    }
+
+    private void flushPowerStats(PowerComponentAggregatedPowerStats stats, long timestamp) {
+        mPowerStats.durationMs = timestamp - mLastUpdateTimestamp;
+        stats.addPowerStats(mPowerStats, timestamp);
+
+        Arrays.fill(mPowerStats.stats, 0);
+        mPowerStats.uidStats.clear();
+        mLastUpdateTimestamp = timestamp;
+    }
+
+    private void computeUidPowerEstimates(PowerComponentAggregatedPowerStats stats,
+            List<Integer> uids) {
+        List<Sensor> sensorList = mSensorManager.getSensorList(Sensor.TYPE_ALL);
+        int[] uidSensorDurationPositions = new int[sensorList.size()];
+        double[] sensorPower = new double[sensorList.size()];
+        for (int i = sensorList.size() - 1; i >= 0; i--) {
+            Sensor sensor = sensorList.get(i);
+            uidSensorDurationPositions[i] =
+                    mStatsLayout.getUidSensorDurationPosition(sensor.getHandle());
+            sensorPower[i] = sensor.getPower() / MILLIS_IN_HOUR;
+        }
+
+        for (int i = mPlan.uidStateEstimates.size() - 1; i >= 0; i--) {
+            UidStateEstimate uidStateEstimate = mPlan.uidStateEstimates.get(i);
+            List<UidStateProportionalEstimate> proportionalEstimates =
+                    uidStateEstimate.proportionalEstimates;
+            for (int j = proportionalEstimates.size() - 1; j >= 0; j--) {
+                UidStateProportionalEstimate proportionalEstimate = proportionalEstimates.get(j);
+                for (int k = uids.size() - 1; k >= 0; k--) {
+                    int uid = uids.get(k);
+                    if (!stats.getUidStats(mTmpUidStatsArray, uid,
+                            proportionalEstimate.stateValues)) {
+                        continue;
+                    }
+                    double power = 0;
+                    for (int m = 0; m < uidSensorDurationPositions.length; m++) {
+                        int position = uidSensorDurationPositions[m];
+                        if (position == PowerStatsLayout.UNSUPPORTED
+                                || mTmpUidStatsArray[position] == 0) {
+                            continue;
+                        }
+                        power += sensorPower[m] * mTmpUidStatsArray[position];
+                    }
+                    if (power == 0) {
+                        continue;
+                    }
+
+                    mStatsLayout.setUidPowerEstimate(mTmpUidStatsArray, power);
+                    stats.setUidStats(uid, proportionalEstimate.stateValues, mTmpUidStatsArray);
+
+                    Intermediates intermediates = (Intermediates) uidStateEstimate
+                            .combinedDeviceStateEstimate.intermediates;
+                    if (intermediates == null) {
+                        intermediates = new Intermediates();
+                        uidStateEstimate.combinedDeviceStateEstimate.intermediates = intermediates;
+                    }
+                    intermediates.power += power;
+                }
+            }
+        }
+    }
+
+    private void computeDevicePowerEstimates(PowerComponentAggregatedPowerStats stats) {
+        for (int i = mPlan.combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
+            CombinedDeviceStateEstimate estimation =
+                    mPlan.combinedDeviceStateEstimations.get(i);
+            if (estimation.intermediates == null) {
+                continue;
+            }
+
+            if (!stats.getDeviceStats(mTmpDeviceStatsArray, estimation.stateValues)) {
+                continue;
+            }
+
+            mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray,
+                    ((Intermediates) estimation.intermediates).power);
+            stats.setDeviceStats(estimation.stateValues, mTmpDeviceStatsArray);
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
index bca81f52..c21f783 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -64,6 +64,7 @@
 import static com.android.internal.util.FrameworkStatsLog.TIME_ZONE_DETECTOR_STATE__DETECTION_MODE__UNKNOWN;
 import static com.android.server.am.MemoryStatUtil.readMemoryStatFromFilesystem;
 import static com.android.server.stats.Flags.addMobileBytesTransferByProcStatePuller;
+import static com.android.server.stats.Flags.applyNetworkStatsPollRateLimit;
 import static com.android.server.stats.pull.IonMemoryUtil.readProcessSystemIonHeapSizesFromDebugfs;
 import static com.android.server.stats.pull.IonMemoryUtil.readSystemIonHeapSizeFromDebugfs;
 import static com.android.server.stats.pull.ProcfsMemoryUtil.getProcessCmdlines;
@@ -298,6 +299,13 @@
      */
     private static final long NETSTATS_UID_DEFAULT_BUCKET_DURATION_MS = HOURS.toMillis(2);
 
+    /**
+     * Polling NetworkStats is a heavy operation and it should be done sparingly. Atom pulls may
+     * happen in bursts, but these should be infrequent. The poll rate limit ensures that data is
+     * sufficiently fresh (i.e. not stale) while reducing system load during atom pull bursts.
+     */
+    private static final long NETSTATS_POLL_RATE_LIMIT_MS = 15000;
+
     private static final int CPU_TIME_PER_THREAD_FREQ_MAX_NUM_FREQUENCIES = 8;
     private static final int OP_FLAGS_PULLED = OP_FLAG_SELF | OP_FLAG_TRUSTED_PROXIED;
     private static final String COMMON_PERMISSION_PREFIX = "android.permission.";
@@ -415,6 +423,9 @@
     @GuardedBy("mDataBytesTransferLock")
     private final ArrayList<NetworkStatsExt> mNetworkStatsBaselines = new ArrayList<>();
 
+    @GuardedBy("mDataBytesTransferLock")
+    private long mLastNetworkStatsPollTime = -NETSTATS_POLL_RATE_LIMIT_MS;
+
     // Listener for monitoring subscriptions changed event.
     private StatsSubscriptionsListener mStatsSubscriptionsListener;
     // List that stores SubInfo of subscriptions that ever appeared since boot.
@@ -1063,24 +1074,26 @@
         // Initialize NetworkStats baselines.
         synchronized (mDataBytesTransferLock) {
             mNetworkStatsBaselines.addAll(
-                    collectNetworkStatsSnapshotForAtom(FrameworkStatsLog.WIFI_BYTES_TRANSFER));
+                    collectNetworkStatsSnapshotForAtomLocked(
+                            FrameworkStatsLog.WIFI_BYTES_TRANSFER));
             mNetworkStatsBaselines.addAll(
-                    collectNetworkStatsSnapshotForAtom(
+                    collectNetworkStatsSnapshotForAtomLocked(
                             FrameworkStatsLog.WIFI_BYTES_TRANSFER_BY_FG_BG));
             mNetworkStatsBaselines.addAll(
-                    collectNetworkStatsSnapshotForAtom(FrameworkStatsLog.MOBILE_BYTES_TRANSFER));
-            mNetworkStatsBaselines.addAll(collectNetworkStatsSnapshotForAtom(
+                    collectNetworkStatsSnapshotForAtomLocked(
+                            FrameworkStatsLog.MOBILE_BYTES_TRANSFER));
+            mNetworkStatsBaselines.addAll(collectNetworkStatsSnapshotForAtomLocked(
                     FrameworkStatsLog.MOBILE_BYTES_TRANSFER_BY_FG_BG));
-            mNetworkStatsBaselines.addAll(collectNetworkStatsSnapshotForAtom(
+            mNetworkStatsBaselines.addAll(collectNetworkStatsSnapshotForAtomLocked(
                     FrameworkStatsLog.BYTES_TRANSFER_BY_TAG_AND_METERED));
             mNetworkStatsBaselines.addAll(
-                    collectNetworkStatsSnapshotForAtom(
+                    collectNetworkStatsSnapshotForAtomLocked(
                             FrameworkStatsLog.DATA_USAGE_BYTES_TRANSFER));
             mNetworkStatsBaselines.addAll(
-                    collectNetworkStatsSnapshotForAtom(
+                    collectNetworkStatsSnapshotForAtomLocked(
                             FrameworkStatsLog.OEM_MANAGED_BYTES_TRANSFER));
             if (canQueryTypeProxy) {
-                mNetworkStatsBaselines.addAll(collectNetworkStatsSnapshotForAtom(
+                mNetworkStatsBaselines.addAll(collectNetworkStatsSnapshotForAtomLocked(
                         FrameworkStatsLog.PROXY_BYTES_TRANSFER_BY_FG_BG));
             }
         }
@@ -1243,12 +1256,14 @@
         );
     }
 
+    @GuardedBy("mDataBytesTransferLock")
     @NonNull
-    private List<NetworkStatsExt> collectNetworkStatsSnapshotForAtom(int atomTag) {
+    private List<NetworkStatsExt> collectNetworkStatsSnapshotForAtomLocked(int atomTag) {
         List<NetworkStatsExt> ret = new ArrayList<>();
         switch (atomTag) {
             case FrameworkStatsLog.WIFI_BYTES_TRANSFER: {
-                final NetworkStats stats = getUidNetworkStatsSnapshotForTransport(TRANSPORT_WIFI);
+                final NetworkStats stats = getUidNetworkStatsSnapshotForTransportLocked(
+                        TRANSPORT_WIFI);
                 if (stats != null) {
                     ret.add(new NetworkStatsExt(sliceNetworkStatsByUid(stats),
                             new int[]{TRANSPORT_WIFI}, /*slicedByFgbg=*/false));
@@ -1256,7 +1271,8 @@
                 break;
             }
             case FrameworkStatsLog.WIFI_BYTES_TRANSFER_BY_FG_BG: {
-                final NetworkStats stats = getUidNetworkStatsSnapshotForTransport(TRANSPORT_WIFI);
+                final NetworkStats stats = getUidNetworkStatsSnapshotForTransportLocked(
+                        TRANSPORT_WIFI);
                 if (stats != null) {
                     ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
                             new int[]{TRANSPORT_WIFI}, /*slicedByFgbg=*/true));
@@ -1265,7 +1281,7 @@
             }
             case FrameworkStatsLog.MOBILE_BYTES_TRANSFER: {
                 final NetworkStats stats =
-                        getUidNetworkStatsSnapshotForTransport(TRANSPORT_CELLULAR);
+                        getUidNetworkStatsSnapshotForTransportLocked(TRANSPORT_CELLULAR);
                 if (stats != null) {
                     ret.add(new NetworkStatsExt(sliceNetworkStatsByUid(stats),
                             new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/false));
@@ -1274,7 +1290,7 @@
             }
             case FrameworkStatsLog.MOBILE_BYTES_TRANSFER_BY_FG_BG: {
                 final NetworkStats stats =
-                        getUidNetworkStatsSnapshotForTransport(TRANSPORT_CELLULAR);
+                        getUidNetworkStatsSnapshotForTransportLocked(TRANSPORT_CELLULAR);
                 if (stats != null) {
                     ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
                             new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/true));
@@ -1282,7 +1298,7 @@
                 break;
             }
             case FrameworkStatsLog.PROXY_BYTES_TRANSFER_BY_FG_BG: {
-                final NetworkStats stats = getUidNetworkStatsSnapshotForTemplate(
+                final NetworkStats stats = getUidNetworkStatsSnapshotForTemplateLocked(
                         new NetworkTemplate.Builder(MATCH_PROXY).build(),  /*includeTags=*/false);
                 if (stats != null) {
                     ret.add(new NetworkStatsExt(sliceNetworkStatsByUidTagAndMetered(stats),
@@ -1294,9 +1310,9 @@
                 break;
             }
             case FrameworkStatsLog.BYTES_TRANSFER_BY_TAG_AND_METERED: {
-                final NetworkStats wifiStats = getUidNetworkStatsSnapshotForTemplate(
+                final NetworkStats wifiStats = getUidNetworkStatsSnapshotForTemplateLocked(
                         new NetworkTemplate.Builder(MATCH_WIFI).build(), /*includeTags=*/true);
-                final NetworkStats cellularStats = getUidNetworkStatsSnapshotForTemplate(
+                final NetworkStats cellularStats = getUidNetworkStatsSnapshotForTemplateLocked(
                         new NetworkTemplate.Builder(MATCH_MOBILE)
                                 .setMeteredness(METERED_YES).build(), /*includeTags=*/true);
                 if (wifiStats != null && cellularStats != null) {
@@ -1311,12 +1327,12 @@
             }
             case FrameworkStatsLog.DATA_USAGE_BYTES_TRANSFER: {
                 for (final SubInfo subInfo : mHistoricalSubs) {
-                    ret.addAll(getDataUsageBytesTransferSnapshotForSub(subInfo));
+                    ret.addAll(getDataUsageBytesTransferSnapshotForSubLocked(subInfo));
                 }
                 break;
             }
             case FrameworkStatsLog.OEM_MANAGED_BYTES_TRANSFER: {
-                ret.addAll(getDataUsageBytesTransferSnapshotForOemManaged());
+                ret.addAll(getDataUsageBytesTransferSnapshotForOemManagedLocked());
                 break;
             }
             default:
@@ -1325,8 +1341,9 @@
         return ret;
     }
 
+    @GuardedBy("mDataBytesTransferLock")
     private int pullDataBytesTransferLocked(int atomTag, @NonNull List<StatsEvent> pulledData) {
-        final List<NetworkStatsExt> current = collectNetworkStatsSnapshotForAtom(atomTag);
+        final List<NetworkStatsExt> current = collectNetworkStatsSnapshotForAtomLocked(atomTag);
 
         if (current == null) {
             Slog.e(TAG, "current snapshot is null for " + atomTag + ", return.");
@@ -1459,8 +1476,9 @@
         }
     }
 
+    @GuardedBy("mDataBytesTransferLock")
     @NonNull
-    private List<NetworkStatsExt> getDataUsageBytesTransferSnapshotForOemManaged() {
+    private List<NetworkStatsExt> getDataUsageBytesTransferSnapshotForOemManagedLocked() {
         final List<Pair<Integer, Integer>> matchRulesAndTransports = List.of(
                 new Pair(MATCH_ETHERNET, TRANSPORT_ETHERNET),
                 new Pair(MATCH_MOBILE, TRANSPORT_CELLULAR),
@@ -1479,7 +1497,8 @@
                 // Thus, specifying networks through their identifiers are not needed.
                 final NetworkTemplate template = new NetworkTemplate.Builder(matchRule)
                         .setOemManaged(oemManaged).build();
-                final NetworkStats stats = getUidNetworkStatsSnapshotForTemplate(template, false);
+                final NetworkStats stats = getUidNetworkStatsSnapshotForTemplateLocked(
+                        template, false);
                 final Integer transport = ruleAndTransport.second;
                 if (stats != null) {
                     ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
@@ -1496,8 +1515,9 @@
     /**
      * Create a snapshot of NetworkStats for a given transport.
      */
+    @GuardedBy("mDataBytesTransferLock")
     @Nullable
-    private NetworkStats getUidNetworkStatsSnapshotForTransport(int transport) {
+    private NetworkStats getUidNetworkStatsSnapshotForTransportLocked(int transport) {
         NetworkTemplate template = null;
         switch (transport) {
             case TRANSPORT_CELLULAR:
@@ -1510,7 +1530,7 @@
             default:
                 Log.wtf(TAG, "Unexpected transport.");
         }
-        return getUidNetworkStatsSnapshotForTemplate(template, /*includeTags=*/false);
+        return getUidNetworkStatsSnapshotForTemplateLocked(template, /*includeTags=*/false);
     }
 
     /**
@@ -1534,8 +1554,9 @@
      * Note that this should be only used to calculate diff since the snapshot might contains
      * some traffic before boot.
      */
+    @GuardedBy("mDataBytesTransferLock")
     @Nullable
-    private NetworkStats getUidNetworkStatsSnapshotForTemplate(
+    private NetworkStats getUidNetworkStatsSnapshotForTemplateLocked(
             @NonNull NetworkTemplate template, boolean includeTags) {
         final long elapsedMillisSinceBoot = SystemClock.elapsedRealtime();
         final long currentTimeInMillis = MICROSECONDS.toMillis(SystemClock.currentTimeMicro());
@@ -1547,13 +1568,19 @@
         final long startTime = currentTimeInMillis - elapsedMillisSinceBoot - bucketDuration;
         final long endTime = currentTimeInMillis + bucketDuration;
 
-        // TODO (b/156313635): This is short-term hack to allow perfd gets updated networkStats
-        //  history when query in every second in order to show realtime statistics. However,
-        //  this is not a good long-term solution since NetworkStatsService will make frequent
-        //  I/O and also block main thread when polling.
-        //  Consider making perfd queries NetworkStatsService directly.
-        if (template.getMatchRule() == MATCH_WIFI && template.getSubscriberIds().isEmpty()) {
-            getNetworkStatsManager().forceUpdate();
+        // NetworkStatsManager#forceUpdate updates stats for all networks
+        if (applyNetworkStatsPollRateLimit()) {
+            // The new way: rate-limit force-polling for all NetworkStats queries
+            if (elapsedMillisSinceBoot - mLastNetworkStatsPollTime >= NETSTATS_POLL_RATE_LIMIT_MS) {
+                mLastNetworkStatsPollTime = elapsedMillisSinceBoot;
+                getNetworkStatsManager().forceUpdate();
+            }
+        } else {
+            // The old way: force-poll only on WiFi queries. Data for other queries can be stale
+            // if there was no recent poll beforehand (e.g. for WiFi or scheduled poll)
+            if (template.getMatchRule() == MATCH_WIFI && template.getSubscriberIds().isEmpty()) {
+                getNetworkStatsManager().forceUpdate();
+            }
         }
 
         final android.app.usage.NetworkStats queryNonTaggedStats =
@@ -1572,8 +1599,9 @@
         return nonTaggedStats.add(taggedStats);
     }
 
+    @GuardedBy("mDataBytesTransferLock")
     @NonNull
-    private List<NetworkStatsExt> getDataUsageBytesTransferSnapshotForSub(
+    private List<NetworkStatsExt> getDataUsageBytesTransferSnapshotForSubLocked(
             @NonNull SubInfo subInfo) {
         final List<NetworkStatsExt> ret = new ArrayList<>();
         for (final int ratType : getAllCollapsedRatTypes()) {
@@ -1583,7 +1611,7 @@
                             .setRatType(ratType)
                             .setMeteredness(METERED_YES).build();
             final NetworkStats stats =
-                    getUidNetworkStatsSnapshotForTemplate(template, /*includeTags=*/false);
+                    getUidNetworkStatsSnapshotForTemplateLocked(template, /*includeTags=*/false);
             if (stats != null) {
                 ret.add(new NetworkStatsExt(sliceNetworkStatsByFgbg(stats),
                         new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/true,
@@ -5273,6 +5301,13 @@
 
         @Override
         public void onSubscriptionsChanged() {
+            synchronized (mDataBytesTransferLock) {
+                onSubscriptionsChangedLocked();
+            }
+        }
+
+        @GuardedBy("mDataBytesTransferLock")
+        private void onSubscriptionsChangedLocked() {
             final List<SubscriptionInfo> currentSubs = mSm.getCompleteActiveSubscriptionInfoList();
             for (final SubscriptionInfo sub : currentSubs) {
                 final SubInfo match = CollectionUtils.find(mHistoricalSubs,
@@ -5295,12 +5330,11 @@
                         subscriberId, sub.isOpportunistic());
                 Slog.i(TAG, "subId " + subId + " added into historical sub list");
 
-                synchronized (mDataBytesTransferLock) {
-                    mHistoricalSubs.add(subInfo);
-                    // Since getting snapshot when pulling will also include data before boot,
-                    // query stats as baseline to prevent double count is needed.
-                    mNetworkStatsBaselines.addAll(getDataUsageBytesTransferSnapshotForSub(subInfo));
-                }
+                mHistoricalSubs.add(subInfo);
+                // Since getting snapshot when pulling will also include data before boot,
+                // query stats as baseline to prevent double count is needed.
+                mNetworkStatsBaselines.addAll(
+                        getDataUsageBytesTransferSnapshotForSubLocked(subInfo));
             }
         }
     }
diff --git a/services/core/java/com/android/server/stats/stats_flags.aconfig b/services/core/java/com/android/server/stats/stats_flags.aconfig
index 6faa273..f360837 100644
--- a/services/core/java/com/android/server/stats/stats_flags.aconfig
+++ b/services/core/java/com/android/server/stats/stats_flags.aconfig
@@ -8,3 +8,11 @@
     bug: "309512867"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "apply_network_stats_poll_rate_limit"
+    namespace: "statsd"
+    description: "Apply a rate limit for polling network stats when pulling relevant atoms"
+    bug: "352495181"
+    is_fixed_read_only: true
+}
diff --git a/services/core/java/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessor.java b/services/core/java/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessor.java
index aa2b74e..58c3ba5 100644
--- a/services/core/java/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessor.java
+++ b/services/core/java/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessor.java
@@ -56,12 +56,17 @@
         // enables immediate failover to a secondary provider, one that might provide valid IDs for
         // the same location, which should provide better behavior than just ignoring the event.
         if (hasInvalidZones(event)) {
-            TimeZoneProviderStatus providerStatus = new TimeZoneProviderStatus.Builder(
-                    event.getTimeZoneProviderStatus())
-                    .setTimeZoneResolutionOperationStatus(OPERATION_STATUS_FAILED)
-                    .build();
-            return TimeZoneProviderEvent.createUncertainEvent(
-                    event.getCreationElapsedMillis(), providerStatus);
+            TimeZoneProviderStatus providerStatus = event.getTimeZoneProviderStatus();
+            TimeZoneProviderStatus.Builder providerStatusBuilder;
+            if (providerStatus != null) {
+                providerStatusBuilder = new TimeZoneProviderStatus.Builder(providerStatus);
+            } else {
+                providerStatusBuilder = new TimeZoneProviderStatus.Builder();
+            }
+            return TimeZoneProviderEvent.createUncertainEvent(event.getCreationElapsedMillis(),
+                    providerStatusBuilder
+                            .setTimeZoneResolutionOperationStatus(OPERATION_STATUS_FAILED)
+                            .build());
         }
 
         return event;
diff --git a/services/core/java/com/android/server/vibrator/VibratorControlService.java b/services/core/java/com/android/server/vibrator/VibratorControlService.java
index 4da6585..de5e662 100644
--- a/services/core/java/com/android/server/vibrator/VibratorControlService.java
+++ b/services/core/java/com/android/server/vibrator/VibratorControlService.java
@@ -41,6 +41,7 @@
 import android.os.SystemClock;
 import android.os.VibrationAttributes;
 import android.os.VibrationEffect;
+import android.os.vibrator.Flags;
 import android.util.IndentingPrintWriter;
 import android.util.IntArray;
 import android.util.Slog;
@@ -265,9 +266,15 @@
                 return null;
             }
 
+            if (Flags.throttleVibrationParamsRequests() && mVibrationParamRequest != null
+                    && mVibrationParamRequest.usage == usage) {
+                // Reuse existing future for ongoing request with same usage.
+                return mVibrationParamRequest.future;
+            }
+
             try {
                 endOngoingRequestVibrationParamsLocked(/* wasCancelled= */ true);
-                mVibrationParamRequest = new VibrationParamRequest(uid);
+                mVibrationParamRequest = new VibrationParamRequest(uid, usage);
                 vibratorController.requestVibrationParams(vibrationType, timeoutInMillis,
                         mVibrationParamRequest.token);
                 return mVibrationParamRequest.future;
@@ -533,10 +540,12 @@
         public final CompletableFuture<Void> future = new CompletableFuture<>();
         public final IBinder token = new Binder();
         public final int uid;
+        public final @VibrationAttributes.Usage int usage;
         public final long uptimeMs;
 
-        VibrationParamRequest(int uid) {
+        VibrationParamRequest(int uid, @VibrationAttributes.Usage int usage) {
             this.uid = uid;
+            this.usage = usage;
             uptimeMs = SystemClock.uptimeMillis();
         }
 
diff --git a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
index 72c7be3..ba2594a 100644
--- a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
+++ b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
@@ -3844,6 +3844,7 @@
             pw.print("  mPadding="); pw.println(wpSize.mPadding);
         });
         pw.print("  mCropHint="); pw.println(wallpaper.cropHint);
+        if (multiCrop()) pw.print("  mCropHints="); pw.println(wallpaper.mCropHints);
         pw.print("  mName=");  pw.println(wallpaper.name);
         pw.print("  mAllowBackup="); pw.println(wallpaper.allowBackup);
         pw.print("  mWallpaperComponent="); pw.println(wallpaper.wallpaperComponent);
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index d039b04..7d70ea1 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -495,7 +495,7 @@
     final String launchedFromPackage; // always the package who started the activity.
     @Nullable
     final String launchedFromFeatureId; // always the feature in launchedFromPackage
-    private final int mLaunchSourceType; // original launch source type
+    int mLaunchSourceType; // latest launch source type
     final Intent intent;    // the original intent that generated us
     final String shortComponentName; // the short component name of the intent
     final String resolvedType; // as per original caller;
@@ -2451,6 +2451,10 @@
         return mLaunchSourceType == type;
     }
 
+    void updateLaunchSourceType(int launchFromUid, WindowProcessController caller) {
+        mLaunchSourceType = determineLaunchSourceType(launchFromUid, caller);
+    }
+
     private int determineLaunchSourceType(int launchFromUid, WindowProcessController caller) {
         if (launchFromUid == Process.SYSTEM_UID || launchFromUid == Process.ROOT_UID) {
             return LAUNCH_SOURCE_TYPE_SYSTEM;
@@ -4290,7 +4294,8 @@
     }
 
     void finishRelaunching() {
-        mLetterboxUiController.setRelaunchingAfterRequestedOrientationChanged(false);
+        mAppCompatController.getAppCompatOrientationOverrides()
+                .setRelaunchingAfterRequestedOrientationChanged(false);
         mTaskSupervisor.getActivityMetricsLogger().notifyActivityRelaunched(this);
 
         if (mPendingRelaunchCount > 0) {
@@ -8181,7 +8186,8 @@
                 mLastReportedConfiguration.getMergedConfiguration())) {
             ensureActivityConfiguration(false /* ignoreVisibility */);
             if (mPendingRelaunchCount > originalRelaunchingCount) {
-                mLetterboxUiController.setRelaunchingAfterRequestedOrientationChanged(true);
+                mAppCompatController.getAppCompatOrientationOverrides()
+                        .setRelaunchingAfterRequestedOrientationChanged(true);
             }
             if (mTransitionController.inPlayingTransition(this)) {
                 mTransitionController.mValidateActivityCompat.add(this);
@@ -8397,7 +8403,8 @@
      */
     @ActivityInfo.SizeChangesSupportMode
     private int supportsSizeChanges() {
-        if (mLetterboxUiController.shouldOverrideForceNonResizeApp()) {
+        if (mAppCompatController.getAppCompatResizeOverrides()
+                .shouldOverrideForceNonResizeApp()) {
             return SIZE_CHANGES_UNSUPPORTED_OVERRIDE;
         }
 
@@ -8405,7 +8412,8 @@
             return SIZE_CHANGES_SUPPORTED_METADATA;
         }
 
-        if (mLetterboxUiController.shouldOverrideForceResizeApp()) {
+        if (mAppCompatController.getAppCompatResizeOverrides()
+                .shouldOverrideForceResizeApp()) {
             return SIZE_CHANGES_SUPPORTED_OVERRIDE;
         }
 
@@ -10488,7 +10496,7 @@
                 mAppCompatController.getAppCompatOrientationOverrides()
                         .shouldIgnoreOrientationRequestLoop());
         proto.write(SHOULD_OVERRIDE_FORCE_RESIZE_APP,
-                mLetterboxUiController.shouldOverrideForceResizeApp());
+                mAppCompatController.getAppCompatResizeOverrides().shouldOverrideForceResizeApp());
         proto.write(SHOULD_ENABLE_USER_ASPECT_RATIO_SETTINGS,
                 mAppCompatController.getAppCompatAspectRatioOverrides()
                         .shouldEnableUserAspectRatioSettings());
diff --git a/services/core/java/com/android/server/wm/ActivityStarter.java b/services/core/java/com/android/server/wm/ActivityStarter.java
index 5bfe9d7..c89f3a3 100644
--- a/services/core/java/com/android/server/wm/ActivityStarter.java
+++ b/services/core/java/com/android/server/wm/ActivityStarter.java
@@ -1793,6 +1793,9 @@
                     activity.destroyIfPossible("Removes redundant singleInstance");
                 }
             }
+            if (mLastStartActivityRecord != null) {
+                targetTaskTop.mLaunchSourceType = mLastStartActivityRecord.mLaunchSourceType;
+            }
             targetTaskTop.mTransitionController.collect(targetTaskTop);
             recordTransientLaunchIfNeeded(targetTaskTop);
             // Recycle the target task for this launch.
diff --git a/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java b/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
index b0d8925..afdbc0a 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
@@ -2801,6 +2801,13 @@
                                 targetActivity, activityOptions);
                     }
 
+                    if (callingPid > 0) {
+                        final WindowProcessController wpc = mService.mProcessMap
+                                .getProcess(callingPid);
+                        if (wpc != null) {
+                            targetActivity.updateLaunchSourceType(callingUid, wpc);
+                        }
+                    }
                     mService.getActivityStartController().postStartActivityProcessingForLastStarter(
                             task.getTopNonFinishingActivity(), ActivityManager.START_TASK_TO_FRONT,
                             task.getRootTask());
diff --git a/services/core/java/com/android/server/wm/AppCompatAspectRatioOverrides.java b/services/core/java/com/android/server/wm/AppCompatAspectRatioOverrides.java
index 05d4c82..25cb134 100644
--- a/services/core/java/com/android/server/wm/AppCompatAspectRatioOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatAspectRatioOverrides.java
@@ -36,6 +36,7 @@
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.AppCompatConfiguration.LETTERBOX_POSITION_MULTIPLIER_CENTER;
 import static com.android.server.wm.AppCompatConfiguration.MIN_FIXED_ORIENTATION_LETTERBOX_ASPECT_RATIO;
+import static com.android.server.wm.AppCompatUtils.isChangeEnabled;
 
 import android.annotation.NonNull;
 import android.content.pm.PackageManager;
@@ -115,7 +116,7 @@
      */
     boolean shouldOverrideMinAspectRatio() {
         return mAllowMinAspectRatioOverrideOptProp.shouldEnableWithOptInOverrideAndOptOutProperty(
-                isCompatChangeEnabled(OVERRIDE_MIN_ASPECT_RATIO));
+                isChangeEnabled(mActivityRecord, OVERRIDE_MIN_ASPECT_RATIO));
     }
 
     /**
@@ -154,7 +155,7 @@
     }
 
     boolean isSystemOverrideToFullscreenEnabled() {
-        return isCompatChangeEnabled(OVERRIDE_ANY_ORIENTATION_TO_USER)
+        return isChangeEnabled(mActivityRecord, OVERRIDE_ANY_ORIENTATION_TO_USER)
                 && !mAllowOrientationOverrideOptProp.isFalse()
                 && (mUserAspectRatioState.mUserAspectRatio == USER_MIN_ASPECT_RATIO_UNSET
                 || mUserAspectRatioState.mUserAspectRatio == USER_MIN_ASPECT_RATIO_FULLSCREEN);
@@ -302,10 +303,6 @@
         private int mUserAspectRatio = USER_MIN_ASPECT_RATIO_UNSET;
     }
 
-    private boolean isCompatChangeEnabled(long overrideChangeId) {
-        return mActivityRecord.info.isChangeEnabled(overrideChangeId);
-    }
-
     private Resources getResources() {
         return mActivityRecord.mWmService.mContext.getResources();
     }
diff --git a/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java b/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
index 93a8663..aeaaffd 100644
--- a/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
@@ -30,6 +30,7 @@
 
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_ATM;
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
+import static com.android.server.wm.AppCompatUtils.isChangeEnabled;
 
 import android.annotation.NonNull;
 import android.app.CameraCompatTaskInfo.FreeformCameraCompatMode;
@@ -99,7 +100,8 @@
     boolean shouldOverrideMinAspectRatioForCamera() {
         return isCameraActive() && mAllowMinAspectRatioOverrideOptProp
                 .shouldEnableWithOptInOverrideAndOptOutProperty(
-                        isCompatChangeEnabled(OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA));
+                        isChangeEnabled(mActivityRecord,
+                                OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA));
     }
 
     /**
@@ -115,7 +117,7 @@
      */
     boolean shouldRefreshActivityForCameraCompat() {
         return mCameraCompatAllowRefreshOptProp.shouldEnableWithOptOutOverrideAndProperty(
-                isCompatChangeEnabled(OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH));
+                isChangeEnabled(mActivityRecord, OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH));
     }
 
     /**
@@ -134,7 +136,7 @@
      */
     boolean shouldRefreshActivityViaPauseForCameraCompat() {
         return mCameraCompatEnableRefreshViaPauseOptProp.shouldEnableWithOverrideAndProperty(
-                isCompatChangeEnabled(OVERRIDE_CAMERA_COMPAT_ENABLE_REFRESH_VIA_PAUSE));
+                isChangeEnabled(mActivityRecord, OVERRIDE_CAMERA_COMPAT_ENABLE_REFRESH_VIA_PAUSE));
     }
 
     /**
@@ -150,7 +152,7 @@
      */
     boolean shouldForceRotateForCameraCompat() {
         return mCameraCompatAllowForceRotationOptProp.shouldEnableWithOptOutOverrideAndProperty(
-                isCompatChangeEnabled(OVERRIDE_CAMERA_COMPAT_DISABLE_FORCE_ROTATION));
+                isChangeEnabled(mActivityRecord, OVERRIDE_CAMERA_COMPAT_DISABLE_FORCE_ROTATION));
     }
 
     /**
@@ -168,7 +170,7 @@
      * </ul>
      */
     boolean shouldApplyFreeformTreatmentForCameraCompat() {
-        return Flags.cameraCompatForFreeform() && !isCompatChangeEnabled(
+        return Flags.cameraCompatForFreeform() && !isChangeEnabled(mActivityRecord,
                 OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT);
     }
 
@@ -191,7 +193,7 @@
     }
 
     boolean isOverrideOrientationOnlyForCameraEnabled() {
-        return isCompatChangeEnabled(OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA);
+        return isChangeEnabled(mActivityRecord, OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA);
     }
 
     /**
@@ -227,10 +229,6 @@
         mAppCompatCameraOverridesState.mFreeformCameraCompatMode = freeformCameraCompatMode;
     }
 
-    private boolean isCompatChangeEnabled(long overrideChangeId) {
-        return mActivityRecord.info.isChangeEnabled(overrideChangeId);
-    }
-
     static class AppCompatCameraOverridesState {
         // Whether activity "refresh" was requested but not finished in
         // ActivityRecord#activityResumedLocked following the camera compat force rotation in
diff --git a/services/core/java/com/android/server/wm/AppCompatController.java b/services/core/java/com/android/server/wm/AppCompatController.java
index 998d65d..54223b6 100644
--- a/services/core/java/com/android/server/wm/AppCompatController.java
+++ b/services/core/java/com/android/server/wm/AppCompatController.java
@@ -87,6 +87,11 @@
         return mAppCompatOverrides.getAppCompatAspectRatioOverrides();
     }
 
+    @NonNull
+    AppCompatResizeOverrides getAppCompatResizeOverrides() {
+        return mAppCompatOverrides.getAppCompatResizeOverrides();
+    }
+
     @Nullable
     AppCompatCameraPolicy getAppCompatCameraPolicy() {
         if (mActivityRecord.mDisplayContent != null) {
diff --git a/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java b/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java
index 0adf825..bd01351 100644
--- a/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java
@@ -20,14 +20,20 @@
 import static android.content.pm.ActivityInfo.OVERRIDE_ENABLE_COMPAT_IGNORE_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED;
 import static android.content.pm.ActivityInfo.OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION;
 import static android.content.pm.ActivityInfo.OVERRIDE_LANDSCAPE_ORIENTATION_TO_REVERSE_LANDSCAPE;
+import static android.content.pm.ActivityInfo.OVERRIDE_RESPECT_REQUESTED_ORIENTATION;
 import static android.content.pm.ActivityInfo.OVERRIDE_UNDEFINED_ORIENTATION_TO_NOSENSOR;
 import static android.content.pm.ActivityInfo.OVERRIDE_UNDEFINED_ORIENTATION_TO_PORTRAIT;
+import static android.content.pm.ActivityInfo.OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION;
+import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
+import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE;
 import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_IGNORING_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED;
+import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_ORIENTATION_OVERRIDE;
 import static android.view.WindowManager.PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION;
 
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_ATM;
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.AppCompatUtils.asLazy;
+import static com.android.server.wm.AppCompatUtils.isChangeEnabled;
 
 import android.annotation.NonNull;
 
@@ -54,6 +60,10 @@
     private final OptPropFactory.OptProp mIgnoreRequestedOrientationOptProp;
     @NonNull
     private final OptPropFactory.OptProp mAllowIgnoringOrientationRequestWhenLoopDetectedOptProp;
+    @NonNull
+    private final OptPropFactory.OptProp mAllowOrientationOverrideOptProp;
+    @NonNull
+    private final OptPropFactory.OptProp mAllowDisplayOrientationOverrideOptProp;
 
     @NonNull
     final OrientationOverridesState mOrientationOverridesState;
@@ -74,6 +84,17 @@
         mAllowIgnoringOrientationRequestWhenLoopDetectedOptProp = optPropBuilder.create(
                 PROPERTY_COMPAT_ALLOW_IGNORING_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED,
                 isPolicyForIgnoringRequestedOrientationEnabled);
+        mAllowOrientationOverrideOptProp = optPropBuilder.create(
+                PROPERTY_COMPAT_ALLOW_ORIENTATION_OVERRIDE);
+        mAllowDisplayOrientationOverrideOptProp = optPropBuilder.create(
+                PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE,
+                () -> mActivityRecord.mDisplayContent != null
+                        && mActivityRecord.getTask() != null
+                        && mActivityRecord.mDisplayContent.getIgnoreOrientationRequest()
+                        && !mActivityRecord.getTask().inMultiWindowMode()
+                        && mActivityRecord.mDisplayContent.getNaturalOrientation()
+                            == ORIENTATION_LANDSCAPE
+        );
     }
 
     boolean shouldEnableIgnoreOrientationRequest() {
@@ -81,6 +102,10 @@
                 isCompatChangeEnabled(OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION));
     }
 
+    boolean isOverrideRespectRequestedOrientationEnabled() {
+        return isChangeEnabled(mActivityRecord, OVERRIDE_RESPECT_REQUESTED_ORIENTATION);
+    }
+
     /**
      * Whether an app is calling {@link android.app.Activity#setRequestedOrientation}
      * in a loop and orientation request should be ignored.
@@ -113,6 +138,26 @@
     }
 
     /**
+     * Whether should fix display orientation to landscape natural orientation when a task is
+     * fullscreen and the display is ignoring orientation requests.
+     *
+     * <p>This treatment is enabled when the following conditions are met:
+     * <ul>
+     *     <li>Opt-out component property isn't enabled
+     *     <li>Opt-in per-app override is enabled
+     *     <li>Task is in fullscreen.
+     *     <li>{@link DisplayContent#getIgnoreOrientationRequest} is enabled
+     *     <li>Natural orientation of the display is landscape.
+     * </ul>
+     */
+    boolean shouldUseDisplayLandscapeNaturalOrientation() {
+        return mAllowDisplayOrientationOverrideOptProp
+                .shouldEnableWithOptInOverrideAndOptOutProperty(
+                        isChangeEnabled(mActivityRecord,
+                                OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION));
+    }
+
+    /**
      * Sets whether an activity is relaunching after the app has called {@link
      * android.app.Activity#setRequestedOrientation}.
      */
@@ -125,6 +170,10 @@
         return mOrientationOverridesState.mIsRelaunchingAfterRequestedOrientationChanged;
     }
 
+    boolean isAllowOrientationOverrideOptOut() {
+        return mAllowOrientationOverrideOptProp.isFalse();
+    }
+
     @VisibleForTesting
     int getSetOrientationRequestCounter() {
         return mOrientationOverridesState.mSetOrientationRequestCounter;
diff --git a/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java b/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java
index 17f0d97..c5506de 100644
--- a/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java
+++ b/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java
@@ -86,7 +86,8 @@
             return SCREEN_ORIENTATION_PORTRAIT;
         }
 
-        if (mAppCompatOverrides.isAllowOrientationOverrideOptOut()) {
+        if (mAppCompatOverrides.getAppCompatOrientationOverrides()
+                .isAllowOrientationOverrideOptOut()) {
             return candidate;
         }
 
diff --git a/services/core/java/com/android/server/wm/AppCompatOverrides.java b/services/core/java/com/android/server/wm/AppCompatOverrides.java
index cde48d6..4450011 100644
--- a/services/core/java/com/android/server/wm/AppCompatOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatOverrides.java
@@ -16,15 +16,6 @@
 
 package com.android.server.wm;
 
-import static android.content.pm.ActivityInfo.FORCE_NON_RESIZE_APP;
-import static android.content.pm.ActivityInfo.FORCE_RESIZE_APP;
-import static android.content.pm.ActivityInfo.OVERRIDE_RESPECT_REQUESTED_ORIENTATION;
-import static android.content.pm.ActivityInfo.OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION;
-import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
-import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE;
-import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_ORIENTATION_OVERRIDE;
-import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES;
-
 import android.annotation.NonNull;
 
 import com.android.server.wm.utils.OptPropFactory;
@@ -35,15 +26,6 @@
 public class AppCompatOverrides {
 
     @NonNull
-    private final ActivityRecord mActivityRecord;
-
-    @NonNull
-    private final OptPropFactory.OptProp mAllowOrientationOverrideOptProp;
-    @NonNull
-    private final OptPropFactory.OptProp mAllowDisplayOrientationOverrideOptProp;
-    @NonNull
-    private final OptPropFactory.OptProp mAllowForceResizeOverrideOptProp;
-    @NonNull
     private final AppCompatOrientationOverrides mAppCompatOrientationOverrides;
     @NonNull
     private final AppCompatCameraOverrides mAppCompatCameraOverrides;
@@ -51,39 +33,24 @@
     private final AppCompatAspectRatioOverrides mAppCompatAspectRatioOverrides;
     @NonNull
     private final AppCompatFocusOverrides mAppCompatFocusOverrides;
+    @NonNull
+    private final AppCompatResizeOverrides mAppCompatResizeOverrides;
 
     AppCompatOverrides(@NonNull ActivityRecord activityRecord,
             @NonNull AppCompatConfiguration appCompatConfiguration,
             @NonNull OptPropFactory optPropBuilder) {
-        mActivityRecord = activityRecord;
-
-        mAppCompatCameraOverrides = new AppCompatCameraOverrides(mActivityRecord,
+        mAppCompatCameraOverrides = new AppCompatCameraOverrides(activityRecord,
                 appCompatConfiguration, optPropBuilder);
-        mAppCompatOrientationOverrides = new AppCompatOrientationOverrides(mActivityRecord,
+        mAppCompatOrientationOverrides = new AppCompatOrientationOverrides(activityRecord,
                 appCompatConfiguration, optPropBuilder, mAppCompatCameraOverrides);
         // TODO(b/341903757) Remove BooleanSuppliers after fixing dependency with reachability.
         mAppCompatAspectRatioOverrides = new AppCompatAspectRatioOverrides(activityRecord,
                 appCompatConfiguration, optPropBuilder,
                 activityRecord.mLetterboxUiController::isDisplayFullScreenAndInPosture,
                 activityRecord.mLetterboxUiController::getHorizontalPositionMultiplier);
-        mAppCompatFocusOverrides = new AppCompatFocusOverrides(mActivityRecord,
+        mAppCompatFocusOverrides = new AppCompatFocusOverrides(activityRecord,
                 appCompatConfiguration, optPropBuilder);
-
-        mAllowOrientationOverrideOptProp = optPropBuilder.create(
-                PROPERTY_COMPAT_ALLOW_ORIENTATION_OVERRIDE);
-
-        mAllowDisplayOrientationOverrideOptProp = optPropBuilder.create(
-                PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE,
-                () -> mActivityRecord.mDisplayContent != null
-                        && mActivityRecord.getTask() != null
-                        && mActivityRecord.mDisplayContent.getIgnoreOrientationRequest()
-                        && !mActivityRecord.getTask().inMultiWindowMode()
-                        && mActivityRecord.mDisplayContent.getNaturalOrientation()
-                            == ORIENTATION_LANDSCAPE
-        );
-
-        mAllowForceResizeOverrideOptProp = optPropBuilder.create(
-                PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+        mAppCompatResizeOverrides = new AppCompatResizeOverrides(activityRecord, optPropBuilder);
     }
 
     @NonNull
@@ -106,66 +73,8 @@
         return mAppCompatFocusOverrides;
     }
 
-    boolean isAllowOrientationOverrideOptOut() {
-        return mAllowOrientationOverrideOptProp.isFalse();
-    }
-
-    boolean isOverrideRespectRequestedOrientationEnabled() {
-        return isCompatChangeEnabled(OVERRIDE_RESPECT_REQUESTED_ORIENTATION);
-    }
-
-    /**
-     * Whether should fix display orientation to landscape natural orientation when a task is
-     * fullscreen and the display is ignoring orientation requests.
-     *
-     * <p>This treatment is enabled when the following conditions are met:
-     * <ul>
-     *     <li>Opt-out component property isn't enabled
-     *     <li>Opt-in per-app override is enabled
-     *     <li>Task is in fullscreen.
-     *     <li>{@link DisplayContent#getIgnoreOrientationRequest} is enabled
-     *     <li>Natural orientation of the display is landscape.
-     * </ul>
-     */
-    boolean shouldUseDisplayLandscapeNaturalOrientation() {
-        return mAllowDisplayOrientationOverrideOptProp
-                .shouldEnableWithOptInOverrideAndOptOutProperty(
-                        isCompatChangeEnabled(OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION));
-    }
-
-    /**
-     * Whether we should apply the force resize per-app override. When this override is applied it
-     * forces the packages it is applied to to be resizable. It won't change whether the app can be
-     * put into multi-windowing mode, but allow the app to resize without going into size-compat
-     * mode when the window container resizes, such as display size change or screen rotation.
-     *
-     * <p>This method returns {@code true} when the following conditions are met:
-     * <ul>
-     *     <li>Opt-out component property isn't enabled
-     *     <li>Per-app override is enabled
-     * </ul>
-     */
-    boolean shouldOverrideForceResizeApp() {
-        return mAllowForceResizeOverrideOptProp.shouldEnableWithOptInOverrideAndOptOutProperty(
-                isCompatChangeEnabled(FORCE_RESIZE_APP));
-    }
-
-    /**
-     * Whether we should apply the force non resize per-app override. When this override is applied
-     * it forces the packages it is applied to to be non-resizable.
-     *
-     * <p>This method returns {@code true} when the following conditions are met:
-     * <ul>
-     *     <li>Opt-out component property isn't enabled
-     *     <li>Per-app override is enabled
-     * </ul>
-     */
-    boolean shouldOverrideForceNonResizeApp() {
-        return mAllowForceResizeOverrideOptProp.shouldEnableWithOptInOverrideAndOptOutProperty(
-                isCompatChangeEnabled(FORCE_NON_RESIZE_APP));
-    }
-
-    private boolean isCompatChangeEnabled(long overrideChangeId) {
-        return mActivityRecord.info.isChangeEnabled(overrideChangeId);
+    @NonNull
+    AppCompatResizeOverrides getAppCompatResizeOverrides() {
+        return mAppCompatResizeOverrides;
     }
 }
diff --git a/services/core/java/com/android/server/wm/AppCompatResizeOverrides.java b/services/core/java/com/android/server/wm/AppCompatResizeOverrides.java
new file mode 100644
index 0000000..60c1825
--- /dev/null
+++ b/services/core/java/com/android/server/wm/AppCompatResizeOverrides.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright (C) 2024 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.server.wm;
+
+import static android.content.pm.ActivityInfo.FORCE_NON_RESIZE_APP;
+import static android.content.pm.ActivityInfo.FORCE_RESIZE_APP;
+import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES;
+
+import static com.android.server.wm.AppCompatUtils.isChangeEnabled;
+
+import android.annotation.NonNull;
+
+import com.android.server.wm.utils.OptPropFactory;
+
+/**
+ * Encapsulate app compat logic about resizability.
+ */
+class AppCompatResizeOverrides {
+
+    @NonNull
+    private final ActivityRecord mActivityRecord;
+
+    @NonNull
+    private final OptPropFactory.OptProp mAllowForceResizeOverrideOptProp;
+
+    AppCompatResizeOverrides(@NonNull ActivityRecord activityRecord,
+            @NonNull OptPropFactory optPropBuilder) {
+        mActivityRecord = activityRecord;
+        mAllowForceResizeOverrideOptProp = optPropBuilder.create(
+                PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+    }
+
+    /**
+     * Whether we should apply the force resize per-app override. When this override is applied it
+     * forces the packages it is applied to to be resizable. It won't change whether the app can be
+     * put into multi-windowing mode, but allow the app to resize without going into size-compat
+     * mode when the window container resizes, such as display size change or screen rotation.
+     *
+     * <p>This method returns {@code true} when the following conditions are met:
+     * <ul>
+     *     <li>Opt-out component property isn't enabled
+     *     <li>Per-app override is enabled
+     * </ul>
+     */
+    boolean shouldOverrideForceResizeApp() {
+        return mAllowForceResizeOverrideOptProp.shouldEnableWithOptInOverrideAndOptOutProperty(
+                isChangeEnabled(mActivityRecord, FORCE_RESIZE_APP));
+    }
+
+    /**
+     * Whether we should apply the force non resize per-app override. When this override is applied
+     * it forces the packages it is applied to to be non-resizable.
+     *
+     * <p>This method returns {@code true} when the following conditions are met:
+     * <ul>
+     *     <li>Opt-out component property isn't enabled
+     *     <li>Per-app override is enabled
+     * </ul>
+     */
+    boolean shouldOverrideForceNonResizeApp() {
+        return mAllowForceResizeOverrideOptProp.shouldEnableWithOptInOverrideAndOptOutProperty(
+                isChangeEnabled(mActivityRecord, FORCE_NON_RESIZE_APP));
+    }
+}
diff --git a/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java b/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java
index cbb210f..9996bbc 100644
--- a/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java
+++ b/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java
@@ -27,7 +27,6 @@
 import static com.android.server.wm.LaunchParamsUtil.applyLayoutGravity;
 import static com.android.server.wm.LaunchParamsUtil.calculateLayoutBounds;
 
-
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.ActivityOptions;
diff --git a/services/core/java/com/android/server/wm/DisplayArea.java b/services/core/java/com/android/server/wm/DisplayArea.java
index 75724eb..86f69cd 100644
--- a/services/core/java/com/android/server/wm/DisplayArea.java
+++ b/services/core/java/com/android/server/wm/DisplayArea.java
@@ -267,7 +267,8 @@
                 // between fullscreen and PiP would work well. Checking TaskFragment rather than
                 // Task to ensure that Activity Embedding is excluded.
                 && activity.getTaskFragment().getWindowingMode() == WINDOWING_MODE_FULLSCREEN
-                && activity.mLetterboxUiController.isOverrideRespectRequestedOrientationEnabled();
+                && activity.mAppCompatController.getAppCompatOrientationOverrides()
+                    .isOverrideRespectRequestedOrientationEnabled();
     }
 
     boolean getIgnoreOrientationRequest() {
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 475b473..403c307 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -2939,8 +2939,9 @@
 
         if (!handlesOrientationChangeFromDescendant(orientation)) {
             ActivityRecord topActivity = topRunningActivity(/* considerKeyguardState= */ true);
-            if (topActivity != null && topActivity.mLetterboxUiController
-                    .shouldUseDisplayLandscapeNaturalOrientation()) {
+            if (topActivity != null && topActivity.mAppCompatController
+                    .getAppCompatOrientationOverrides()
+                        .shouldUseDisplayLandscapeNaturalOrientation()) {
                 ProtoLog.v(WM_DEBUG_ORIENTATION,
                         "Display id=%d is ignoring orientation request for %d, return %d"
                         + " following a per-app override for %s",
diff --git a/services/core/java/com/android/server/wm/DragDropController.java b/services/core/java/com/android/server/wm/DragDropController.java
index 6abef8b..c849a37 100644
--- a/services/core/java/com/android/server/wm/DragDropController.java
+++ b/services/core/java/com/android/server/wm/DragDropController.java
@@ -16,7 +16,6 @@
 
 package com.android.server.wm;
 
-import static android.content.ClipDescription.EXTRA_HIDE_DRAG_SOURCE_TASK_ID;
 import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.view.View.DRAG_FLAG_GLOBAL;
 import static android.view.View.DRAG_FLAG_GLOBAL_SAME_APPLICATION;
@@ -228,7 +227,7 @@
                         final Display display = displayContent.getDisplay();
                         touchFocusTransferredFuture = mCallback.get().registerInputChannel(
                                 mDragState, display, mService.mInputManager,
-                                callingWin.mInputChannel);
+                                callingWin.mInputChannelToken);
                     } else {
                         // Skip surface logic for a drag triggered by an AccessibilityAction
                         mDragState.broadcastDragStartedLocked(touchX, touchY);
diff --git a/services/core/java/com/android/server/wm/DragState.java b/services/core/java/com/android/server/wm/DragState.java
index ba74f50..59435b8 100644
--- a/services/core/java/com/android/server/wm/DragState.java
+++ b/services/core/java/com/android/server/wm/DragState.java
@@ -456,10 +456,6 @@
         }
     }
 
-    InputChannel getInputChannel() {
-        return mInputInterceptor == null ? null : mInputInterceptor.mClientChannel;
-    }
-
     InputWindowHandle getInputWindowHandle() {
         return mInputInterceptor == null ? null : mInputInterceptor.mDragWindowHandle;
     }
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index 872b4e1..7a0fd3e 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -60,7 +60,6 @@
 import com.android.internal.policy.IKeyguardDismissCallback;
 import com.android.server.inputmethod.InputMethodManagerInternal;
 import com.android.server.policy.WindowManagerPolicy;
-import com.android.window.flags.Flags;
 
 import java.io.PrintWriter;
 
@@ -198,14 +197,6 @@
             setWakeTransitionReady();
             return;
         }
-        EventLogTags.writeWmSetKeyguardShown(
-                displayId,
-                keyguardShowing ? 1 : 0,
-                aodShowing ? 1 : 0,
-                state.mKeyguardGoingAway ? 1 : 0,
-                state.mOccluded ? 1 : 0,
-                "setKeyguardShown");
-
         // Update the task snapshot if the screen will not be turned off. To make sure that the
         // unlocking animation can animate consistent content. The conditions are:
         // - Either AOD or keyguard changes to be showing. So if the states change individually,
@@ -224,6 +215,7 @@
 
         state.mKeyguardShowing = keyguardShowing;
         state.mAodShowing = aodShowing;
+        state.writeEventLog("setKeyguardShown");
 
         if (keyguardChanged) {
             // Irrelevant to AOD.
@@ -232,19 +224,13 @@
                 state.mDismissalRequested = false;
             }
             if (goingAwayRemoved
-                    || (Flags.keyguardAppearTransition() && keyguardShowing
-                            && !Display.isOffState(dc.getDisplayInfo().state))) {
+                    || (keyguardShowing && !Display.isOffState(dc.getDisplayInfo().state))) {
                 // Keyguard decided to show or stopped going away. Send a transition to animate back
                 // to the locked state before holding the sleep token again
-                final DisplayContent transitionDc = Flags.keyguardAppearTransition()
-                        ? dc
-                        : mRootWindowContainer.getDefaultDisplay();
-                transitionDc.requestTransitionAndLegacyPrepare(
+                dc.requestTransitionAndLegacyPrepare(
                         TRANSIT_TO_FRONT, TRANSIT_FLAG_KEYGUARD_APPEARING);
-                if (Flags.keyguardAppearTransition()) {
-                    dc.mWallpaperController.adjustWallpaperWindows();
-                }
-                transitionDc.executeAppTransition();
+                dc.mWallpaperController.adjustWallpaperWindows();
+                dc.executeAppTransition();
             }
         }
 
@@ -284,13 +270,7 @@
         mService.deferWindowLayout();
         state.mKeyguardGoingAway = true;
         try {
-            EventLogTags.writeWmSetKeyguardShown(
-                    displayId,
-                    state.mKeyguardShowing ? 1 : 0,
-                    state.mAodShowing ? 1 : 0,
-                    1 /* keyguardGoingAway */,
-                    state.mOccluded ? 1 : 0,
-                    "keyguardGoingAway");
+            state.writeEventLog("keyguardGoingAway");
             final int transitFlags = convertTransitFlags(flags);
             final DisplayContent dc = mRootWindowContainer.getDefaultDisplay();
             dc.prepareAppTransition(TRANSIT_KEYGUARD_GOING_AWAY, transitFlags);
@@ -436,8 +416,9 @@
         }
 
         final TransitionController tc = mRootWindowContainer.mTransitionController;
+        final KeyguardDisplayState state = getDisplayState(displayId);
 
-        final boolean occluded = getDisplayState(displayId).mOccluded;
+        final boolean occluded = state.mOccluded;
         final boolean performTransition = isKeyguardLocked(displayId);
         final boolean executeTransition = performTransition && !tc.isCollecting();
 
@@ -481,7 +462,7 @@
     /**
      * Called when keyguard going away state changed.
      */
-    private void handleKeyguardGoingAwayChanged(DisplayContent dc) {
+    private void handleDismissInsecureKeyguard(DisplayContent dc) {
         mService.deferWindowLayout();
         try {
             dc.prepareAppTransition(TRANSIT_KEYGUARD_GOING_AWAY, 0 /* transitFlags */);
@@ -646,6 +627,16 @@
             mSleepTokenAcquirer.release(mDisplayId);
         }
 
+        void writeEventLog(String reason) {
+            EventLogTags.writeWmSetKeyguardShown(
+                    mDisplayId,
+                    mKeyguardShowing ? 1 : 0,
+                    mAodShowing ? 1 : 0,
+                    mKeyguardGoingAway ? 1 : 0,
+                    mOccluded ? 1 : 0,
+                    reason);
+        }
+
         /**
          * Updates keyguard status if the top task could be visible. The top task may occlude
          * keyguard, request to dismiss keyguard or make insecure keyguard go away based on its
@@ -715,23 +706,16 @@
             }
 
             boolean hasChange = false;
-            if (lastOccluded != mOccluded) {
-                if (mDisplayId == DEFAULT_DISPLAY) {
-                    EventLogTags.writeWmSetKeyguardShown(
-                            mDisplayId,
-                            mKeyguardShowing ? 1 : 0,
-                            mAodShowing ? 1 : 0,
-                            mKeyguardGoingAway ? 1 : 0,
-                            mOccluded ? 1 : 0,
-                            "updateVisibility");
-                }
+            if (!lastKeyguardGoingAway && mKeyguardGoingAway) {
+                writeEventLog("dismissIfInsecure");
+                controller.handleDismissInsecureKeyguard(display);
+                hasChange = true;
+            } else if (lastOccluded != mOccluded) {
                 controller.handleOccludedChanged(mDisplayId, mTopOccludesActivity);
                 hasChange = true;
-            } else if (!lastKeyguardGoingAway && mKeyguardGoingAway) {
-                controller.handleKeyguardGoingAwayChanged(display);
-                hasChange = true;
             }
-            // Collect the participates for shell transition, so that transition won't happen too
+
+            // Collect the participants for shell transition, so that transition won't happen too
             // early since the transition was set ready.
             if (hasChange && top != null && (mOccluded || mKeyguardGoingAway)) {
                 display.mTransitionController.collect(top);
diff --git a/services/core/java/com/android/server/wm/LetterboxUiController.java b/services/core/java/com/android/server/wm/LetterboxUiController.java
index 73f3655..291eab1 100644
--- a/services/core/java/com/android/server/wm/LetterboxUiController.java
+++ b/services/core/java/com/android/server/wm/LetterboxUiController.java
@@ -116,67 +116,6 @@
         }
     }
 
-    /**
-     * Whether we should apply the force resize per-app override. When this override is applied it
-     * forces the packages it is applied to to be resizable. It won't change whether the app can be
-     * put into multi-windowing mode, but allow the app to resize without going into size-compat
-     * mode when the window container resizes, such as display size change or screen rotation.
-     *
-     * <p>This method returns {@code true} when the following conditions are met:
-     * <ul>
-     *     <li>Opt-out component property isn't enabled
-     *     <li>Per-app override is enabled
-     * </ul>
-     */
-    boolean shouldOverrideForceResizeApp() {
-        return getAppCompatOverrides().shouldOverrideForceResizeApp();
-    }
-
-    /**
-     * Whether we should apply the force non resize per-app override. When this override is applied
-     * it forces the packages it is applied to to be non-resizable.
-     *
-     * <p>This method returns {@code true} when the following conditions are met:
-     * <ul>
-     *     <li>Opt-out component property isn't enabled
-     *     <li>Per-app override is enabled
-     * </ul>
-     */
-    boolean shouldOverrideForceNonResizeApp() {
-        return getAppCompatOverrides().shouldOverrideForceNonResizeApp();
-    }
-
-    /**
-     * Sets whether an activity is relaunching after the app has called {@link
-     * android.app.Activity#setRequestedOrientation}.
-     */
-    void setRelaunchingAfterRequestedOrientationChanged(boolean isRelaunching) {
-        getAppCompatOverrides().getAppCompatOrientationOverrides()
-                .setRelaunchingAfterRequestedOrientationChanged(isRelaunching);
-    }
-
-
-    boolean isOverrideRespectRequestedOrientationEnabled() {
-        return getAppCompatOverrides().isOverrideRespectRequestedOrientationEnabled();
-    }
-
-    /**
-     * Whether should fix display orientation to landscape natural orientation when a task is
-     * fullscreen and the display is ignoring orientation requests.
-     *
-     * <p>This treatment is enabled when the following conditions are met:
-     * <ul>
-     *     <li>Opt-out component property isn't enabled
-     *     <li>Opt-in per-app override is enabled
-     *     <li>Task is in fullscreen.
-     *     <li>{@link DisplayContent#getIgnoreOrientationRequest} is enabled
-     *     <li>Natural orientation of the display is landscape.
-     * </ul>
-     */
-    boolean shouldUseDisplayLandscapeNaturalOrientation() {
-        return getAppCompatOverrides().shouldUseDisplayLandscapeNaturalOrientation();
-    }
-
     boolean hasWallpaperBackgroundForLetterbox() {
         return mShowWallpaperForLetterboxBackground;
     }
@@ -386,11 +325,6 @@
                 : mAppCompatConfiguration.getLetterboxVerticalPositionMultiplier(tabletopMode);
     }
 
-    float getFixedOrientationLetterboxAspectRatio(@NonNull Configuration parentConfiguration) {
-        return mActivityRecord.mAppCompatController.getAppCompatAspectRatioOverrides()
-                .getFixedOrientationLetterboxAspectRatio(parentConfiguration);
-    }
-
     boolean isLetterboxEducationEnabled() {
         return mAppCompatConfiguration.getIsEducationEnabled();
     }
@@ -800,11 +734,6 @@
         return null;
     }
 
-    boolean getIsRelaunchingAfterRequestedOrientationChanged() {
-        return getAppCompatOverrides().getAppCompatOrientationOverrides()
-                .getIsRelaunchingAfterRequestedOrientationChanged();
-    }
-
     private void adjustBoundsForTaskbar(final WindowState mainWindow, final Rect bounds) {
         // Rounded corners should be displayed above the taskbar. When taskbar is hidden,
         // an insets frame is equal to a navigation bar which shouldn't affect position of
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index f9772f4..1c9aaf9 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -1132,6 +1132,9 @@
             final Task oldParentTask = oldParent.asTask();
             if (oldParentTask != null) {
                 forAllActivities(oldParentTask::cleanUpActivityReferences);
+
+                // Update the task description of the previous parent as well
+                oldParentTask.updateTaskDescription();
             }
 
             if (newParent == null || !newParent.inPinnedWindowingMode()) {
@@ -1163,6 +1166,9 @@
                 } catch (RemoteException e) {
                 }
             }
+
+            // Update the ancestor tasks' task description after reparenting
+            updateTaskDescription();
         }
 
         // First time we are adding the task to the system.
@@ -2281,7 +2287,7 @@
         // Apply crop to root tasks only and clear the crops of the descendant tasks.
         int width = 0;
         int height = 0;
-        if (isRootTask()) {
+        if (isRootTask() && !mTransitionController.mIsWaitingForDisplayEnabled) {
             final Rect taskBounds = getBounds();
             width = taskBounds.width();
             height = taskBounds.height();
@@ -3353,7 +3359,7 @@
 
         //TODO (AM refactor): Just use local once updateEffectiveIntent is run during all child
         //                    order changes.
-        final Task topTask = top != null ? top.getTask() : this;
+        final Task topTask = top != null && top.getTask() != null ? top.getTask() : this;
         info.resizeMode = topTask.mResizeMode;
         info.topActivityType = topTask.getActivityType();
         info.displayCutoutInsets = topTask.getDisplayCutoutInsets();
@@ -3407,6 +3413,7 @@
         info.isSleeping = shouldSleepActivities();
         info.isTopActivityTransparent = top != null && !top.fillsParent();
         info.isTopActivityStyleFloating = top != null && top.isStyleFloating();
+        info.lastNonFullscreenBounds = topTask.mLastNonFullscreenBounds;
         appCompatTaskInfo.topActivityLetterboxVerticalPosition = TaskInfo.PROPERTY_VALUE_UNSET;
         appCompatTaskInfo.topActivityLetterboxHorizontalPosition = TaskInfo.PROPERTY_VALUE_UNSET;
         appCompatTaskInfo.topActivityLetterboxWidth = TaskInfo.PROPERTY_VALUE_UNSET;
@@ -6138,9 +6145,8 @@
 
     @Override
     void onChildPositionChanged(WindowContainer child) {
-        dispatchTaskInfoChangedIfNeeded(false /* force */);
-
         if (!mChildren.contains(child)) {
+            dispatchTaskInfoChangedIfNeeded(false /* force */);
             return;
         }
         if (child.asTask() != null) {
@@ -6152,6 +6158,10 @@
             // Send for TaskFragmentParentInfo#hasDirectActivity change.
             sendTaskFragmentParentInfoChangedIfNeeded();
         }
+
+        // Update the ancestor tasks' task description after any children have reparented
+        updateTaskDescription();
+        dispatchTaskInfoChangedIfNeeded(false /* force */);
     }
 
     void reparent(TaskDisplayArea newParent, boolean onTop) {
diff --git a/services/core/java/com/android/server/wm/WindowAnimationSpec.java b/services/core/java/com/android/server/wm/WindowAnimationSpec.java
index 34b9913..2c58c61 100644
--- a/services/core/java/com/android/server/wm/WindowAnimationSpec.java
+++ b/services/core/java/com/android/server/wm/WindowAnimationSpec.java
@@ -97,10 +97,10 @@
 
     /**
      * @return If a window animation has outsets applied to it.
-     * @see Animation#hasExtension()
+     * @see Animation#getExtensionEdges()
      */
     public boolean hasExtension() {
-        return mAnimation.hasExtension();
+        return mAnimation.getExtensionEdges() != 0;
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/WindowManagerInternal.java b/services/core/java/com/android/server/wm/WindowManagerInternal.java
index 8ae1cf0..48a5050 100644
--- a/services/core/java/com/android/server/wm/WindowManagerInternal.java
+++ b/services/core/java/com/android/server/wm/WindowManagerInternal.java
@@ -41,7 +41,6 @@
 import android.view.IInputFilter;
 import android.view.IRemoteAnimationFinishedCallback;
 import android.view.IWindow;
-import android.view.InputChannel;
 import android.view.MagnificationSpec;
 import android.view.RemoteAnimationTarget;
 import android.view.Surface;
@@ -377,10 +376,10 @@
     public interface IDragDropCallback {
         default CompletableFuture<Boolean> registerInputChannel(
                 DragState state, Display display, InputManagerService service,
-                InputChannel source) {
+                IBinder sourceInputChannelToken) {
             return state.register(display)
                 .thenApply(unused ->
-                    service.startDragAndDrop(source, state.getInputChannel()));
+                    service.startDragAndDrop(sourceInputChannelToken, state.getInputToken()));
         }
 
         /**
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index a218068..4db62478 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -9778,7 +9778,7 @@
                 Slog.e(TAG, "Host window not found");
                 return;
             }
-            if (hostWindow.mInputChannel == null) {
+            if (hostWindow.mInputChannelToken == null) {
                 Slog.e(TAG, "Host window does not have an input channel");
                 return;
             }
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 164994c..153d41b 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -638,7 +638,7 @@
     /**
      * Only populated if flag REMOVE_INPUT_CHANNEL_FROM_WINDOWSTATE is disabled.
      */
-    InputChannel mInputChannel;
+    private InputChannel mInputChannel;
 
     /**
      * The token will be assigned to {@link InputWindowHandle#token} if this window can receive
@@ -5373,7 +5373,7 @@
             // change then delay the position update until it has redrawn to avoid any flickers.
             final boolean isLetterboxedAndRelaunching = activityRecord != null
                     && activityRecord.areBoundsLetterboxed()
-                    && activityRecord.mLetterboxUiController
+                    && activityRecord.mAppCompatController.getAppCompatOrientationOverrides()
                         .getIsRelaunchingAfterRequestedOrientationChanged();
             if (surfaceResizedWithoutMoveAnimation || isLetterboxedAndRelaunching) {
                 applyWithNextDraw(mSetSurfacePositionConsumer);
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java
index 9e46f2f..efcc23f 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/DefaultImeVisibilityApplierTest.java
@@ -70,14 +70,11 @@
 public class DefaultImeVisibilityApplierTest extends InputMethodManagerServiceTestBase {
     private DefaultImeVisibilityApplier mVisibilityApplier;
 
-    private int mUserId = 0;
-
     @Before
     public void setUp() throws RemoteException {
         super.setUp();
         synchronized (ImfLock.class) {
             mVisibilityApplier = mInputMethodManagerService.getVisibilityApplierLocked();
-            mUserId = mInputMethodManagerService.getCurrentImeUserIdLocked();
             mInputMethodManagerService.setAttachedClientForTesting(requireNonNull(
                     mInputMethodManagerService.getClientStateLocked(mMockInputMethodClient)));
         }
@@ -248,7 +245,7 @@
                 mMockRemoteInputConnection /* inputConnection */,
                 mMockRemoteAccessibilityInputConnection /* remoteAccessibilityInputConnection */,
                 mTargetSdkVersion /* unverifiedTargetSdkVersion */,
-                mCallingUserId /* userId */,
+                mUserId /* userId */,
                 mMockImeOnBackInvokedDispatcher /* imeDispatcher */);
     }
 }
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImeVisibilityStateComputerTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImeVisibilityStateComputerTest.java
index 337d5c1..dd3b33e 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImeVisibilityStateComputerTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImeVisibilityStateComputerTest.java
@@ -49,6 +49,7 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.server.wm.ImeTargetChangeListener;
 import com.android.server.wm.WindowManagerInternal;
 
@@ -93,33 +94,37 @@
 
     @Test
     public void testRequestImeVisibility_showImplicit() {
-        initImeTargetWindowState(mWindowToken);
-        boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(),
-                InputMethodManager.SHOW_IMPLICIT);
-        mComputer.requestImeVisibility(mWindowToken, res);
+        synchronized (ImfLock.class) {
+            initImeTargetWindowState(mWindowToken);
+            boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(),
+                    InputMethodManager.SHOW_IMPLICIT);
+            mComputer.requestImeVisibility(mWindowToken, res);
 
-        final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
-        assertThat(state).isNotNull();
-        assertThat(state.hasEditorFocused()).isTrue();
-        assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
-        assertThat(state.isRequestedImeVisible()).isTrue();
+            final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
+            assertThat(state).isNotNull();
+            assertThat(state.hasEditorFocused()).isTrue();
+            assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
+            assertThat(state.isRequestedImeVisible()).isTrue();
 
-        assertThat(mComputer.mRequestedShowExplicitly).isFalse();
+            assertThat(mComputer.mRequestedShowExplicitly).isFalse();
+        }
     }
 
     @Test
     public void testRequestImeVisibility_showExplicit() {
-        initImeTargetWindowState(mWindowToken);
-        boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(), 0 /* showFlags */);
-        mComputer.requestImeVisibility(mWindowToken, res);
+        synchronized (ImfLock.class) {
+            initImeTargetWindowState(mWindowToken);
+            boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(), 0 /* showFlags */);
+            mComputer.requestImeVisibility(mWindowToken, res);
 
-        final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
-        assertThat(state).isNotNull();
-        assertThat(state.hasEditorFocused()).isTrue();
-        assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
-        assertThat(state.isRequestedImeVisible()).isTrue();
+            final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
+            assertThat(state).isNotNull();
+            assertThat(state.hasEditorFocused()).isTrue();
+            assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
+            assertThat(state.isRequestedImeVisible()).isTrue();
 
-        assertThat(mComputer.mRequestedShowExplicitly).isTrue();
+            assertThat(mComputer.mRequestedShowExplicitly).isTrue();
+        }
     }
 
     /**
@@ -128,12 +133,14 @@
      */
     @Test
     public void testRequestImeVisibility_showExplicit_thenShowImplicit() {
-        initImeTargetWindowState(mWindowToken);
-        mComputer.onImeShowFlags(ImeTracker.Token.empty(), 0 /* showFlags */);
-        assertThat(mComputer.mRequestedShowExplicitly).isTrue();
+        synchronized (ImfLock.class) {
+            initImeTargetWindowState(mWindowToken);
+            mComputer.onImeShowFlags(ImeTracker.Token.empty(), 0 /* showFlags */);
+            assertThat(mComputer.mRequestedShowExplicitly).isTrue();
 
-        mComputer.onImeShowFlags(null, InputMethodManager.SHOW_IMPLICIT);
-        assertThat(mComputer.mRequestedShowExplicitly).isTrue();
+            mComputer.onImeShowFlags(null, InputMethodManager.SHOW_IMPLICIT);
+            assertThat(mComputer.mRequestedShowExplicitly).isTrue();
+        }
     }
 
     /**
@@ -142,162 +149,181 @@
      */
     @Test
     public void testRequestImeVisibility_showForced_thenShowExplicit() {
-        initImeTargetWindowState(mWindowToken);
-        mComputer.onImeShowFlags(ImeTracker.Token.empty(), InputMethodManager.SHOW_FORCED);
-        assertThat(mComputer.mShowForced).isTrue();
+        synchronized (ImfLock.class) {
+            initImeTargetWindowState(mWindowToken);
+            mComputer.onImeShowFlags(ImeTracker.Token.empty(), InputMethodManager.SHOW_FORCED);
+            assertThat(mComputer.mShowForced).isTrue();
 
-        mComputer.onImeShowFlags(ImeTracker.Token.empty(), 0 /* showFlags */);
-        assertThat(mComputer.mShowForced).isTrue();
+            mComputer.onImeShowFlags(ImeTracker.Token.empty(), 0 /* showFlags */);
+            assertThat(mComputer.mShowForced).isTrue();
+        }
     }
 
     @Test
     public void testRequestImeVisibility_showImplicit_a11yNoImePolicy() {
-        // Precondition: set AccessibilityService#SHOW_MODE_HIDDEN policy
-        mComputer.getImePolicy().setA11yRequestNoSoftKeyboard(SHOW_MODE_HIDDEN);
+        synchronized (ImfLock.class) {
+            // Precondition: set AccessibilityService#SHOW_MODE_HIDDEN policy
+            mComputer.getImePolicy().setA11yRequestNoSoftKeyboard(SHOW_MODE_HIDDEN);
 
-        initImeTargetWindowState(mWindowToken);
-        boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(),
-                InputMethodManager.SHOW_IMPLICIT);
-        mComputer.requestImeVisibility(mWindowToken, res);
+            initImeTargetWindowState(mWindowToken);
+            boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(),
+                    InputMethodManager.SHOW_IMPLICIT);
+            mComputer.requestImeVisibility(mWindowToken, res);
 
-        final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
-        assertThat(state).isNotNull();
-        assertThat(state.hasEditorFocused()).isTrue();
-        assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
-        assertThat(state.isRequestedImeVisible()).isFalse();
+            final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
+            assertThat(state).isNotNull();
+            assertThat(state.hasEditorFocused()).isTrue();
+            assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
+            assertThat(state.isRequestedImeVisible()).isFalse();
 
-        assertThat(mComputer.mRequestedShowExplicitly).isFalse();
+            assertThat(mComputer.mRequestedShowExplicitly).isFalse();
+        }
     }
 
     @Test
     public void testRequestImeVisibility_showImplicit_imeHiddenPolicy() {
-        // Precondition: set IME hidden display policy before calling showSoftInput
-        mComputer.getImePolicy().setImeHiddenByDisplayPolicy(true);
+        synchronized (ImfLock.class) {
+            // Precondition: set IME hidden display policy before calling showSoftInput
+            mComputer.getImePolicy().setImeHiddenByDisplayPolicy(true);
 
-        initImeTargetWindowState(mWindowToken);
-        boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(),
-                InputMethodManager.SHOW_IMPLICIT);
-        mComputer.requestImeVisibility(mWindowToken, res);
+            initImeTargetWindowState(mWindowToken);
+            boolean res = mComputer.onImeShowFlags(ImeTracker.Token.empty(),
+                    InputMethodManager.SHOW_IMPLICIT);
+            mComputer.requestImeVisibility(mWindowToken, res);
 
-        final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
-        assertThat(state).isNotNull();
-        assertThat(state.hasEditorFocused()).isTrue();
-        assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
-        assertThat(state.isRequestedImeVisible()).isFalse();
+            final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
+            assertThat(state).isNotNull();
+            assertThat(state.hasEditorFocused()).isTrue();
+            assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
+            assertThat(state.isRequestedImeVisible()).isFalse();
 
-        assertThat(mComputer.mRequestedShowExplicitly).isFalse();
+            assertThat(mComputer.mRequestedShowExplicitly).isFalse();
+        }
     }
 
     @Test
     public void testRequestImeVisibility_hideNotAlways() {
-        // Precondition: ensure IME has shown before hiding request.
-        mComputer.setInputShown(true);
+        synchronized (ImfLock.class) {
+            // Precondition: ensure IME has shown before hiding request.
+            mComputer.setInputShown(true);
 
-        initImeTargetWindowState(mWindowToken);
-        assertThat(mComputer.canHideIme(ImeTracker.Token.empty(),
-                InputMethodManager.HIDE_NOT_ALWAYS)).isTrue();
-        mComputer.requestImeVisibility(mWindowToken, false);
+            initImeTargetWindowState(mWindowToken);
+            assertThat(mComputer.canHideIme(ImeTracker.Token.empty(),
+                    InputMethodManager.HIDE_NOT_ALWAYS)).isTrue();
+            mComputer.requestImeVisibility(mWindowToken, false);
 
-        final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
-        assertThat(state).isNotNull();
-        assertThat(state.hasEditorFocused()).isTrue();
-        assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
-        assertThat(state.isRequestedImeVisible()).isFalse();
+            final ImeTargetWindowState state = mComputer.getWindowStateOrNull(mWindowToken);
+            assertThat(state).isNotNull();
+            assertThat(state.hasEditorFocused()).isTrue();
+            assertThat(state.getSoftInputModeState()).isEqualTo(SOFT_INPUT_STATE_UNCHANGED);
+            assertThat(state.isRequestedImeVisible()).isFalse();
+        }
     }
 
     @Test
     public void testComputeImeDisplayId() {
-        final ImeTargetWindowState state = mComputer.getOrCreateWindowState(mWindowToken);
+        synchronized (ImfLock.class) {
+            final ImeTargetWindowState state = mComputer.getOrCreateWindowState(mWindowToken);
 
-        mImeDisplayPolicy = DISPLAY_IME_POLICY_LOCAL;
-        mComputer.computeImeDisplayId(state, DEFAULT_DISPLAY);
-        assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isFalse();
-        assertThat(state.getImeDisplayId()).isEqualTo(DEFAULT_DISPLAY);
+            mImeDisplayPolicy = DISPLAY_IME_POLICY_LOCAL;
+            mComputer.computeImeDisplayId(state, DEFAULT_DISPLAY);
+            assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isFalse();
+            assertThat(state.getImeDisplayId()).isEqualTo(DEFAULT_DISPLAY);
 
-        mComputer.computeImeDisplayId(state, 10 /* displayId */);
-        assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isFalse();
-        assertThat(state.getImeDisplayId()).isEqualTo(10);
+            mComputer.computeImeDisplayId(state, 10 /* displayId */);
+            assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isFalse();
+            assertThat(state.getImeDisplayId()).isEqualTo(10);
 
-        mImeDisplayPolicy = DISPLAY_IME_POLICY_HIDE;
-        mComputer.computeImeDisplayId(state, 10 /* displayId */);
-        assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isTrue();
-        assertThat(state.getImeDisplayId()).isEqualTo(INVALID_DISPLAY);
+            mImeDisplayPolicy = DISPLAY_IME_POLICY_HIDE;
+            mComputer.computeImeDisplayId(state, 10 /* displayId */);
+            assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isTrue();
+            assertThat(state.getImeDisplayId()).isEqualTo(INVALID_DISPLAY);
 
-        mImeDisplayPolicy = DISPLAY_IME_POLICY_FALLBACK_DISPLAY;
-        mComputer.computeImeDisplayId(state, 10 /* displayId */);
-        assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isFalse();
-        assertThat(state.getImeDisplayId()).isEqualTo(FALLBACK_DISPLAY_ID);
+            mImeDisplayPolicy = DISPLAY_IME_POLICY_FALLBACK_DISPLAY;
+            mComputer.computeImeDisplayId(state, 10 /* displayId */);
+            assertThat(mComputer.getImePolicy().isImeHiddenByDisplayPolicy()).isFalse();
+            assertThat(state.getImeDisplayId()).isEqualTo(FALLBACK_DISPLAY_ID);
+        }
     }
 
     @Test
     public void testComputeState_lastImeRequestedVisible_preserved_When_StateUnChanged() {
-        // Assume the last IME targeted window has requested IME visible
-        final IBinder lastImeTargetWindowToken = new Binder();
-        mInputMethodManagerService.mLastImeTargetWindow = lastImeTargetWindowToken;
-        mComputer.requestImeVisibility(lastImeTargetWindowToken, true);
-        final ImeTargetWindowState lastState = mComputer.getWindowStateOrNull(
-                lastImeTargetWindowToken);
-        assertThat(lastState.isRequestedImeVisible()).isTrue();
+        synchronized (ImfLock.class) {
+            // Assume the last IME targeted window has requested IME visible
+            final IBinder lastImeTargetWindowToken = new Binder();
+            mComputer.setLastImeTargetWindow(lastImeTargetWindowToken);
+            mComputer.requestImeVisibility(lastImeTargetWindowToken, true);
+            final ImeTargetWindowState lastState = mComputer.getWindowStateOrNull(
+                    lastImeTargetWindowToken);
+            assertThat(lastState.isRequestedImeVisible()).isTrue();
 
-        // Verify when focusing the next window with STATE_UNCHANGED flag, the last IME
-        // visibility state will be preserved to the current window state.
-        final ImeTargetWindowState stateWithUnChangedFlag = initImeTargetWindowState(mWindowToken);
-        mComputer.computeState(stateWithUnChangedFlag, true /* allowVisible */);
-        assertThat(stateWithUnChangedFlag.isRequestedImeVisible()).isEqualTo(
-                lastState.isRequestedImeVisible());
+            // Verify when focusing the next window with STATE_UNCHANGED flag, the last IME
+            // visibility state will be preserved to the current window state.
+            final ImeTargetWindowState stateWithUnChangedFlag = initImeTargetWindowState(
+                    mWindowToken);
+            mComputer.computeState(stateWithUnChangedFlag, true /* allowVisible */);
+            assertThat(stateWithUnChangedFlag.isRequestedImeVisible()).isEqualTo(
+                    lastState.isRequestedImeVisible());
+        }
     }
 
     @Test
     public void testOnInteractiveChanged() {
-        mComputer.getOrCreateWindowState(mWindowToken);
-        // Precondition: ensure IME has shown before hiding request.
-        mComputer.requestImeVisibility(mWindowToken, true);
-        mComputer.setInputShown(true);
+        synchronized (ImfLock.class) {
+            mComputer.getOrCreateWindowState(mWindowToken);
+            // Precondition: ensure IME has shown before hiding request.
+            mComputer.requestImeVisibility(mWindowToken, true);
+            mComputer.setInputShown(true);
 
-        // No need any visibility change When initially shows IME on the device was interactive.
-        ImeVisibilityStateComputer.ImeVisibilityResult result = mComputer.onInteractiveChanged(
-                mWindowToken, true /* interactive */);
-        assertThat(result).isNull();
+            // No need any visibility change When initially shows IME on the device was interactive.
+            ImeVisibilityStateComputer.ImeVisibilityResult result = mComputer.onInteractiveChanged(
+                    mWindowToken, true /* interactive */);
+            assertThat(result).isNull();
 
-        // Show the IME screenshot to capture the last IME visible state when the device inactive.
-        result = mComputer.onInteractiveChanged(mWindowToken, false /* interactive */);
-        assertThat(result).isNotNull();
-        assertThat(result.getState()).isEqualTo(STATE_SHOW_IME_SNAPSHOT);
-        assertThat(result.getReason()).isEqualTo(SHOW_IME_SCREENSHOT_FROM_IMMS);
+            // Show the IME screenshot to capture the last IME visible state when the device
+            // inactive.
+            result = mComputer.onInteractiveChanged(mWindowToken, false /* interactive */);
+            assertThat(result).isNotNull();
+            assertThat(result.getState()).isEqualTo(STATE_SHOW_IME_SNAPSHOT);
+            assertThat(result.getReason()).isEqualTo(SHOW_IME_SCREENSHOT_FROM_IMMS);
 
-        // Remove the IME screenshot when the device became interactive again.
-        result = mComputer.onInteractiveChanged(mWindowToken, true /* interactive */);
-        assertThat(result).isNotNull();
-        assertThat(result.getState()).isEqualTo(STATE_REMOVE_IME_SNAPSHOT);
-        assertThat(result.getReason()).isEqualTo(REMOVE_IME_SCREENSHOT_FROM_IMMS);
+            // Remove the IME screenshot when the device became interactive again.
+            result = mComputer.onInteractiveChanged(mWindowToken, true /* interactive */);
+            assertThat(result).isNotNull();
+            assertThat(result.getState()).isEqualTo(STATE_REMOVE_IME_SNAPSHOT);
+            assertThat(result.getReason()).isEqualTo(REMOVE_IME_SCREENSHOT_FROM_IMMS);
+        }
     }
 
     @Test
     public void testOnApplyImeVisibilityFromComputer() {
-        final IBinder testImeTargetOverlay = new Binder();
-        final IBinder testImeInputTarget = new Binder();
+        synchronized (ImfLock.class) {
+            final IBinder testImeTargetOverlay = new Binder();
+            final IBinder testImeInputTarget = new Binder();
 
-        // Simulate a test IME input target was visible.
-        mListener.onImeInputTargetVisibilityChanged(testImeInputTarget, true, false);
+            // Simulate a test IME input target was visible.
+            mListener.onImeInputTargetVisibilityChanged(testImeInputTarget, true, false);
 
-        // Simulate a test IME layering target overlay fully occluded the IME input target.
-        mListener.onImeTargetOverlayVisibilityChanged(testImeTargetOverlay,
-                TYPE_APPLICATION_OVERLAY, true, false);
-        mListener.onImeInputTargetVisibilityChanged(testImeInputTarget, false, false);
-        final ArgumentCaptor<IBinder> targetCaptor = ArgumentCaptor.forClass(IBinder.class);
-        final ArgumentCaptor<ImeVisibilityResult> resultCaptor = ArgumentCaptor.forClass(
-                ImeVisibilityResult.class);
-        verify(mInputMethodManagerService).onApplyImeVisibilityFromComputer(targetCaptor.capture(),
-                notNull() /* statsToken */, resultCaptor.capture());
-        final IBinder imeInputTarget = targetCaptor.getValue();
-        final ImeVisibilityResult result = resultCaptor.getValue();
+            // Simulate a test IME layering target overlay fully occluded the IME input target.
+            mListener.onImeTargetOverlayVisibilityChanged(testImeTargetOverlay,
+                    TYPE_APPLICATION_OVERLAY, true, false);
+            mListener.onImeInputTargetVisibilityChanged(testImeInputTarget, false, false);
+            final ArgumentCaptor<IBinder> targetCaptor = ArgumentCaptor.forClass(IBinder.class);
+            final ArgumentCaptor<ImeVisibilityResult> resultCaptor = ArgumentCaptor.forClass(
+                    ImeVisibilityResult.class);
+            verify(mInputMethodManagerService).onApplyImeVisibilityFromComputerLocked(
+                    targetCaptor.capture(), notNull() /* statsToken */, resultCaptor.capture());
+            final IBinder imeInputTarget = targetCaptor.getValue();
+            final ImeVisibilityResult result = resultCaptor.getValue();
 
-        // Verify the computer will callback hiding IME state to IMMS.
-        assertThat(imeInputTarget).isEqualTo(testImeInputTarget);
-        assertThat(result.getState()).isEqualTo(STATE_HIDE_IME_EXPLICIT);
-        assertThat(result.getReason()).isEqualTo(HIDE_WHEN_INPUT_TARGET_INVISIBLE);
+            // Verify the computer will callback hiding IME state to IMMS.
+            assertThat(imeInputTarget).isEqualTo(testImeInputTarget);
+            assertThat(result.getState()).isEqualTo(STATE_HIDE_IME_EXPLICIT);
+            assertThat(result.getReason()).isEqualTo(HIDE_WHEN_INPUT_TARGET_INVISIBLE);
+        }
     }
 
+    @GuardedBy("ImfLock.class")
     private ImeTargetWindowState initImeTargetWindowState(IBinder windowToken) {
         final ImeTargetWindowState state = new ImeTargetWindowState(SOFT_INPUT_STATE_UNCHANGED,
                 0, true, true, true);
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImmutableSparseArrayTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImmutableSparseArrayTest.java
new file mode 100644
index 0000000..944b7c6
--- /dev/null
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/ImmutableSparseArrayTest.java
@@ -0,0 +1,278 @@
+/*
+ * Copyright (C) 2024 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.server.inputmethod;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.assertThrows;
+
+import android.annotation.NonNull;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.ArrayList;
+
+@RunWith(AndroidJUnit4.class)
+public final class ImmutableSparseArrayTest {
+
+    @Test
+    public void testEmptyObject() {
+        final ImmutableSparseArray<Object> empty = ImmutableSparseArray.empty();
+
+        assertThat(empty.size()).isEqualTo(0);
+        verifyCommonBehaviors(empty);
+    }
+
+    @Test
+    public void testEmptyMethod() {
+        assertThat(ImmutableSparseArray.empty()).isSameInstanceAs(ImmutableSparseArray.empty());
+    }
+
+    @Test
+    public void testCloneWithPutOrSelf_appendingFromEmpty() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = -2;  // intentionally negative
+        final Object value2 = new Object();
+        final int key3 = -3;  // intentionally negative
+        final Object value3 = new Object();
+        final int key4 = 4;
+        final Object value4 = new Object();
+
+        final ImmutableSparseArray<Object> oneItemArray = ImmutableSparseArray.empty()
+                .cloneWithPutOrSelf(key1, value1);
+        verifyCommonBehaviors(oneItemArray);
+        assertThat(oneItemArray.size()).isEqualTo(1);
+        assertThat(oneItemArray.get(key1)).isSameInstanceAs(value1);
+
+        final ImmutableSparseArray<Object> twoItemArray =
+                oneItemArray.cloneWithPutOrSelf(key2, value2);
+        assertThat(twoItemArray).isNotSameInstanceAs(oneItemArray);
+        verifyCommonBehaviors(twoItemArray);
+        assertThat(twoItemArray.size()).isEqualTo(2);
+        assertThat(twoItemArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(twoItemArray.get(key2)).isSameInstanceAs(value2);
+
+        final ImmutableSparseArray<Object> threeItemArray =
+                twoItemArray.cloneWithPutOrSelf(key3, value3);
+        assertThat(threeItemArray).isNotSameInstanceAs(twoItemArray);
+        verifyCommonBehaviors(threeItemArray);
+        assertThat(threeItemArray.size()).isEqualTo(3);
+        assertThat(threeItemArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(threeItemArray.get(key2)).isSameInstanceAs(value2);
+        assertThat(threeItemArray.get(key3)).isSameInstanceAs(value3);
+
+        final ImmutableSparseArray<Object> fourItemArray =
+                threeItemArray.cloneWithPutOrSelf(key4, value4);
+        assertThat(fourItemArray).isNotSameInstanceAs(threeItemArray);
+        verifyCommonBehaviors(fourItemArray);
+        assertThat(fourItemArray.size()).isEqualTo(4);
+        assertThat(fourItemArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(fourItemArray.get(key2)).isSameInstanceAs(value2);
+        assertThat(fourItemArray.get(key3)).isSameInstanceAs(value3);
+        assertThat(fourItemArray.get(key4)).isSameInstanceAs(value4);
+    }
+
+    @Test
+    public void testCloneWithPutOrSelf_returnSelf() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1);
+        assertThat(array.cloneWithPutOrSelf(key1, value1)).isSameInstanceAs(array);
+    }
+
+    @Test
+    public void testCloneWithPutOrSelf_updateExistingValue() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final Object value2updated = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3);
+
+        final var updatedArray = array.cloneWithPutOrSelf(key2, value2updated);
+        verifyCommonBehaviors(updatedArray);
+
+        assertThat(updatedArray.size()).isEqualTo(3);
+        assertThat(updatedArray.get(key1)).isSameInstanceAs(value1);
+        assertThat(updatedArray.get(key2)).isSameInstanceAs(value2updated);
+        assertThat(updatedArray.get(key3)).isSameInstanceAs(value3);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_empty() {
+        final ImmutableSparseArray<Object> empty = ImmutableSparseArray.empty();
+        assertThat(empty.cloneWithRemoveOrSelf(0)).isSameInstanceAs(empty);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_singleInstance() {
+        final int key = 1;
+        final Object value = new Object();
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key, value);
+        assertThat(array.cloneWithRemoveOrSelf(key)).isSameInstanceAs(ImmutableSparseArray.empty());
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_firstItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3)
+                .cloneWithRemoveOrSelf(key1);
+        verifyCommonBehaviors(array);
+
+        assertThat(array.size()).isEqualTo(2);
+        assertThat(array.get(key1)).isNull();
+        assertThat(array.get(key2)).isSameInstanceAs(value2);
+        assertThat(array.get(key3)).isSameInstanceAs(value3);
+        assertThat(array.keyAt(0)).isEqualTo(key2);
+        assertThat(array.keyAt(1)).isEqualTo(key3);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_lastItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3)
+                .cloneWithRemoveOrSelf(key3);
+        verifyCommonBehaviors(array);
+
+        assertThat(array.size()).isEqualTo(2);
+        assertThat(array.get(key1)).isSameInstanceAs(value1);
+        assertThat(array.get(key2)).isSameInstanceAs(value2);
+        assertThat(array.get(key3)).isNull();
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_middleItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3)
+                .cloneWithRemoveOrSelf(key2);
+        verifyCommonBehaviors(array);
+
+        assertThat(array.size()).isEqualTo(2);
+        assertThat(array.get(key1)).isSameInstanceAs(value1);
+        assertThat(array.get(key2)).isNull();
+        assertThat(array.get(key3)).isSameInstanceAs(value3);
+    }
+
+    @Test
+    public void testCloneWithRemoveOrSelf_nonExistentItem() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+        final int key4 = 4;
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3);
+
+        assertThat(array.cloneWithRemoveOrSelf(key4)).isSameInstanceAs(array);
+    }
+
+    @Test
+    public void testForEach() {
+        final int key1 = 1;
+        final Object value1 = new Object();
+        final int key2 = 2;
+        final Object value2 = new Object();
+        final int key3 = 3;
+        final Object value3 = new Object();
+
+        final ImmutableSparseArray<Object> array = ImmutableSparseArray
+                .empty()
+                .cloneWithPutOrSelf(key1, value1)
+                .cloneWithPutOrSelf(key2, value2)
+                .cloneWithPutOrSelf(key3, value3);
+
+        final ArrayList<Object> list = new ArrayList<>();
+        array.forEach(list::add);
+        assertThat(list).containsExactlyElementsIn(new Object[]{ value1, value2, value3 })
+                .inOrder();
+    }
+
+
+    private void verifyCommonBehaviors(@NonNull ImmutableSparseArray<Object> sparseArray) {
+        verifyInvalidKeyBehaviors(sparseArray);
+        verifyOutOfBoundsBehaviors(sparseArray);
+    }
+
+    private void verifyInvalidKeyBehaviors(@NonNull ImmutableSparseArray<Object> sparseArray) {
+        final int invalid_key = -123456678;
+        assertThat(sparseArray.get(invalid_key)).isNull();
+        assertThat(sparseArray.indexOfKey(invalid_key)).isEqualTo(-1);
+    }
+
+    private void verifyOutOfBoundsBehaviors(@NonNull ImmutableSparseArray<Object> sparseArray) {
+        final int size = sparseArray.size();
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.keyAt(size));
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.valueAt(size));
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.keyAt(-1));
+        assertThrows(ArrayIndexOutOfBoundsException.class,
+                () -> sparseArray.valueAt(-1));
+    }
+}
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java
index 4d28b3c..1e3b7e9 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodBindingControllerTest.java
@@ -140,8 +140,7 @@
         final InputMethodInfo info;
         synchronized (ImfLock.class) {
             mBindingController.setSelectedMethodId(TEST_IME_ID);
-            info = InputMethodSettingsRepository.get(mCallingUserId).getMethodMap()
-                    .get(TEST_IME_ID);
+            info = InputMethodSettingsRepository.get(mUserId).getMethodMap().get(TEST_IME_ID);
         }
         assertThat(info).isNotNull();
         assertThat(info.getId()).isEqualTo(TEST_IME_ID);
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
index ec9bfa7..461697c 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceTestBase.java
@@ -50,7 +50,6 @@
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.ServiceManager;
-import android.os.UserHandle;
 import android.util.ArraySet;
 import android.view.InputChannel;
 import android.view.inputmethod.EditorInfo;
@@ -128,7 +127,7 @@
     protected Context mContext;
     protected MockitoSession mMockingSession;
     protected int mTargetSdkVersion;
-    protected int mCallingUserId;
+    protected int mUserId;
     protected EditorInfo mEditorInfo;
     protected IInputMethodInvoker mMockInputMethodInvoker;
     protected InputMethodManagerService mInputMethodManagerService;
@@ -165,12 +164,12 @@
                         .spyStatic(AdditionalSubtypeUtils.class)
                         .startMocking();
 
-        mContext = spy(InstrumentationRegistry.getInstrumentation().getContext());
+        mContext = spy(InstrumentationRegistry.getInstrumentation().getTargetContext());
 
         mTargetSdkVersion = mContext.getApplicationInfo().targetSdkVersion;
         mIsLargeScreen = mContext.getResources().getConfiguration()
                 .isLayoutSizeAtLeast(Configuration.SCREENLAYOUT_SIZE_LARGE);
-        mCallingUserId = UserHandle.getCallingUserId();
+        mUserId = mContext.getUserId();
         mEditorInfo = new EditorInfo();
         mEditorInfo.packageName = TEST_EDITOR_PKG_NAME;
 
@@ -202,7 +201,7 @@
         // Injecting and mocked InputMethodBindingController and InputMethod.
         mMockInputMethodInvoker = IInputMethodInvoker.create(mMockInputMethod);
         mInputManagerGlobalSession = InputManagerGlobal.createTestSession(mMockIInputManager);
-        when(mMockInputMethodBindingController.getUserId()).thenReturn(mCallingUserId);
+        when(mMockInputMethodBindingController.getUserId()).thenReturn(mUserId);
         synchronized (ImfLock.class) {
             when(mMockInputMethodBindingController.getCurMethod())
                     .thenReturn(mMockInputMethodInvoker);
@@ -222,7 +221,7 @@
                 .thenReturn(new int[] {0});
         when(mMockUserManagerInternal.getUserIds()).thenReturn(new int[] {0});
         when(mMockActivityManagerInternal.isSystemReady()).thenReturn(true);
-        when(mMockActivityManagerInternal.getCurrentUserId()).thenReturn(mCallingUserId);
+        when(mMockActivityManagerInternal.getCurrentUserId()).thenReturn(mUserId);
         when(mMockPackageManagerInternal.getPackageUid(anyString(), anyLong(), anyInt()))
                 .thenReturn(Binder.getCallingUid());
         when(mMockPackageManagerInternal.isSameApp(anyString(), anyLong(), anyInt(), anyInt()))
@@ -272,14 +271,13 @@
 
         // Certain tests rely on TEST_IME_ID that is installed with AndroidTest.xml.
         // TODO(b/352615651): Consider just synthesizing test InputMethodInfo then injecting it.
-        AdditionalSubtypeMapRepository.ensureInitializedAndGet(mCallingUserId);
+        AdditionalSubtypeMapRepository.initializeIfNecessary(mUserId);
         final var settings = InputMethodManagerService.queryInputMethodServicesInternal(mContext,
-                mCallingUserId, AdditionalSubtypeMapRepository.get(mCallingUserId),
-                DirectBootAwareness.AUTO);
-        InputMethodSettingsRepository.put(mCallingUserId, settings);
+                mUserId, AdditionalSubtypeMapRepository.get(mUserId), DirectBootAwareness.AUTO);
+        InputMethodSettingsRepository.put(mUserId, settings);
 
         // Emulate that the user initialization is done.
-        mInputMethodManagerService.getUserData(mCallingUserId).mBackgroundLoadLatch.countDown();
+        mInputMethodManagerService.getUserData(mUserId).mBackgroundLoadLatch.countDown();
 
         // After this boot phase, services can broadcast Intents.
         lifecycle.onBootPhase(SystemService.PHASE_ACTIVITY_MANAGER_READY);
@@ -291,7 +289,7 @@
 
     @After
     public void tearDown() {
-        InputMethodSettingsRepository.remove(mCallingUserId);
+        InputMethodSettingsRepository.remove(mUserId);
 
         if (mInputMethodManagerService != null) {
             mInputMethodManagerService.mInputMethodDeviceConfigs.destroy();
@@ -347,8 +345,8 @@
         synchronized (ImfLock.class) {
             ClientState cs = mInputMethodManagerService.getClientStateLocked(client);
             cs.mCurSession = new InputMethodManagerService.SessionState(cs,
-                    mMockInputMethodInvoker, mMockInputMethodSession, mock(
-                    InputChannel.class));
+                    mMockInputMethodInvoker, mMockInputMethodSession, mock(InputChannel.class),
+                    mUserId);
         }
     }
 }
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceWindowGainedFocusTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceWindowGainedFocusTest.java
index ffc4df8..c5b5668 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceWindowGainedFocusTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/server/inputmethod/InputMethodManagerServiceWindowGainedFocusTest.java
@@ -263,7 +263,7 @@
                 mMockRemoteInputConnection /* inputConnection */,
                 mMockRemoteAccessibilityInputConnection /* remoteAccessibilityInputConnection */,
                 mTargetSdkVersion /* unverifiedTargetSdkVersion */,
-                mCallingUserId /* userId */,
+                mUserId /* userId */,
                 mMockImeOnBackInvokedDispatcher /* imeDispatcher */);
     }
 
diff --git a/services/tests/PackageManagerServiceTests/server/src/com/android/server/pm/PackageManagerServiceUtilsTest.java b/services/tests/PackageManagerServiceTests/server/src/com/android/server/pm/PackageManagerServiceUtilsTest.java
new file mode 100644
index 0000000..4a2396d
--- /dev/null
+++ b/services/tests/PackageManagerServiceTests/server/src/com/android/server/pm/PackageManagerServiceUtilsTest.java
@@ -0,0 +1,341 @@
+/*
+ * Copyright (C) 2024 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.server.pm;
+
+import static org.testng.Assert.assertThrows;
+
+import android.content.pm.PackageInfoLite;
+import android.platform.test.annotations.Presubmit;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.internal.pm.parsing.pkg.PackageImpl;
+import com.android.server.pm.pkg.AndroidPackage;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.UUID;
+
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class PackageManagerServiceUtilsTest {
+
+    private static final String PACKAGE_NAME = "com.android.app";
+    private static final File CODE_PATH =
+            InstrumentationRegistry.getInstrumentation().getContext().getFilesDir();
+
+    @Test
+    public void testCheckDowngrade_packageSetting_versionCodeSmaller_throwException()
+            throws Exception {
+        final PackageSetting before = createPackageSetting();
+        before.setLongVersionCode(2);
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_packageSetting_baseRevisionCodeSmaller_throwException()
+            throws Exception {
+        final PackageSetting before = createPackageSetting();
+        before.setLongVersionCode(1);
+        before.setBaseRevisionCode(2);
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.baseRevisionCode = 1;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_packageSetting_splitArraySizeIsDifferent_throwException()
+            throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final String[] splitNames = new String[] { splitOne, splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionOne };
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne, revisionTwo };
+
+        final PackageSetting before = createPackageSetting();
+        before.setLongVersionCode(1);
+        before.setSplitNames(splitNames);
+        before.setSplitRevisionCodes(beforeSplitRevisionCodes);
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = splitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_packageSetting_splitRevisionCodeSmaller_throwException()
+            throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final int revisionThree = 360;
+        final String[] splitNames = new String[] { splitOne, splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionTwo, revisionThree};
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne, revisionTwo };
+
+        final PackageSetting before = createPackageSetting();
+        before.setLongVersionCode(1);
+        before.setSplitNames(splitNames);
+        before.setSplitRevisionCodes(beforeSplitRevisionCodes);
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = splitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_packageSetting_sameSplitNameRevisionsBigger()
+            throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final int revisionThree = 360;
+        final String[] splitNames = new String[] { splitOne, splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionOne, revisionTwo};
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne, revisionThree };
+
+        final PackageSetting before = createPackageSetting();
+        before.setLongVersionCode(1);
+        before.setSplitNames(splitNames);
+        before.setSplitRevisionCodes(beforeSplitRevisionCodes);
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = splitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        PackageManagerServiceUtils.checkDowngrade(before, after);
+    }
+
+    @Test
+    public void testCheckDowngrade_packageSetting_hasDifferentSplitNames() throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final int revisionThree = 360;
+        final String[] beforeSplitNames = new String[] { splitOne, splitTwo };
+        final String[] afterSplitNames = new String[] { splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionOne, revisionTwo};
+        final int[] afterSplitRevisionCodes = new int[] { revisionThree };
+
+        final PackageSetting before = createPackageSetting();
+        before.setLongVersionCode(1);
+        before.setSplitNames(beforeSplitNames);
+        before.setSplitRevisionCodes(beforeSplitRevisionCodes);
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = afterSplitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        PackageManagerServiceUtils.checkDowngrade(before, after);
+    }
+
+    @Test
+    public void testCheckDowngrade_packageSetting_newSplitName() throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final String[] beforeSplitNames = new String[] { splitOne };
+        final String[] afterSplitNames = new String[] { splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionTwo };
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne };
+
+        final PackageSetting before = createPackageSetting();
+        before.setLongVersionCode(1);
+        before.setSplitNames(beforeSplitNames);
+        before.setSplitRevisionCodes(beforeSplitRevisionCodes);
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = afterSplitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        PackageManagerServiceUtils.checkDowngrade(before, after);
+    }
+
+    @Test
+    public void testCheckDowngrade_androidPackage_versionCodeSmaller_throwException()
+            throws Exception {
+        final AndroidPackage before = PackageImpl.forTesting(PACKAGE_NAME).hideAsParsed()
+                .setVersionCode(2).hideAsFinal();
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_androidPackage_baseRevisionCodeSmaller_throwException()
+            throws Exception {
+        final AndroidPackage before = PackageImpl.forTesting(PACKAGE_NAME).setBaseRevisionCode(2)
+                .hideAsParsed().setVersionCode(1).hideAsFinal();
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.baseRevisionCode = 1;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_androidPackage_splitArraySizeIsDifferent_throwException()
+            throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final String[] splitNames = new String[] { splitOne, splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionOne };
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne, revisionTwo };
+
+        final AndroidPackage before = PackageImpl.forTesting(PACKAGE_NAME)
+                .asSplit(splitNames, /* splitCodePaths= */ null,
+                        beforeSplitRevisionCodes, /* splitDependencies= */ null)
+                .hideAsParsed().setVersionCode(1).hideAsFinal();
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = splitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_androidPackage_splitRevisionCodeSmaller_throwException()
+            throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final int revisionThree = 360;
+        final String[] splitNames = new String[] { splitOne, splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionTwo, revisionThree};
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne, revisionTwo };
+
+        final AndroidPackage before = PackageImpl.forTesting(PACKAGE_NAME)
+                .asSplit(splitNames, /* splitCodePaths= */ null,
+                        beforeSplitRevisionCodes, /* splitDependencies= */ null)
+                .hideAsParsed().setVersionCode(1).hideAsFinal();
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = splitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        assertThrows(PackageManagerException.class,
+                () -> PackageManagerServiceUtils.checkDowngrade(before, after));
+    }
+
+    @Test
+    public void testCheckDowngrade_androidPackage_sameSplitNameRevisionsBigger()
+            throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final int revisionThree = 360;
+        final String[] splitNames = new String[] { splitOne, splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionOne, revisionTwo};
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne, revisionThree };
+
+        final AndroidPackage before = PackageImpl.forTesting(PACKAGE_NAME)
+                .asSplit(splitNames, /* splitCodePaths= */ null,
+                        beforeSplitRevisionCodes, /* splitDependencies= */ null)
+                .hideAsParsed().setVersionCode(1).hideAsFinal();
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = splitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        PackageManagerServiceUtils.checkDowngrade(before, after);
+    }
+
+    @Test
+    public void testCheckDowngrade_androidPackage_hasDifferentSplitNames() throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final int revisionThree = 360;
+        final String[] beforeSplitNames = new String[] { splitOne, splitTwo };
+        final String[] afterSplitNames = new String[] { splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionOne, revisionTwo};
+        final int[] afterSplitRevisionCodes = new int[] { revisionThree };
+
+        final AndroidPackage before = PackageImpl.forTesting(PACKAGE_NAME)
+                .asSplit(beforeSplitNames, /* splitCodePaths= */ null,
+                        beforeSplitRevisionCodes, /* splitDependencies= */ null)
+                .hideAsParsed().setVersionCode(1).hideAsFinal();
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = afterSplitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        PackageManagerServiceUtils.checkDowngrade(before, after);
+    }
+
+    @Test
+    public void testCheckDowngrade_androidPackage_newSplitName() throws Exception {
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        final String[] beforeSplitNames = new String[] { splitOne };
+        final String[] afterSplitNames = new String[] { splitTwo };
+        final int[] beforeSplitRevisionCodes = new int[] { revisionTwo };
+        final int[] afterSplitRevisionCodes = new int[] { revisionOne };
+
+        final AndroidPackage before = PackageImpl.forTesting(PACKAGE_NAME)
+                .asSplit(beforeSplitNames, /* splitCodePaths= */ null,
+                        beforeSplitRevisionCodes, /* splitDependencies= */ null)
+                .hideAsParsed().setVersionCode(1).hideAsFinal();
+        final PackageInfoLite after = new PackageInfoLite();
+        after.versionCode = 1;
+        after.splitNames = afterSplitNames;
+        after.splitRevisionCodes = afterSplitRevisionCodes;
+
+        PackageManagerServiceUtils.checkDowngrade(before, after);
+    }
+
+    private PackageSetting createPackageSetting() {
+        return new PackageSetting(PACKAGE_NAME, PACKAGE_NAME, CODE_PATH, /* pkgFlags= */ 0,
+                /* privateFlags= */ 0 , UUID.randomUUID());
+    }
+}
diff --git a/services/tests/PackageManagerServiceTests/server/src/com/android/server/pm/PackageManagerSettingsTests.java b/services/tests/PackageManagerServiceTests/server/src/com/android/server/pm/PackageManagerSettingsTests.java
index dec4634..d7af443 100644
--- a/services/tests/PackageManagerServiceTests/server/src/com/android/server/pm/PackageManagerSettingsTests.java
+++ b/services/tests/PackageManagerServiceTests/server/src/com/android/server/pm/PackageManagerSettingsTests.java
@@ -1092,7 +1092,7 @@
     }
 
     @Test
-    public void testNoPkg_writeReadSplitVersions() {
+    public void testNoPkgDifferentRevisions_writeReadSplitVersions() {
         Settings settings = makeSettings();
         PackageSetting packageSetting = createPackageSetting(PACKAGE_NAME_1);
         packageSetting.setAppId(Process.FIRST_APPLICATION_UID);
@@ -1117,6 +1117,54 @@
     }
 
     @Test
+    public void testNoPkgSameRevisions_writeReadSplitVersions() {
+        Settings settings = makeSettings();
+        PackageSetting packageSetting = createPackageSetting(PACKAGE_NAME_1);
+        packageSetting.setAppId(Process.FIRST_APPLICATION_UID);
+
+        final String splitOne = "one";
+        final String splitTwo = "two";
+        final int revisionOne = 311;
+        packageSetting.setSplitNames(new String[] { splitOne, splitTwo});
+        packageSetting.setSplitRevisionCodes(new int[] { revisionOne, revisionOne});
+        settings.mPackages.put(PACKAGE_NAME_1, packageSetting);
+
+        settings.writeLPr(computer, /* sync= */ true);
+        settings.mPackages.clear();
+
+        assertThat(settings.readLPw(computer, createFakeUsers()), is(true));
+        PackageSetting resultSetting = settings.getPackageLPr(PACKAGE_NAME_1);
+        assertThat(resultSetting.getSplitNames()[0], is(splitOne));
+        assertThat(resultSetting.getSplitNames()[1], is(splitTwo));
+        assertThat(resultSetting.getSplitRevisionCodes()[0], is(revisionOne));
+        assertThat(resultSetting.getSplitRevisionCodes()[1], is(revisionOne));
+    }
+
+    @Test
+    public void testNoPkgSameSplitNames_writeReadSplitVersions() {
+        Settings settings = makeSettings();
+        PackageSetting packageSetting = createPackageSetting(PACKAGE_NAME_1);
+        packageSetting.setAppId(Process.FIRST_APPLICATION_UID);
+
+        final String splitOne = "one";
+        final int revisionOne = 311;
+        final int revisionTwo = 330;
+        packageSetting.setSplitNames(new String[] { splitOne, splitOne});
+        packageSetting.setSplitRevisionCodes(new int[] { revisionOne, revisionTwo});
+        settings.mPackages.put(PACKAGE_NAME_1, packageSetting);
+
+        settings.writeLPr(computer, /* sync= */ true);
+        settings.mPackages.clear();
+
+        assertThat(settings.readLPw(computer, createFakeUsers()), is(true));
+        PackageSetting resultSetting = settings.getPackageLPr(PACKAGE_NAME_1);
+        assertThat(resultSetting.getSplitNames().length, is(1));
+        assertThat(resultSetting.getSplitRevisionCodes().length, is(1));
+        assertThat(resultSetting.getSplitNames()[0], is(splitOne));
+        assertThat(resultSetting.getSplitRevisionCodes()[0], is(revisionTwo));
+    }
+
+    @Test
     public void testWriteReadArchiveState() {
         Settings settings = makeSettings();
         PackageSetting packageSetting = createPackageSetting(PACKAGE_NAME_1);
diff --git a/services/tests/performancehinttests/Android.bp b/services/tests/performancehinttests/Android.bp
new file mode 100644
index 0000000..1692921c
--- /dev/null
+++ b/services/tests/performancehinttests/Android.bp
@@ -0,0 +1,34 @@
+package {
+    default_team: "trendy_team_games",
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_test {
+    name: "PerformanceHintTests",
+    srcs: [
+        "src/**/*.java",
+    ],
+    static_libs: [
+        "androidx.test.runner",
+        "androidx.test.rules",
+        "flag-junit",
+        "junit",
+        "mockito-target-minus-junit4",
+        "platform-test-annotations",
+        "services.core",
+        "truth",
+    ],
+    libs: [
+        "android.test.base",
+    ],
+    test_suites: [
+        "automotive-tests",
+        "device-tests",
+    ],
+    platform_apis: true,
+    certificate: "platform",
+    dxflags: ["--multi-dex"],
+    optimize: {
+        enabled: false,
+    },
+}
diff --git a/services/tests/performancehinttests/AndroidManifest.xml b/services/tests/performancehinttests/AndroidManifest.xml
new file mode 100644
index 0000000..d955234
--- /dev/null
+++ b/services/tests/performancehinttests/AndroidManifest.xml
@@ -0,0 +1,23 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+     package="com.android.server.power.hinttests">
+
+    <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
+         android:targetPackage="com.android.server.power.hinttests"
+         android:label="ADPF Performance Hint Service Test"/>
+</manifest>
diff --git a/services/tests/performancehinttests/AndroidTest.xml b/services/tests/performancehinttests/AndroidTest.xml
new file mode 100644
index 0000000..578b7d6
--- /dev/null
+++ b/services/tests/performancehinttests/AndroidTest.xml
@@ -0,0 +1,34 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 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.
+-->
+<configuration description="Runs Performance Hint Tests.">
+    <option name="test-suite-tag" value="apct" />
+    <option name="test-suite-tag" value="apct-instrumentation" />
+
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="install-arg" value="-t" />
+        <option name="test-file-name" value="PerformanceHintTests.apk" />
+    </target_preparer>
+
+    <option name="test-tag" value="PerformanceHintTests" />
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.server.power.hinttests" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+        <option name="hidden-api-checks" value="false"/>
+        <option name="exclude-annotation" value="androidx.test.filters.FlakyTest" />
+    </test>
+</configuration>
diff --git a/services/tests/servicestests/src/com/android/server/power/hint/OWNERS b/services/tests/performancehinttests/OWNERS
similarity index 100%
rename from services/tests/servicestests/src/com/android/server/power/hint/OWNERS
rename to services/tests/performancehinttests/OWNERS
diff --git a/services/tests/performancehinttests/TEST_MAPPING b/services/tests/performancehinttests/TEST_MAPPING
new file mode 100644
index 0000000..faffe35
--- /dev/null
+++ b/services/tests/performancehinttests/TEST_MAPPING
@@ -0,0 +1,19 @@
+{
+  "ravenwood-postsubmit": [
+    {
+      "name": "PerformanceHintTestsRavenwood",
+      "host": true,
+      "options": [
+        {"exclude-annotation": "android.platform.test.annotations.DisabledOnRavenwood"}
+      ]
+    }
+  ],
+  "postsubmit": [
+    {
+      "name": "PerformanceHintTests",
+      "options": [
+        {"exclude-annotation": "org.junit.Ignore"}
+      ]
+    }
+  ]
+}
diff --git a/services/tests/servicestests/src/com/android/server/power/hint/HintManagerServiceTest.java b/services/tests/performancehinttests/src/com/android/server/power/hint/HintManagerServiceTest.java
similarity index 98%
rename from services/tests/servicestests/src/com/android/server/power/hint/HintManagerServiceTest.java
rename to services/tests/performancehinttests/src/com/android/server/power/hint/HintManagerServiceTest.java
index 1decd36..7d04470 100644
--- a/services/tests/servicestests/src/com/android/server/power/hint/HintManagerServiceTest.java
+++ b/services/tests/performancehinttests/src/com/android/server/power/hint/HintManagerServiceTest.java
@@ -596,7 +596,7 @@
                 ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND, 0, 0);
         service.mUidObserver.onUidStateChanged(UID,
                 ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND, 0, 0);
-        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(500) + TimeUnit.MILLISECONDS.toNanos(
+        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1000) + TimeUnit.MILLISECONDS.toNanos(
                 CLEAN_UP_UID_DELAY_MILLIS));
         verify(mNativeWrapperMock, never()).halSetThreads(eq(sessionPtr1), any());
         reset(mNativeWrapperMock);
@@ -653,7 +653,7 @@
                 ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND, 0, 0);
         service.mUidObserver.onUidStateChanged(UID,
                 ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND, 0, 0);
-        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(500) + TimeUnit.MILLISECONDS.toNanos(
+        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1000) + TimeUnit.MILLISECONDS.toNanos(
                 CLEAN_UP_UID_DELAY_MILLIS));
         verify(mNativeWrapperMock, never()).halSetThreads(eq(sessionPtr1), any());
         verify(mNativeWrapperMock, never()).halSetThreads(eq(sessionPtr2), any());
@@ -666,7 +666,7 @@
         service.mUidObserver.onUidStateChanged(UID,
                 ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND, 0, 0);
         // wait for the async uid state change to trigger resume and setThreads
-        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(500));
+        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1000));
         verify(mNativeWrapperMock, times(1)).halSetThreads(eq(sessionPtr2), eq(expectedTids2));
         reset(mNativeWrapperMock);
 
@@ -675,7 +675,7 @@
         LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(100));
         service.mUidObserver.onUidStateChanged(UID,
                 ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND, 0, 0);
-        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(500) + TimeUnit.MILLISECONDS.toNanos(
+        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1000) + TimeUnit.MILLISECONDS.toNanos(
                 CLEAN_UP_UID_DELAY_MILLIS));
         verify(mNativeWrapperMock, times(1)).halPauseHintSession(eq(sessionPtr1));
         verify(mNativeWrapperMock, never()).halSetThreads(eq(sessionPtr1), any());
@@ -684,7 +684,7 @@
         verifyAllHintsEnabled(session2, false);
         service.mUidObserver.onUidStateChanged(UID,
                 ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND, 0, 0);
-        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(500));
+        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1000));
         verifyAllHintsEnabled(session1, false);
         verifyAllHintsEnabled(session2, true);
         reset(mNativeWrapperMock);
@@ -705,7 +705,7 @@
         LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(100));
         service.mUidObserver.onUidStateChanged(UID,
                 ActivityManager.PROCESS_STATE_IMPORTANT_BACKGROUND, 0, 0);
-        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(500) + TimeUnit.MILLISECONDS.toNanos(
+        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1000) + TimeUnit.MILLISECONDS.toNanos(
                 CLEAN_UP_UID_DELAY_MILLIS));
         verify(mNativeWrapperMock, times(1)).halPauseHintSession(eq(sessionPtr1));
         verify(mNativeWrapperMock, never()).halSetThreads(eq(sessionPtr1), any());
@@ -721,7 +721,7 @@
 
         service.mUidObserver.onUidStateChanged(UID,
                 ActivityManager.PROCESS_STATE_IMPORTANT_FOREGROUND, 0, 0);
-        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(500) + TimeUnit.MILLISECONDS.toNanos(
+        LockSupport.parkNanos(TimeUnit.MILLISECONDS.toNanos(1000) + TimeUnit.MILLISECONDS.toNanos(
                 CLEAN_UP_UID_DELAY_MILLIS));
         verify(mNativeWrapperMock, times(1)).halSetThreads(eq(sessionPtr1),
                 eq(SESSION_TIDS_A));
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java
new file mode 100644
index 0000000..8d2849b
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java
@@ -0,0 +1,183 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.when;
+
+import android.hardware.power.stats.EnergyConsumerType;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.platform.test.ravenwood.RavenwoodRule;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.function.IntSupplier;
+
+public class AmbientDisplayPowerStatsProcessorTest {
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setNumDisplays(2)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 0, 180.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 1, 360.0);
+
+    private static final double PRECISION = 0.1;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private ScreenPowerStatsCollector.ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
+
+    private final Injector mInjector = new Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return new PowerStatsUidResolver();
+        }
+
+        @Override
+        public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
+            return 0;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> VOLTAGE_MV;
+        }
+
+        @Override
+        public int getDisplayCount() {
+            return 2;
+        }
+
+        @Override
+        public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
+            return mScreenUsageTimeRetriever;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    @Test
+    public void processPowerStats() {
+        PowerComponentAggregatedPowerStats stats = collectAndAggregatePowerStats();
+
+        assertPowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 16.2);
+        assertPowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 5.4);
+        assertPowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_ON, 0);
+        assertPowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_ON, 0);
+    }
+
+    private PowerComponentAggregatedPowerStats collectAndAggregatePowerStats() {
+        ScreenPowerStatsProcessor screenPowerStatsProcessor =
+                new ScreenPowerStatsProcessor(mStatsRule.getPowerProfile());
+        AmbientDisplayPowerStatsProcessor ambientDisplayPowerStatsProcessor =
+                new AmbientDisplayPowerStatsProcessor();
+
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                .trackUidStates(STATE_POWER, STATE_SCREEN)
+                .setProcessor(screenPowerStatsProcessor);
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .setProcessor(ambientDisplayPowerStatsProcessor);
+
+        AggregatedPowerStats stats = new AggregatedPowerStats(config);
+
+        stats.setDeviceState(STATE_POWER, POWER_STATE_OTHER, 0);
+        stats.setDeviceState(STATE_SCREEN, SCREEN_STATE_OTHER, 0);
+
+        ScreenPowerStatsCollector collector = new ScreenPowerStatsCollector(mInjector);
+        collector.setEnabled(true);
+
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
+                .thenReturn(new int[0]);
+
+        stats.addPowerStats(collector.collectStats(), 1000);
+
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
+                .thenReturn(60_000L);
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(1))
+                .thenReturn(120_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(0))
+                .thenReturn(180_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(1))
+                .thenReturn(240_000L);
+        stats.setDeviceState(STATE_POWER, POWER_STATE_BATTERY, 101_000);
+        stats.setDeviceState(STATE_SCREEN, SCREEN_STATE_ON, 401_000);
+
+        // Slightly larger than 600_000 total screen time, to simulate a sight race
+        // between state changes and power stats collection
+        stats.addPowerStats(collector.collectStats(), 612_000);
+
+        stats.finish(612_000);
+
+        return stats.getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY);
+    }
+
+    private void assertPowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats,
+            int powerState, int screenState, double expectedPowerEstimate) {
+        PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
+        PowerStatsLayout layout = new PowerStatsLayout(descriptor);
+        long[] stats = new long[descriptor.statsArrayLength];
+        aggregatedStats.getDeviceStats(stats, new int[]{powerState, screenState});
+        assertThat(layout.getDevicePowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedPowerEstimate);
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
new file mode 100644
index 0000000..817fdcb
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
@@ -0,0 +1,207 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.when;
+
+import android.hardware.power.stats.EnergyConsumerType;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.Handler;
+import android.platform.test.ravenwood.RavenwoodRule;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.function.IntSupplier;
+
+public class ScreenPowerStatsCollectorTest {
+    private static final int APP_UID1 = 42;
+    private static final int APP_UID2 = 24;
+    private static final int ISOLATED_UID = 99123;
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_SCREEN, 1000);
+
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+    @Mock
+    private ScreenPowerStatsCollector.ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
+
+    private final Injector mInjector = new Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return mPowerStatsUidResolver;
+        }
+
+        @Override
+        public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
+            return 0;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> 3500;
+        }
+
+        @Override
+        public int getDisplayCount() {
+            return 2;
+        }
+
+        @Override
+        public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
+            return mScreenUsageTimeRetriever;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        when(mPowerStatsUidResolver.mapUid(anyInt())).thenAnswer(invocation -> {
+            int uid = invocation.getArgument(0);
+            if (uid == ISOLATED_UID) {
+                return APP_UID2;
+            } else {
+                return uid;
+            }
+        });
+    }
+
+    @Test
+    public void collectStats() {
+        ScreenPowerStatsCollector collector = new ScreenPowerStatsCollector(mInjector);
+        collector.setEnabled(true);
+
+        // Establish a baseline
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
+                .thenReturn(new int[]{77});
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
+                .thenReturn(new long[]{10_000});
+
+        doAnswer(inv -> {
+            ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
+                    inv.getArgument(0);
+            callback.onUidTopActivityTime(APP_UID1, 1000);
+            callback.onUidTopActivityTime(APP_UID2, 2000);
+            return null;
+        }).when(mScreenUsageTimeRetriever).retrieveTopActivityTimes(any(
+                ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback.class));
+
+        collector.collectStats();
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
+                .thenReturn(new long[]{45_000});
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
+                .thenReturn(60_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_DARK))
+                .thenReturn(10_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_MEDIUM))
+                .thenReturn(20_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_BRIGHT))
+                .thenReturn(30_000L);
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(1))
+                .thenReturn(120_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(0))
+                .thenReturn(180_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(1))
+                .thenReturn(240_000L);
+        doAnswer(inv -> {
+            ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
+                    inv.getArgument(0);
+            callback.onUidTopActivityTime(APP_UID1, 3000);
+            callback.onUidTopActivityTime(APP_UID2, 5000);
+            callback.onUidTopActivityTime(ISOLATED_UID, 7000);
+            return null;
+        }).when(mScreenUsageTimeRetriever).retrieveTopActivityTimes(any(
+                ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback.class));
+
+
+        PowerStats powerStats = collector.collectStats();
+
+        ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout();
+        layout.fromExtras(powerStats.descriptor.extras);
+
+        // (45000 - 10000) / 3500
+        assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
+                .isEqualTo(10_000);
+
+        assertThat(layout.getScreenOnDuration(powerStats.stats, 0))
+                .isEqualTo(60_000);
+        assertThat(layout.getBrightnessLevelDuration(powerStats.stats, 0,
+                BatteryStats.SCREEN_BRIGHTNESS_DARK))
+                .isEqualTo(10_000);
+        assertThat(layout.getBrightnessLevelDuration(powerStats.stats, 0,
+                BatteryStats.SCREEN_BRIGHTNESS_MEDIUM))
+                .isEqualTo(20_000);
+        assertThat(layout.getBrightnessLevelDuration(powerStats.stats, 0,
+                BatteryStats.SCREEN_BRIGHTNESS_BRIGHT))
+                .isEqualTo(30_000);
+        assertThat(layout.getScreenOnDuration(powerStats.stats, 1))
+                .isEqualTo(120_000);
+        assertThat(layout.getScreenDozeDuration(powerStats.stats, 0))
+                .isEqualTo(180_000);
+        assertThat(layout.getScreenDozeDuration(powerStats.stats, 1))
+                .isEqualTo(240_000);
+
+        assertThat(powerStats.uidStats.size()).isEqualTo(2);
+        // 3000 - 1000
+        assertThat(layout.getUidTopActivityDuration(powerStats.uidStats.get(APP_UID1)))
+                .isEqualTo(2000);
+        // (5000 - 2000) + 7000
+        assertThat(layout.getUidTopActivityDuration(powerStats.uidStats.get(APP_UID2)))
+                .isEqualTo(10000);
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java
new file mode 100644
index 0000000..9fde61a
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java
@@ -0,0 +1,287 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import static android.os.BatteryConsumer.PROCESS_STATE_ANY;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.hardware.power.stats.EnergyConsumerType;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.Handler;
+import android.os.Process;
+import android.platform.test.ravenwood.RavenwoodRule;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.function.IntSupplier;
+
+public class ScreenPowerStatsProcessorTest {
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setNumDisplays(2)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 0, 180.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 1, 360.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON, 0, 480.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON, 1, 720.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_FULL, 0, 4800.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON, 1, 7200.0)
+            .initMeasuredEnergyStatsLocked();
+
+    private static final double PRECISION = 0.1;
+    private static final int APP_UID1 = Process.FIRST_APPLICATION_UID + 42;
+    private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 101;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private ScreenPowerStatsCollector.ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
+
+    private final Injector mInjector = new Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return new PowerStatsUidResolver();
+        }
+
+        @Override
+        public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
+            return 0;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> VOLTAGE_MV;
+        }
+
+        @Override
+        public int getDisplayCount() {
+            return 2;
+        }
+
+        @Override
+        public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
+            return mScreenUsageTimeRetriever;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    @Test
+    public void processPowerStats_powerProfile() {
+        PowerComponentAggregatedPowerStats stats = collectAndAggregatePowerStats(false);
+
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_ON, 195.5, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0, 0.6);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_ON, 97.8, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0, 0);
+
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_ON, 78.2);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_ON, 39.1);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_ON, 117.3);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_ON, 58.7);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+    }
+
+    @Test
+    public void processPowerStats_energyConsumer() {
+        PowerComponentAggregatedPowerStats stats = collectAndAggregatePowerStats(true);
+
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_ON, 261.9, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0, 7.2);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_ON, 130.9, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0, 0);
+
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_ON, 104.8);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_ON, 52.4);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_ON, 157.1);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_ON, 78.6);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+    }
+
+    private PowerComponentAggregatedPowerStats collectAndAggregatePowerStats(
+            boolean energyConsumer) {
+        ScreenPowerStatsProcessor processor =
+                new ScreenPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        ScreenPowerStatsCollector collector = new ScreenPowerStatsCollector(mInjector);
+        collector.setEnabled(true);
+
+        if (energyConsumer) {
+            when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
+                    .thenReturn(new int[]{77});
+            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
+                    .thenReturn(new long[]{10_000});
+        } else {
+            when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
+                    .thenReturn(new int[0]);
+        }
+
+        doAnswer(inv -> {
+            ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
+                    inv.getArgument(0);
+            callback.onUidTopActivityTime(APP_UID1, 1000);
+            callback.onUidTopActivityTime(APP_UID2, 2000);
+            return null;
+        }).when(mScreenUsageTimeRetriever).retrieveTopActivityTimes(any(
+                ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback.class));
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 1000);
+
+        if (energyConsumer) {
+            // 400 mAh represented as microWattSeconds
+            long energyUws = 400L * 3600 * VOLTAGE_MV;
+            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
+                    .thenReturn(new long[]{10_000 + energyUws});
+        }
+
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
+                .thenReturn(60_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_DARK))
+                .thenReturn(10_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_MEDIUM))
+                .thenReturn(20_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_BRIGHT))
+                .thenReturn(30_000L);
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(1))
+                .thenReturn(120_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(0))
+                .thenReturn(180_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(1))
+                .thenReturn(240_000L);
+        doAnswer(inv -> {
+            ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
+                    inv.getArgument(0);
+            callback.onUidTopActivityTime(APP_UID1, 3000);
+            callback.onUidTopActivityTime(APP_UID2, 5000);
+            return null;
+        }).when(mScreenUsageTimeRetriever).retrieveTopActivityTimes(any(
+                ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback.class));
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_BATTERY, 201_000);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 601_000);
+
+        // Slightly larger than 600_000 total screen time, to simulate a sight race
+        // between state changes and power stats collection
+        aggregatedStats.addPowerStats(collector.collectStats(), 612_000);
+
+        aggregatedStats.getConfig().getProcessor().finish(aggregatedStats, 180_000);
+        return aggregatedStats;
+    }
+
+    private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
+            ScreenPowerStatsProcessor processor) {
+        AggregatedPowerStatsConfig.PowerComponent config =
+                new AggregatedPowerStatsConfig.PowerComponent(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                        .trackUidStates(STATE_POWER, STATE_SCREEN)
+                        .setProcessor(processor);
+
+        PowerComponentAggregatedPowerStats aggregatedStats =
+                new PowerComponentAggregatedPowerStats(
+                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+
+        return aggregatedStats;
+    }
+
+    private void assertDevicePowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats,
+            int powerState, int screenState, double expectedScreenPowerEstimate,
+            double expectedDozePowerEstimate) {
+        PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
+        ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout(descriptor);
+        long[] stats = new long[descriptor.statsArrayLength];
+        aggregatedStats.getDeviceStats(stats, new int[]{powerState, screenState});
+        assertThat(layout.getDevicePowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedScreenPowerEstimate);
+        assertThat(layout.getScreenDozePowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedDozePowerEstimate);
+    }
+
+    private void assertUidPowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats, int uid,
+            int powerState, int screenState, double expectedScreenPowerEstimate) {
+        PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
+        ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout(descriptor);
+        long[] stats = new long[descriptor.uidStatsArrayLength];
+        aggregatedStats.getUidStats(stats, uid,
+                new int[]{powerState, screenState, PROCESS_STATE_ANY});
+        assertThat(layout.getUidPowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedScreenPowerEstimate);
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/SensorPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/SensorPowerStatsProcessorTest.java
new file mode 100644
index 0000000..7000487
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/SensorPowerStatsProcessorTest.java
@@ -0,0 +1,241 @@
+/*
+ * Copyright (C) 2024 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.server.power.stats;
+
+import static android.os.BatteryConsumer.PROCESS_STATE_BACKGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_CACHED;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.hardware.Sensor;
+import android.hardware.SensorManager;
+import android.hardware.input.InputSensorInfo;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.Process;
+import android.platform.test.ravenwood.RavenwoodRule;
+
+import com.android.internal.os.MonotonicClock;
+import com.android.internal.os.PowerStats;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.List;
+
+public class SensorPowerStatsProcessorTest {
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .initMeasuredEnergyStatsLocked();
+
+    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 + 101;
+    private static final int SENSOR_HANDLE_1 = 77;
+    private static final int SENSOR_HANDLE_2 = 88;
+    private static final int SENSOR_HANDLE_3 = 99;
+
+    @Mock
+    private SensorManager mSensorManager;
+
+    private MonotonicClock mMonotonicClock;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        mMonotonicClock = new MonotonicClock(0, mStatsRule.getMockClock());
+        Sensor sensor1 = createSensor(SENSOR_HANDLE_1, Sensor.TYPE_STEP_COUNTER,
+                Sensor.STRING_TYPE_STEP_COUNTER, "dancing", 100);
+        Sensor sensor2 = createSensor(SENSOR_HANDLE_2, Sensor.TYPE_MOTION_DETECT,
+                "com.example", "tango", 200);
+        Sensor sensor3 = createSensor(SENSOR_HANDLE_3, Sensor.TYPE_MOTION_DETECT,
+                "com.example", "waltz", 300);
+        when(mSensorManager.getSensorList(Sensor.TYPE_ALL)).thenReturn(
+                List.of(sensor1, sensor2, sensor3));
+    }
+
+    @Test
+    public void testPowerEstimation() {
+        SensorPowerStatsProcessor processor = new SensorPowerStatsProcessor(() -> mSensorManager);
+
+        PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(processor);
+
+        processor.noteStateChange(stats, buildHistoryItem(0, true, APP_UID1, SENSOR_HANDLE_1));
+
+        // Turn the screen off after 2.5 seconds
+        stats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        stats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        stats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE, 5000);
+
+        processor.noteStateChange(stats, buildHistoryItem(6000, false, APP_UID1, SENSOR_HANDLE_1));
+        processor.noteStateChange(stats, buildHistoryItem(7000, true, APP_UID2, SENSOR_HANDLE_1));
+        processor.noteStateChange(stats, buildHistoryItem(8000, true, APP_UID2, SENSOR_HANDLE_2));
+        processor.noteStateChange(stats, buildHistoryItem(9000, false, APP_UID2, SENSOR_HANDLE_1));
+
+        processor.finish(stats, 10000);
+
+        PowerStats.Descriptor descriptor = stats.getPowerStatsDescriptor();
+        SensorPowerStatsLayout statsLayout = new SensorPowerStatsLayout();
+        statsLayout.fromExtras(descriptor.extras);
+
+        String dump = stats.toString();
+        assertThat(dump).contains(" step_counter: ");
+        assertThat(dump).contains(" com.example.tango: ");
+
+        long[] uidStats = new long[descriptor.uidStatsArrayLength];
+
+        // For UID1:
+        // SENSOR1 was on for 6000 ms.
+        //   Estimated power: 6000 * 100 = 0.167 mAh
+        //     split between three different states
+        //          fg screen-on: 6000 * 2500/10000
+        //          bg screen-off: 6000 * 2500/10000
+        //          fgs screen-off: 6000 * 5000/10000
+        double expectedPower1 = 0.166666;
+        stats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 2500 / 10000);
+        stats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 2500 / 10000);
+        stats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 5000 / 10000);
+
+        // For UID2:
+        // SENSOR1 was on for 2000 ms.
+        //   Estimated power: 2000 * 100 = 0.0556 mAh
+        //     split between three different states
+        //          cached screen-on: 2000 * 2500/10000
+        //          cached screen-off: 2000 * 7500/10000
+        // SENSOR2 was on for 2000 ms.
+        //   Estimated power: 2000 * 200 = 0.11111 mAh
+        //     split between three different states
+        //          cached screen-on: 2000 * 2500/10000
+        //          cached screen-off: 2000 * 7500/10000
+        double expectedPower2 = 0.05555 + 0.11111;
+        stats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 2500 / 10000);
+        stats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 7500 / 10000);
+
+        long[] deviceStats = new long[descriptor.statsArrayLength];
+
+        stats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of((expectedPower1 + expectedPower2) * 2500 / 10000);
+
+        stats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of((expectedPower1 + expectedPower2) * 7500 / 10000);
+    }
+
+    private BatteryStats.HistoryItem buildHistoryItem(int timestamp, boolean stateOn,
+            int uid, int sensor) {
+        mStatsRule.setTime(timestamp, timestamp);
+        BatteryStats.HistoryItem historyItem = new BatteryStats.HistoryItem();
+        historyItem.time = mMonotonicClock.monotonicTime();
+        historyItem.states = stateOn ? BatteryStats.HistoryItem.STATE_SENSOR_ON_FLAG : 0;
+        if (stateOn) {
+            historyItem.eventCode = BatteryStats.HistoryItem.EVENT_STATE_CHANGE
+                    | BatteryStats.HistoryItem.EVENT_FLAG_START;
+        } else {
+            historyItem.eventCode = BatteryStats.HistoryItem.EVENT_STATE_CHANGE
+                    | BatteryStats.HistoryItem.EVENT_FLAG_FINISH;
+        }
+        historyItem.eventTag = historyItem.localEventTag;
+        historyItem.eventTag.uid = uid;
+        historyItem.eventTag.string = "sensor:0x" + Integer.toHexString(sensor);
+        return historyItem;
+    }
+
+    private int[] states(int... states) {
+        return states;
+    }
+
+    private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
+            SensorPowerStatsProcessor processor) {
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SENSORS)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessor(processor);
+
+        AggregatedPowerStats aggregatedPowerStats = new AggregatedPowerStats(config);
+        PowerComponentAggregatedPowerStats powerComponentStats =
+                aggregatedPowerStats.getPowerComponentStats(
+                        BatteryConsumer.POWER_COMPONENT_SENSORS);
+        processor.start(powerComponentStats, 0);
+
+        powerComponentStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
+        powerComponentStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+        powerComponentStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND, 0);
+        powerComponentStats.setUidState(APP_UID2, STATE_PROCESS_STATE, PROCESS_STATE_CACHED, 0);
+
+        return powerComponentStats;
+    }
+
+    private Sensor createSensor(int handle, int type, String stringType, String name, float power) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            Sensor sensor = mock(Sensor.class);
+            when(sensor.getHandle()).thenReturn(handle);
+            when(sensor.getType()).thenReturn(type);
+            when(sensor.getStringType()).thenReturn(stringType);
+            when(sensor.getName()).thenReturn(name);
+            when(sensor.getPower()).thenReturn(power);
+            return sensor;
+        } else {
+            return new Sensor(new InputSensorInfo(name, "vendor", 0 /* version */,
+                    handle, type, 100.0f /*maxRange */, 0.02f /* resolution */,
+                    (float) power, 1000 /* minDelay */, 0 /* fifoReservedEventCount */,
+                    0 /* fifoMaxEventCount */, stringType /* stringType */,
+                    "" /* requiredPermission */, 0 /* maxDelay */, 0 /* flags */, 0 /* id */));
+        }
+    }
+}
diff --git a/services/tests/servicestests/Android.bp b/services/tests/servicestests/Android.bp
index a888dad..a86289b 100644
--- a/services/tests/servicestests/Android.bp
+++ b/services/tests/servicestests/Android.bp
@@ -78,7 +78,6 @@
         // TODO: remove once Android migrates to JUnit 4.12,
         // which provides assertThrows
         "testng",
-        "truth",
         "junit",
         "junit-params",
         "ActivityContext",
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManagerTest.java b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManagerTest.java
index f92eaab..c1b3929 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManagerTest.java
@@ -18,13 +18,13 @@
 
 import static com.android.server.accessibility.Flags.FLAG_ENABLE_A11Y_CHECKER_LOGGING;
 import static com.android.server.accessibility.a11ychecker.AccessibilityCheckerConstants.MIN_DURATION_BETWEEN_CHECKS;
+import static com.android.server.accessibility.a11ychecker.TestUtils.QUALIFIED_TEST_ACTIVITY_NAME;
 import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_A11Y_SERVICE_CLASS_NAME;
 import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_A11Y_SERVICE_SOURCE_PACKAGE_NAME;
 import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_ACTIVITY_NAME;
 import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_DEFAULT_BROWSER;
 import static com.android.server.accessibility.a11ychecker.TestUtils.createAtom;
 import static com.android.server.accessibility.a11ychecker.TestUtils.getMockPackageManagerWithInstalledApps;
-import static com.android.server.accessibility.a11ychecker.TestUtils.getTestAccessibilityEvent;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -114,7 +114,7 @@
 
         Set<A11yCheckerProto.AccessibilityCheckResultReported> results =
                 mAccessibilityCheckerManager.maybeRunA11yChecker(
-                        List.of(mockNodeInfo1, mockNodeInfo2), getTestAccessibilityEvent(),
+                        List.of(mockNodeInfo1, mockNodeInfo2), QUALIFIED_TEST_ACTIVITY_NAME,
                         new ComponentName(TEST_A11Y_SERVICE_SOURCE_PACKAGE_NAME,
                                 TEST_A11Y_SERVICE_CLASS_NAME), /*userId=*/ 0);
 
@@ -139,7 +139,7 @@
 
         Set<A11yCheckerProto.AccessibilityCheckResultReported> results =
                 mAccessibilityCheckerManager.maybeRunA11yChecker(
-                        List.of(mockNodeInfo), getTestAccessibilityEvent(),
+                        List.of(mockNodeInfo), QUALIFIED_TEST_ACTIVITY_NAME,
                         new ComponentName(TEST_A11Y_SERVICE_SOURCE_PACKAGE_NAME,
                                 TEST_A11Y_SERVICE_CLASS_NAME), /*userId=*/ 0);
 
@@ -160,7 +160,7 @@
 
         Set<A11yCheckerProto.AccessibilityCheckResultReported> results =
                 mAccessibilityCheckerManager.maybeRunA11yChecker(
-                        List.of(mockNodeInfo, mockNodeInfoDuplicate), getTestAccessibilityEvent(),
+                        List.of(mockNodeInfo, mockNodeInfoDuplicate), QUALIFIED_TEST_ACTIVITY_NAME,
                         new ComponentName(TEST_A11Y_SERVICE_SOURCE_PACKAGE_NAME,
                                 TEST_A11Y_SERVICE_CLASS_NAME), /*userId=*/ 0);
 
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java
index 141f174..5b4e72e 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java
@@ -16,11 +16,13 @@
 
 package com.android.server.accessibility.a11ychecker;
 
+import static com.android.server.accessibility.a11ychecker.TestUtils.QUALIFIED_TEST_ACTIVITY_NAME;
 import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_A11Y_SERVICE_CLASS_NAME;
 import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_A11Y_SERVICE_SOURCE_PACKAGE_NAME;
+import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_ACTIVITY_NAME;
+import static com.android.server.accessibility.a11ychecker.TestUtils.TEST_APP_PACKAGE_NAME;
 import static com.android.server.accessibility.a11ychecker.TestUtils.createAtom;
 import static com.android.server.accessibility.a11ychecker.TestUtils.getMockPackageManagerWithInstalledApps;
-import static com.android.server.accessibility.a11ychecker.TestUtils.getTestAccessibilityEvent;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -28,7 +30,6 @@
 
 import android.content.ComponentName;
 import android.content.pm.PackageManager;
-import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.AccessibilityNodeInfo;
 
 import androidx.test.runner.AndroidJUnit4;
@@ -138,11 +139,15 @@
     }
 
     @Test
-    public void getActivityName_hasWindowStateChangedEvent_returnsActivityName() {
-        AccessibilityEvent accessibilityEvent = getTestAccessibilityEvent();
-
+    public void getActivityName_hasValidActivityClassName_returnsActivityName() {
         assertThat(AccessibilityCheckerUtils.getActivityName(mMockPackageManager,
-                accessibilityEvent)).isEqualTo("MainActivity");
+                TEST_APP_PACKAGE_NAME, QUALIFIED_TEST_ACTIVITY_NAME)).isEqualTo(TEST_ACTIVITY_NAME);
+    }
+
+    @Test
+    public void getActivityName_hasInvalidActivityClassName_returnsActivityName() {
+        assertThat(AccessibilityCheckerUtils.getActivityName(mMockPackageManager,
+                TEST_APP_PACKAGE_NAME, "com.NonActivityClass")).isEmpty();
     }
 
     // Makes sure the AccessibilityHierarchyCheck class to enum mapping is up to date with the
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/TestUtils.java b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/TestUtils.java
index ec1a255..acf64b6 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/TestUtils.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/TestUtils.java
@@ -16,6 +16,7 @@
 
 package com.android.server.accessibility.a11ychecker;
 
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.when;
 
@@ -26,6 +27,7 @@
 import android.content.pm.PackageManager;
 import android.view.accessibility.AccessibilityEvent;
 
+import org.mockito.AdditionalMatchers;
 import org.mockito.Mockito;
 
 public class TestUtils {
@@ -46,8 +48,11 @@
         ComponentName testActivityComponentName = new ComponentName(TEST_APP_PACKAGE_NAME,
                 QUALIFIED_TEST_ACTIVITY_NAME);
 
-        when(mockPackageManager.getActivityInfo(testActivityComponentName, 0))
+        when(mockPackageManager.getActivityInfo(eq(testActivityComponentName), eq(0)))
                 .thenReturn(testActivityInfo);
+        when(mockPackageManager.getActivityInfo(
+                AdditionalMatchers.not(eq(testActivityComponentName)), eq(0)))
+                .thenThrow(PackageManager.NameNotFoundException.class);
         when(mockPackageManager.getPackageInfo(TEST_APP_PACKAGE_NAME, 0))
                 .thenReturn(createPackageInfo(TEST_APP_PACKAGE_NAME, TEST_APP_VERSION_CODE,
                         testActivityInfo));
diff --git a/services/tests/servicestests/src/com/android/server/appop/DiscreteAppOpPersistenceTest.java b/services/tests/servicestests/src/com/android/server/appop/DiscreteAppOpPersistenceTest.java
new file mode 100644
index 0000000..bc3a5ca
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/appop/DiscreteAppOpPersistenceTest.java
@@ -0,0 +1,168 @@
+/*
+ * Copyright (C) 2024 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.server.appop;
+
+import static android.app.AppOpsManager.ATTRIBUTION_FLAG_ACCESSOR;
+import static android.app.AppOpsManager.ATTRIBUTION_FLAG_RECEIVER;
+import static android.app.AppOpsManager.OP_FLAGS_ALL_TRUSTED;
+import static android.app.AppOpsManager.OP_FLAG_SELF;
+import static android.app.AppOpsManager.UID_STATE_FOREGROUND;
+import static android.app.AppOpsManager.UID_STATE_FOREGROUND_SERVICE;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.AppOpsManager;
+import android.companion.virtual.VirtualDeviceManager;
+import android.content.Context;
+import android.os.FileUtils;
+import android.os.Process;
+import android.permission.flags.Flags;
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.List;
+
+@RunWith(AndroidJUnit4.class)
+public class DiscreteAppOpPersistenceTest {
+    private DiscreteRegistry mDiscreteRegistry;
+    private final Object mLock = new Object();
+    private File mMockDataDirectory;
+    private final Context mContext =
+            InstrumentationRegistry.getInstrumentation().getTargetContext();
+
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
+    @Before
+    public void setUp() {
+        mMockDataDirectory = mContext.getDir("mock_data", Context.MODE_PRIVATE);
+        mDiscreteRegistry = new DiscreteRegistry(mLock, mMockDataDirectory);
+        mDiscreteRegistry.systemReady();
+    }
+
+    @After
+    public void cleanUp() {
+        mDiscreteRegistry.writeAndClearAccessHistory();
+        FileUtils.deleteContents(mMockDataDirectory);
+    }
+
+    @RequiresFlagsEnabled(Flags.FLAG_DEVICE_AWARE_APP_OP_NEW_SCHEMA_ENABLED)
+    @Test
+    public void defaultDevice_recordAccess_persistToDisk() {
+        int uid = Process.myUid();
+        String packageName = mContext.getOpPackageName();
+        int op = AppOpsManager.OP_CAMERA;
+        String deviceId = VirtualDeviceManager.PERSISTENT_DEVICE_ID_DEFAULT;
+        long accessTime = System.currentTimeMillis();
+        long duration = 60000L;
+        int uidState = UID_STATE_FOREGROUND;
+        int opFlags = OP_FLAGS_ALL_TRUSTED;
+        int attributionFlags = ATTRIBUTION_FLAG_ACCESSOR;
+        int attributionChainId = AppOpsManager.ATTRIBUTION_CHAIN_ID_NONE;
+
+        mDiscreteRegistry.recordDiscreteAccess(uid, packageName, deviceId, op, null, opFlags,
+                uidState, accessTime, duration, attributionFlags, attributionChainId);
+
+        // Verify in-memory object is correct
+        fetchDiscreteOpsAndValidate(uid, packageName, op, deviceId, null, accessTime,
+                duration, uidState, opFlags, attributionFlags, attributionChainId);
+
+        // Write to disk and clear the in-memory object
+        mDiscreteRegistry.writeAndClearAccessHistory();
+
+        // Verify the storage file is created and then verify its content is correct
+        File[] files = FileUtils.listFilesOrEmpty(mMockDataDirectory);
+        assertThat(files.length).isEqualTo(1);
+        fetchDiscreteOpsAndValidate(uid, packageName, op, deviceId, null, accessTime,
+                duration, uidState, opFlags, attributionFlags, attributionChainId);
+    }
+
+    @RequiresFlagsEnabled(Flags.FLAG_DEVICE_AWARE_APP_OP_NEW_SCHEMA_ENABLED)
+    @Test
+    public void externalDevice_recordAccess_persistToDisk() {
+        int uid = Process.myUid();
+        String packageName = mContext.getOpPackageName();
+        int op = AppOpsManager.OP_CAMERA;
+        String deviceId = "companion:1";
+        long accessTime = System.currentTimeMillis();
+        long duration = -1;
+        int uidState = UID_STATE_FOREGROUND_SERVICE;
+        int opFlags = OP_FLAG_SELF;
+        int attributionFlags = ATTRIBUTION_FLAG_RECEIVER;
+        int attributionChainId = 10;
+
+        mDiscreteRegistry.recordDiscreteAccess(uid, packageName, deviceId, op, null, opFlags,
+                uidState, accessTime, duration, attributionFlags, attributionChainId);
+
+        fetchDiscreteOpsAndValidate(uid, packageName, op, deviceId, null, accessTime,
+                duration, uidState, opFlags, attributionFlags, attributionChainId);
+
+        mDiscreteRegistry.writeAndClearAccessHistory();
+
+        File[] files = FileUtils.listFilesOrEmpty(mMockDataDirectory);
+        assertThat(files.length).isEqualTo(1);
+        fetchDiscreteOpsAndValidate(uid, packageName, op, deviceId, null, accessTime,
+                duration, uidState, opFlags, attributionFlags, attributionChainId);
+    }
+
+    private void fetchDiscreteOpsAndValidate(int expectedUid, String expectedPackageName,
+            int expectedOp, String expectedDeviceId, String expectedAttrTag,
+            long expectedAccessTime, long expectedAccessDuration, int expectedUidState,
+            int expectedOpFlags, int expectedAttrFlags, int expectedAttrChainId) {
+        DiscreteRegistry.DiscreteOps discreteOps = mDiscreteRegistry.getAllDiscreteOps();
+
+        assertThat(discreteOps.isEmpty()).isFalse();
+        assertThat(discreteOps.mUids.size()).isEqualTo(1);
+
+        DiscreteRegistry.DiscreteUidOps discreteUidOps = discreteOps.mUids.get(expectedUid);
+        assertThat(discreteUidOps.mPackages.size()).isEqualTo(1);
+
+        DiscreteRegistry.DiscretePackageOps discretePackageOps =
+                discreteUidOps.mPackages.get(expectedPackageName);
+        assertThat(discretePackageOps.mPackageOps.size()).isEqualTo(1);
+
+        DiscreteRegistry.DiscreteOp discreteOp = discretePackageOps.mPackageOps.get(expectedOp);
+        assertThat(discreteOp.mDeviceAttributedOps.size()).isEqualTo(1);
+
+        DiscreteRegistry.DiscreteDeviceOp discreteDeviceOp =
+                discreteOp.mDeviceAttributedOps.get(expectedDeviceId);
+        assertThat(discreteDeviceOp.mAttributedOps.size()).isEqualTo(1);
+
+        List<DiscreteRegistry.DiscreteOpEvent> discreteOpEvents =
+                discreteDeviceOp.mAttributedOps.get(expectedAttrTag);
+        assertThat(discreteOpEvents.size()).isEqualTo(1);
+
+        DiscreteRegistry.DiscreteOpEvent discreteOpEvent = discreteOpEvents.get(0);
+        assertThat(discreteOpEvent.mNoteTime).isEqualTo(expectedAccessTime);
+        assertThat(discreteOpEvent.mNoteDuration).isEqualTo(expectedAccessDuration);
+        assertThat(discreteOpEvent.mUidState).isEqualTo(expectedUidState);
+        assertThat(discreteOpEvent.mOpFlag).isEqualTo(expectedOpFlags);
+        assertThat(discreteOpEvent.mAttributionFlags).isEqualTo(expectedAttrFlags);
+        assertThat(discreteOpEvent.mAttributionChainId).isEqualTo(expectedAttrChainId);
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/hdmi/HdmiCecLocalDeviceTvTest.java b/services/tests/servicestests/src/com/android/server/hdmi/HdmiCecLocalDeviceTvTest.java
index 87b52e6..f98bbf9 100644
--- a/services/tests/servicestests/src/com/android/server/hdmi/HdmiCecLocalDeviceTvTest.java
+++ b/services/tests/servicestests/src/com/android/server/hdmi/HdmiCecLocalDeviceTvTest.java
@@ -28,6 +28,7 @@
 import static com.android.server.hdmi.HdmiControlService.INITIATED_BY_WAKE_UP_MESSAGE;
 import static com.android.server.hdmi.HdmiControlService.STANDBY_SCREEN_OFF;
 import static com.android.server.hdmi.HdmiControlService.WAKE_UP_SCREEN_ON;
+import static com.android.server.hdmi.RequestActiveSourceAction.TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -1792,7 +1793,7 @@
         mTestLooper.dispatchAll();
 
         // Skip the LauncherX API timeout.
-        mTestLooper.moveTimeForward(HdmiConfig.TIMEOUT_MS * 2);
+        mTestLooper.moveTimeForward(TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS);
         mTestLooper.dispatchAll();
 
         assertThat(mNativeWrapper.getResultMessages()).contains(requestActiveSource);
@@ -1825,7 +1826,7 @@
         mTestLooper.dispatchAll();
 
         // Skip the LauncherX API timeout.
-        mTestLooper.moveTimeForward(HdmiConfig.TIMEOUT_MS * 2);
+        mTestLooper.moveTimeForward(TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS);
         mTestLooper.dispatchAll();
 
         assertThat(mNativeWrapper.getResultMessages()).contains(requestActiveSource);
@@ -1861,7 +1862,7 @@
         mTestLooper.dispatchAll();
 
         // Skip the LauncherX API timeout.
-        mTestLooper.moveTimeForward(HdmiConfig.TIMEOUT_MS * 2);
+        mTestLooper.moveTimeForward(TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS);
         mTestLooper.dispatchAll();
 
         assertThat(mNativeWrapper.getResultMessages()).contains(requestActiveSource);
@@ -1904,7 +1905,7 @@
         mTestLooper.dispatchAll();
 
         // Skip the LauncherX API timeout.
-        mTestLooper.moveTimeForward(HdmiConfig.TIMEOUT_MS * 2);
+        mTestLooper.moveTimeForward(TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS);
         mTestLooper.dispatchAll();
 
         assertThat(mNativeWrapper.getResultMessages()).contains(requestActiveSource);
@@ -1941,7 +1942,7 @@
         mHdmiControlService.sendCecCommand(setStreamPathFromTv);
 
         // Skip the LauncherX API timeout.
-        mTestLooper.moveTimeForward(HdmiConfig.TIMEOUT_MS * 2);
+        mTestLooper.moveTimeForward(TIMEOUT_WAIT_FOR_LAUNCHERX_API_CALL_MS);
         mTestLooper.dispatchAll();
 
         assertThat(mNativeWrapper.getResultMessages()).doesNotContain(requestActiveSource);
diff --git a/services/tests/servicestests/src/com/android/server/power/hint/TEST_MAPPING b/services/tests/servicestests/src/com/android/server/power/hint/TEST_MAPPING
deleted file mode 100644
index 874eec7..0000000
--- a/services/tests/servicestests/src/com/android/server/power/hint/TEST_MAPPING
+++ /dev/null
@@ -1,15 +0,0 @@
-{
-  "presubmit": [
-    {
-      "name": "FrameworksServicesTests",
-      "options": [
-        {
-          "include-filter": "com.android.server.power.hint"
-        },
-        {
-          "exclude-annotation": "androidx.test.filters.FlakyTest"
-        }
-      ]
-    }
-  ]
-}
diff --git a/services/tests/timetests/src/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessorTest.java b/services/tests/timetests/src/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessorTest.java
index f3440f7..ea3b409 100644
--- a/services/tests/timetests/src/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessorTest.java
+++ b/services/tests/timetests/src/com/android/server/timezonedetector/location/ZoneInfoDbTimeZoneProviderEventPreProcessorTest.java
@@ -39,13 +39,23 @@
 
     private static final long ARBITRARY_TIME_MILLIS = 11223344;
 
+    private final List<String> mNonExistingTimeZones = Arrays.asList(
+            "SystemV/HST10", "Atlantic/Atlantis", "EUROPE/LONDON", "Etc/GMT-5:30");
     private final ZoneInfoDbTimeZoneProviderEventPreProcessor mPreProcessor =
             new ZoneInfoDbTimeZoneProviderEventPreProcessor();
 
+    private static final TimeZoneProviderStatus ARBITRARY_TIME_ZONE_PROVIDER_STATUS =
+            new TimeZoneProviderStatus.Builder()
+                    .setConnectivityDependencyStatus(DEPENDENCY_STATUS_OK)
+                    .setLocationDetectionDependencyStatus(DEPENDENCY_STATUS_OK)
+                    .setTimeZoneResolutionOperationStatus(OPERATION_STATUS_OK)
+                    .build();
+
     @Test
     public void timeZoneIdsFromZoneInfoDbAreValid() {
         for (String timeZone : TimeZone.getAvailableIDs()) {
-            TimeZoneProviderEvent event = timeZoneProviderEvent(timeZone);
+            TimeZoneProviderEvent event = timeZoneProviderEvent(timeZone,
+                    ARBITRARY_TIME_ZONE_PROVIDER_STATUS);
             assertWithMessage("Time zone %s should be supported", timeZone)
                     .that(mPreProcessor.preProcess(event)).isEqualTo(event);
         }
@@ -53,11 +63,9 @@
 
     @Test
     public void eventWithNonExistingZones_areMappedToUncertainEvent() {
-        List<String> nonExistingTimeZones = Arrays.asList(
-                "SystemV/HST10", "Atlantic/Atlantis", "EUROPE/LONDON", "Etc/GMT-5:30");
-
-        for (String timeZone : nonExistingTimeZones) {
-            TimeZoneProviderEvent event = timeZoneProviderEvent(timeZone);
+        for (String timeZone : mNonExistingTimeZones) {
+            TimeZoneProviderEvent event = timeZoneProviderEvent(timeZone,
+                    ARBITRARY_TIME_ZONE_PROVIDER_STATUS);
 
             TimeZoneProviderStatus expectedProviderStatus =
                     new TimeZoneProviderStatus.Builder(event.getTimeZoneProviderStatus())
@@ -73,14 +81,31 @@
         }
     }
 
-    private static TimeZoneProviderEvent timeZoneProviderEvent(String... timeZoneIds) {
-        TimeZoneProviderStatus providerStatus = new TimeZoneProviderStatus.Builder()
-                .setLocationDetectionDependencyStatus(DEPENDENCY_STATUS_OK)
-                .setConnectivityDependencyStatus(DEPENDENCY_STATUS_OK)
-                .setTimeZoneResolutionOperationStatus(OPERATION_STATUS_OK)
-                .build();
+    @Test
+    public void eventWithNullProviderStatus_areMappedToUncertainEvent() {
+        for (String timeZone : mNonExistingTimeZones) {
+            TimeZoneProviderEvent eventWithNullStatus = timeZoneProviderEvent(timeZone,
+                    /* providerStatus= */ null);
+
+            TimeZoneProviderStatus expectedProviderStatus =
+                    new TimeZoneProviderStatus.Builder()
+                            .setTimeZoneResolutionOperationStatus(OPERATION_STATUS_FAILED)
+                            .build();
+
+            TimeZoneProviderEvent expectedResultEvent =
+                    TimeZoneProviderEvent.createUncertainEvent(
+                            eventWithNullStatus.getCreationElapsedMillis(),
+                            expectedProviderStatus);
+            assertWithMessage(timeZone + " with null time zone provider status")
+                    .that(mPreProcessor.preProcess(eventWithNullStatus))
+                    .isEqualTo(expectedResultEvent);
+        }
+    }
+
+    private static TimeZoneProviderEvent timeZoneProviderEvent(String timeZoneId,
+            TimeZoneProviderStatus providerStatus) {
         TimeZoneProviderSuggestion suggestion = new TimeZoneProviderSuggestion.Builder()
-                .setTimeZoneIds(Arrays.asList(timeZoneIds))
+                .setTimeZoneIds(Arrays.asList(timeZoneId))
                 .setElapsedRealtimeMillis(ARBITRARY_TIME_MILLIS)
                 .build();
         return TimeZoneProviderEvent.createSuggestionEvent(
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index c1f5a01..1b999e4 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -3105,6 +3105,29 @@
     }
 
     @Test
+    @EnableFlags(android.app.Flags.FLAG_LIFETIME_EXTENSION_REFACTOR)
+    public void testLifetimeExtendedCancelledOnClick() throws Exception {
+        // Adds a lifetime extended notification.
+        final NotificationRecord notif = generateNotificationRecord(mTestNotificationChannel, 1,
+                null, false);
+        notif.getSbn().getNotification().flags =
+                Notification.FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY;
+        mService.addNotification(notif);
+        // Verify that the notification is posted and active.
+        assertThat(mBinderService.getActiveNotifications(mPkg).length).isEqualTo(1);
+
+        // Click the notification.
+        final NotificationVisibility nv = NotificationVisibility.obtain(notif.getKey(), 1, 2, true);
+        mService.mNotificationDelegate.onNotificationClick(mUid, Binder.getCallingPid(),
+                notif.getKey(), nv);
+        waitForIdle();
+
+        // The notification has been cancelled.
+        StatusBarNotification[] notifs = mBinderService.getActiveNotifications(mPkg);
+        assertThat(notifs.length).isEqualTo(0);
+    }
+
+    @Test
     public void testCancelNotificationWithTag_fromApp_cannotCancelFgsChild()
             throws Exception {
         when(mAmi.applyForegroundServiceNotification(
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
index b07940a..d7bae45 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
@@ -1044,6 +1044,8 @@
         AudioManagerInternal mAudioManager = mock(AudioManagerInternal.class);
         mZenModeHelper.mAudioManager = mAudioManager;
         setupZenConfig();
+        mTestableLooper.processAllMessages();
+        reset(mAudioManager);
 
         // Turn manual zen mode on
         mZenModeHelper.setManualZenMode(ZEN_MODE_IMPORTANT_INTERRUPTIONS, null, UPDATE_ORIGIN_APP,
@@ -1063,6 +1065,44 @@
     }
 
     @Test
+    public void testSetConfig_updatesAudioForSequentialChangesToZenMode() {
+        AudioManagerInternal mAudioManager = mock(AudioManagerInternal.class);
+        mZenModeHelper.mAudioManager = mAudioManager;
+        setupZenConfig();
+        mTestableLooper.processAllMessages();
+        reset(mAudioManager);
+
+        // Turn manual zen mode on
+        mZenModeHelper.setManualZenMode(
+                ZEN_MODE_IMPORTANT_INTERRUPTIONS,
+                null,
+                UPDATE_ORIGIN_APP,
+                null,
+                "test",
+                CUSTOM_PKG_UID);
+        mZenModeHelper.setManualZenMode(
+                ZEN_MODE_IMPORTANT_INTERRUPTIONS,
+                null,
+                UPDATE_ORIGIN_APP,
+                null,
+                "test",
+                CUSTOM_PKG_UID);
+
+        // audio manager shouldn't do anything until the handler processes its messages
+        verify(mAudioManager, never()).updateRingerModeAffectedStreamsInternal();
+
+        // now process the looper's messages
+        mTestableLooper.processAllMessages();
+
+        // Expect calls to audio manager
+        verify(mAudioManager, times(2)).updateRingerModeAffectedStreamsInternal();
+        verify(mAudioManager, times(1)).setRingerModeInternal(anyInt(), anyString());
+
+        // called during applyZenToRingerMode(), which should be true since zen changed
+        verify(mAudioManager, atLeastOnce()).getRingerModeInternal();
+    }
+
+    @Test
     public void testParcelConfig() {
         mZenModeHelper.setNotificationPolicy(new Policy(PRIORITY_CATEGORY_EVENTS
                         | PRIORITY_CATEGORY_MESSAGES | PRIORITY_CATEGORY_REPEAT_CALLERS
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibratorControlServiceTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibratorControlServiceTest.java
index 8ca8623..c496bbb 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibratorControlServiceTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibratorControlServiceTest.java
@@ -45,6 +45,11 @@
 import android.os.IBinder;
 import android.os.Process;
 import android.os.test.TestLooper;
+import android.os.vibrator.Flags;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.util.SparseArray;
 
 import androidx.test.InstrumentationRegistry;
@@ -63,7 +68,6 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.TimeUnit;
 
 public class VibratorControlServiceTest {
 
@@ -71,6 +75,8 @@
 
     @Rule
     public MockitoRule rule = MockitoJUnit.rule();
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
 
     @Mock private VibrationScaler mMockVibrationScaler;
     @Mock private PackageManagerInternal mPackageManagerInternalMock;
@@ -98,6 +104,7 @@
         mVibratorControlService = new VibratorControlService(
                 InstrumentationRegistry.getContext(), new VibratorControllerHolder(),
                 mMockVibrationScaler, mVibrationSettings, mStatsLoggerMock, mLock);
+        mFakeVibratorController.setVibratorControlService(mVibratorControlService);
     }
 
     @Test
@@ -280,10 +287,10 @@
         CompletableFuture<Void> future =
                 mVibratorControlService.triggerVibrationParamsRequest(UID, USAGE_RINGTONE,
                         timeoutInMillis);
-        try {
-            future.orTimeout(timeoutInMillis, TimeUnit.MILLISECONDS).get();
-        } catch (Throwable ignored) {
-        }
+        mTestLooper.dispatchAll();
+
+        assertThat(future).isNotNull();
+        assertThat(future.isDone()).isTrue();
         assertThat(mFakeVibratorController.didRequestVibrationParams).isTrue();
         assertThat(mFakeVibratorController.requestVibrationType).isEqualTo(
                 ScaleParam.TYPE_RINGTONE);
@@ -315,6 +322,46 @@
         }
     }
 
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_THROTTLE_VIBRATION_PARAMS_REQUESTS)
+    public void testRequestVibrationParams_withOngoingRequestAndSameUsage_returnOngoingFuture() {
+        int timeoutInMillis = 10;
+        mVibratorControlService.registerVibratorController(mFakeVibratorController);
+        CompletableFuture<Void> future =
+                mVibratorControlService.triggerVibrationParamsRequest(UID, USAGE_RINGTONE,
+                        timeoutInMillis);
+        CompletableFuture<Void> future2 =
+                mVibratorControlService.triggerVibrationParamsRequest(UID, USAGE_RINGTONE,
+                        timeoutInMillis);
+        mTestLooper.dispatchAll();
+
+        assertThat(future).isNotNull();
+        assertThat(future).isEqualTo(future2);
+        assertThat(future.isDone()).isTrue();
+        assertThat(mFakeVibratorController.requestVibrationParamsCounter).isEqualTo(1);
+    }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_THROTTLE_VIBRATION_PARAMS_REQUESTS)
+    public void testRequestVibrationParams_withOngoingRequestAndSameUsage_returnNewFuture() {
+        int timeoutInMillis = 10;
+        mVibratorControlService.registerVibratorController(mFakeVibratorController);
+        CompletableFuture<Void> future =
+                mVibratorControlService.triggerVibrationParamsRequest(UID, USAGE_RINGTONE,
+                        timeoutInMillis);
+        CompletableFuture<Void> future2 =
+                mVibratorControlService.triggerVibrationParamsRequest(UID, USAGE_RINGTONE,
+                        timeoutInMillis);
+        mTestLooper.dispatchAll();
+
+        assertThat(future).isNotNull();
+        assertThat(future2).isNotNull();
+        assertThat(future).isNotEqualTo(future2);
+        assertThat(future.isDone()).isTrue();
+        assertThat(future2.isDone()).isTrue();
+        assertThat(mFakeVibratorController.requestVibrationParamsCounter).isEqualTo(2);
+    }
+
     private static int buildVibrationTypesMask(int... types) {
         int typesMask = 0;
         for (int type : types) {
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java
index 1f4a469..bea6917 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java
@@ -1603,7 +1603,7 @@
     @Test
     @RequiresFlagsEnabled(android.os.vibrator.Flags.FLAG_VENDOR_VIBRATION_EFFECTS)
     public void vibrate_vendorEffectsWithPermission_successful() throws Exception {
-        // Deny permission to vibrate with vendor effects
+        // Grant permission to vibrate with vendor effects
         grantPermission(android.Manifest.permission.VIBRATE_VENDOR_EFFECTS);
         mockVibrators(1);
         FakeVibratorControllerProvider fakeVibrator = mVibratorProviders.get(1);
@@ -1767,6 +1767,9 @@
     })
     public void vibrate_withIntensitySettingsAndAdaptiveHaptics_appliesSettingsToVendorEffects()
             throws Exception {
+        // Grant permission to vibrate with vendor effects
+        grantPermission(android.Manifest.permission.VIBRATE_VENDOR_EFFECTS);
+
         setUserSetting(Settings.System.NOTIFICATION_VIBRATION_INTENSITY,
                 Vibrator.VIBRATION_INTENSITY_LOW);
 
@@ -1789,7 +1792,7 @@
 
         assertThat(fakeVibrator.getAllVendorEffects()).hasSize(1);
         VibrationEffect.VendorEffect scaled = fakeVibrator.getAllVendorEffects().get(0);
-        assertThat(scaled.getEffectStrength()).isEqualTo(VibrationEffect.EFFECT_STRENGTH_STRONG);
+        assertThat(scaled.getEffectStrength()).isEqualTo(VibrationEffect.EFFECT_STRENGTH_LIGHT);
         assertThat(scaled.getLinearScale()).isEqualTo(0.4f);
     }
 
diff --git a/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorController.java b/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorController.java
index 0cd88ef..c0e1407 100644
--- a/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorController.java
+++ b/services/tests/vibrator/utils/com/android/server/vibrator/FakeVibratorController.java
@@ -39,6 +39,7 @@
     public boolean didRequestVibrationParams = false;
     public int requestVibrationType = VibrationAttributes.USAGE_UNKNOWN;
     public long requestTimeoutInMillis = 0;
+    public int requestVibrationParamsCounter = 0;
 
     public FakeVibratorController(Looper looper) {
         mHandler = new Handler(looper);
@@ -58,6 +59,7 @@
         didRequestVibrationParams = true;
         requestVibrationType = vibrationType;
         requestTimeoutInMillis = timeoutInMillis;
+        requestVibrationParamsCounter++;
         mHandler.post(() -> {
             if (mVibratorControlService != null) {
                 mVibratorControlService.onRequestVibrationParamsComplete(token, mRequestResult);
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java
index 11f7560..220248c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java
@@ -126,6 +126,10 @@
                 .isCameraActive(any(ActivityRecord.class), anyBoolean());
     }
 
+    void setDisplayNaturalOrientation(@Configuration.Orientation int naturalOrientation) {
+        doReturn(naturalOrientation).when(mDisplayContent).getNaturalOrientation();
+    }
+
     @NonNull
     ActivityRecord top() {
         return mActivityStack.top();
@@ -189,6 +193,11 @@
         mDisplayContent.setIgnoreOrientationRequest(enabled);
     }
 
+    void setTopTaskInMultiWindowMode(boolean inMultiWindowMode) {
+        doReturn(inMultiWindowMode).when(mTaskStack.top())
+                .inMultiWindowMode();
+    }
+
     void setTopActivityAsEmbedded(boolean embedded) {
         doReturn(embedded).when(mActivityStack.top()).isEmbedded();
     }
@@ -225,20 +234,21 @@
     void createNewDisplay() {
         mDisplayContent = new TestDisplayContent.Builder(mAtm, mDisplayWidth, mDisplayHeight)
                 .build();
+        spyOn(mDisplayContent);
         spyOnAppCompatCameraPolicy();
     }
 
     void createNewTask() {
         final Task newTask = new WindowTestsBase.TaskBuilder(mSupervisor)
                 .setDisplay(mDisplayContent).build();
-        mTaskStack.push(newTask);
+        pushTask(newTask);
     }
 
     void createNewTaskWithBaseActivity() {
         final Task newTask = new WindowTestsBase.TaskBuilder(mSupervisor)
                 .setCreateActivity(true)
                 .setDisplay(mDisplayContent).build();
-        mTaskStack.push(newTask);
+        pushTask(newTask);
         pushActivity(newTask.getTopNonFinishingActivity());
     }
 
@@ -430,9 +440,15 @@
         spyOn(activity.mAppCompatController.getAppCompatAspectRatioOverrides());
         spyOn(activity.mAppCompatController.getAppCompatAspectRatioPolicy());
         spyOn(activity.mAppCompatController.getAppCompatFocusOverrides());
+        spyOn(activity.mAppCompatController.getAppCompatResizeOverrides());
         spyOn(activity.mLetterboxUiController);
     }
 
+    private void pushTask(@NonNull Task task) {
+        spyOn(task);
+        mTaskStack.push(task);
+    }
+
     private void spyOnAppCompatCameraPolicy() {
         spyOn(mDisplayContent.mAppCompatCameraPolicy);
         if (mDisplayContent.mAppCompatCameraPolicy.hasDisplayRotationCompatPolicy()) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatAspectRatioOverridesTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatAspectRatioOverridesTest.java
index ddd6d56..a6fd112 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatAspectRatioOverridesTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatAspectRatioOverridesTest.java
@@ -26,6 +26,7 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
 
 import android.compat.testing.PlatformCompatChangeRule;
 import android.platform.test.annotations.Presubmit;
@@ -246,7 +247,6 @@
         });
     }
 
-
     @Test
     @EnableCompatChanges({OVERRIDE_MIN_ASPECT_RATIO})
     public void testshouldOverrideMinAspectRatio_propertyFalse_overrideEnabled_returnsFalse() {
@@ -269,6 +269,24 @@
         });
     }
 
+    @Test
+    public void testGetFixedOrientationLetterboxAspectRatio_splitScreenAspectEnabled() {
+        runTestScenario((robot)-> {
+            robot.applyOnConf((c) -> {
+                c.enableCameraCompatTreatment(/* enabled */ true);
+                c.enableCameraCompatTreatmentAtBuildTime(/* enabled */ true);
+                c.enableCameraCompatSplitScreenAspectRatio(/* enabled */ true);
+                c.enableDisplayAspectRatioEnabledForFixedOrientationLetterbox(/* enabled */ false);
+                c.setFixedOrientationLetterboxAspectRatio(/* aspectRatio */ 1.5f);
+            });
+            robot.activity().createActivityWithComponentInNewTaskAndDisplay();
+            robot.checkFixedOrientationLetterboxAspectRatioForTopParent(/* expected */ 1.5f);
+
+            robot.activity().enableTreatmentForTopActivity(/* enabled */ true);
+            robot.checkAspectRatioForTopParentIsSplitScreenRatio(/* expected */ true);
+        });
+    }
+
     /**
      * Runs a test scenario providing a Robot.
      */
@@ -308,6 +326,28 @@
         }
 
         @NonNull
+        void checkFixedOrientationLetterboxAspectRatioForTopParent(float expected) {
+            assertEquals(expected,
+                    getTopActivityAppCompatAspectRatioOverrides()
+                            .getFixedOrientationLetterboxAspectRatio(
+                                    activity().top().getParent().getConfiguration()),
+                                        FLOAT_TOLLERANCE);
+        }
+
+        void checkAspectRatioForTopParentIsSplitScreenRatio(boolean expected) {
+            final AppCompatAspectRatioOverrides aspectRatioOverrides =
+                    getTopActivityAppCompatAspectRatioOverrides();
+            if (expected) {
+                assertEquals(aspectRatioOverrides.getSplitScreenAspectRatio(),
+                        aspectRatioOverrides.getFixedOrientationLetterboxAspectRatio(
+                                activity().top().getParent().getConfiguration()), FLOAT_TOLLERANCE);
+            } else {
+                assertNotEquals(aspectRatioOverrides.getSplitScreenAspectRatio(),
+                        aspectRatioOverrides.getFixedOrientationLetterboxAspectRatio(
+                                activity().top().getParent().getConfiguration()), FLOAT_TOLLERANCE);
+            }
+        }
+
         private AppCompatAspectRatioOverrides getTopActivityAppCompatAspectRatioOverrides() {
             return activity().top().mAppCompatController.getAppCompatAspectRatioOverrides();
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatConfigurationRobot.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatConfigurationRobot.java
index 00a8771..6592f26 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatConfigurationRobot.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatConfigurationRobot.java
@@ -70,4 +70,14 @@
     void enableCompatFakeFocus(boolean enabled) {
         doReturn(enabled).when(mAppCompatConfiguration).isCompatFakeFocusEnabled();
     }
+
+    void enableDisplayAspectRatioEnabledForFixedOrientationLetterbox(boolean enabled) {
+        doReturn(enabled).when(mAppCompatConfiguration)
+                .getIsDisplayAspectRatioEnabledForFixedOrientationLetterbox();
+    }
+
+    void setFixedOrientationLetterboxAspectRatio(float aspectRatio) {
+        doReturn(aspectRatio).when(mAppCompatConfiguration)
+                .getFixedOrientationLetterboxAspectRatio();
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java
index 634453f..6c0d8c4 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java
@@ -16,6 +16,10 @@
 package com.android.server.wm;
 
 import static android.content.pm.ActivityInfo.OVERRIDE_ENABLE_COMPAT_IGNORE_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED;
+import static android.content.pm.ActivityInfo.OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION;
+import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
+import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
+import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE;
 import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_IGNORING_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
@@ -158,6 +162,72 @@
         });
     }
 
+    @Test
+    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
+    public void testShouldUseDisplayLandscapeNaturalOrientation_override_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.applyOnActivity((a) -> {
+                a.setDisplayNaturalOrientation(ORIENTATION_LANDSCAPE);
+                a.setIgnoreOrientationRequest(true);
+                a.createActivityWithComponent();
+            });
+            robot.checkShouldUseDisplayLandscapeNaturalOrientation(/* expected */ true);
+        });
+    }
+
+    @Test
+    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
+    public void testShouldUseDisplayLandscapeNaturalOrientation_falseProperty_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.prop().disable(PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE);
+            robot.applyOnActivity((a) -> {
+                a.setDisplayNaturalOrientation(ORIENTATION_LANDSCAPE);
+                a.setIgnoreOrientationRequest(true);
+                a.createActivityWithComponent();
+            });
+            robot.checkShouldUseDisplayLandscapeNaturalOrientation(/* expected */ false);
+        });
+    }
+
+    @Test
+    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
+    public void testShouldUseDisplayLandscapeNaturalOrientation_portrait_isFalse() {
+        runTestScenario((robot) -> {
+            robot.applyOnActivity((a) -> {
+                a.setDisplayNaturalOrientation(ORIENTATION_PORTRAIT);
+                a.setIgnoreOrientationRequest(true);
+                a.createActivityWithComponent();
+            });
+            robot.checkShouldUseDisplayLandscapeNaturalOrientation(/* expected */ false);
+        });
+    }
+    @Test
+    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
+    public void testShouldUseDisplayLandscapeNaturalOrientation_noIgnoreRequest_isFalse() {
+        runTestScenario((robot) -> {
+            robot.applyOnActivity((a) -> {
+                a.setDisplayNaturalOrientation(ORIENTATION_LANDSCAPE);
+                a.setIgnoreOrientationRequest(false);
+                a.createActivityWithComponent();
+            });
+            robot.checkShouldUseDisplayLandscapeNaturalOrientation(/* expected */ false);
+        });
+    }
+
+    @Test
+    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
+    public void testShouldUseDisplayLandscapeNaturalOrientation_inMultiWindowMode_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.applyOnActivity((a) -> {
+                a.setDisplayNaturalOrientation(ORIENTATION_LANDSCAPE);
+                a.setIgnoreOrientationRequest(true);
+                a.createActivityWithComponent();
+                a.setTopTaskInMultiWindowMode(/* inMultiWindowMode */ true);
+            });
+            robot.checkShouldUseDisplayLandscapeNaturalOrientation(/* expected */ false);
+        });
+    }
+
     /**
      * Runs a test scenario providing a Robot.
      */
@@ -215,6 +285,11 @@
             }
         }
 
+        void checkShouldUseDisplayLandscapeNaturalOrientation(boolean expected) {
+            assertEquals(expected,
+                    getTopOrientationOverrides().shouldUseDisplayLandscapeNaturalOrientation());
+        }
+
         private AppCompatOrientationOverrides getTopOrientationOverrides() {
             return activity().top().mAppCompatController.getAppCompatOverrides()
                     .getAppCompatOrientationOverrides();
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatResizeOverridesTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatResizeOverridesTest.java
new file mode 100644
index 0000000..8fc1a77
--- /dev/null
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatResizeOverridesTest.java
@@ -0,0 +1,199 @@
+/*
+ * Copyright (C) 2024 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.server.wm;
+
+import static android.content.pm.ActivityInfo.FORCE_NON_RESIZE_APP;
+import static android.content.pm.ActivityInfo.FORCE_RESIZE_APP;
+import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
+import android.compat.testing.PlatformCompatChangeRule;
+import android.platform.test.annotations.Presubmit;
+
+import androidx.annotation.NonNull;
+
+import libcore.junit.util.compat.CoreCompatChangeRule;
+
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestRule;
+import org.junit.runner.RunWith;
+
+import java.util.function.Consumer;
+
+/**
+ * Test class for {@link AppCompatResizeOverrides}.
+ * <p>
+ * Build/Install/Run:
+ * atest WmTests:AppCompatResizeOverridesTest
+ */
+@Presubmit
+@RunWith(WindowTestRunner.class)
+public class AppCompatResizeOverridesTest extends WindowTestsBase {
+
+    @Rule
+    public TestRule compatChangeRule = new PlatformCompatChangeRule();
+
+    @Test
+    @CoreCompatChangeRule.EnableCompatChanges({FORCE_RESIZE_APP})
+    public void testShouldOverrideForceResizeApp_overrideEnabled_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceResizeApp(/* expected */ true);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.EnableCompatChanges({FORCE_RESIZE_APP})
+    public void testShouldOverrideForceResizeApp_propertyTrue_overrideEnabled_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.prop().enable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceResizeApp(/* expected */ true);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.DisableCompatChanges({FORCE_RESIZE_APP})
+    public void testShouldOverrideForceResizeApp_propertyTrue_overrideDisabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.prop().enable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceResizeApp(/* expected */ false);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.DisableCompatChanges({FORCE_RESIZE_APP})
+    public void testShouldOverrideForceResizeApp_overrideDisabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceResizeApp(/* expected */ false);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.EnableCompatChanges({FORCE_RESIZE_APP})
+    public void testShouldOverrideForceResizeApp_propertyFalse_overrideEnabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.prop().disable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceResizeApp(/* expected */ false);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.DisableCompatChanges({FORCE_RESIZE_APP})
+    public void testShouldOverrideForceResizeApp_propertyFalse_noOverride_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.prop().disable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceResizeApp(/* expected */ false);
+        });
+    }
+
+
+    @Test
+    @CoreCompatChangeRule.EnableCompatChanges({FORCE_NON_RESIZE_APP})
+    public void testShouldOverrideForceNonResizeApp_overrideEnabled_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceNonResizeApp(/* expected */ true);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.EnableCompatChanges({FORCE_NON_RESIZE_APP})
+    public void testShouldOverrideForceNonResizeApp_propertyTrue_overrideEnabled_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.prop().enable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceNonResizeApp(/* expected */ true);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.DisableCompatChanges({FORCE_NON_RESIZE_APP})
+    public void testShouldOverrideForceNonResizeApp_propertyTrue_overrideDisabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.prop().enable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceNonResizeApp(/* expected */ false);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.DisableCompatChanges({FORCE_NON_RESIZE_APP})
+    public void testShouldOverrideForceNonResizeApp_overrideDisabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceNonResizeApp(/* expected */ false);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.EnableCompatChanges({FORCE_NON_RESIZE_APP})
+    public void testShouldOverrideForceNonResizeApp_propertyFalse_overrideEnabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.prop().disable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceNonResizeApp(/* expected */ false);
+        });
+    }
+
+    @Test
+    @CoreCompatChangeRule.DisableCompatChanges({FORCE_NON_RESIZE_APP})
+    public void testShouldOverrideForceNonResizeApp_propertyFalse_noOverride_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.prop().disable(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES);
+            robot.activity().createActivityWithComponent();
+            robot.checkShouldOverrideForceNonResizeApp(/* expected */ false);
+        });
+    }
+
+    /**
+     * Runs a test scenario providing a Robot.
+     */
+    void runTestScenario(@NonNull Consumer<ResizeOverridesRobotTest> consumer) {
+        spyOn(mWm.mAppCompatConfiguration);
+        final ResizeOverridesRobotTest robot = new ResizeOverridesRobotTest(mWm, mAtm, mSupervisor);
+        consumer.accept(robot);
+    }
+
+    private static class ResizeOverridesRobotTest extends AppCompatRobotBase {
+
+        ResizeOverridesRobotTest(@NonNull WindowManagerService wm,
+                @NonNull ActivityTaskManagerService atm,
+                @NonNull ActivityTaskSupervisor supervisor) {
+            super(wm, atm, supervisor);
+        }
+
+
+        void checkShouldOverrideForceResizeApp(boolean expected) {
+            Assert.assertEquals(expected, activity().top().mAppCompatController
+                    .getAppCompatResizeOverrides().shouldOverrideForceResizeApp());
+        }
+
+        void checkShouldOverrideForceNonResizeApp(boolean expected) {
+            Assert.assertEquals(expected, activity().top().mAppCompatController
+                    .getAppCompatResizeOverrides().shouldOverrideForceNonResizeApp());
+        }
+    }
+
+}
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatRobotBase.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatRobotBase.java
index 92f246b..6939f97 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatRobotBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatRobotBase.java
@@ -28,6 +28,8 @@
     private static final int DEFAULT_DISPLAY_WIDTH = 1000;
     private static final int DEFAULT_DISPLAY_HEIGHT = 2000;
 
+    static final float FLOAT_TOLLERANCE = 0.01f;
+
     @NonNull
     private final AppCompatActivityRobot mActivityRobot;
     @NonNull
diff --git a/services/tests/wmtests/src/com/android/server/wm/DragDropControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/DragDropControllerTests.java
index 8cdb574..4a9d5c7 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DragDropControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DragDropControllerTests.java
@@ -152,11 +152,10 @@
         // Use a new TestIWindow so we don't collect events for other windows
         final WindowState window = createWindow(
                 null, TYPE_BASE_APPLICATION, activity, name, ownerId, false, new TestIWindow());
-        window.mInputChannel = new InputChannel();
-        window.mInputChannelToken = window.mInputChannel.getToken();
+        InputChannel channel = new InputChannel();
+        window.openInputChannel(channel);
         window.mHasSurface = true;
         mWm.mWindowMap.put(window.mClient.asBinder(), window);
-        mWm.mInputToWindowMap.put(window.mInputChannelToken, window);
         return window;
     }
 
@@ -178,8 +177,8 @@
                 TEST_PID, TEST_UID);
         mWindow = createDropTargetWindow("Drag test window", 0);
         doReturn(mWindow).when(mDisplayContent).getTouchableWinAtPointLocked(0, 0);
-        when(mWm.mInputManager.startDragAndDrop(any(InputChannel.class),
-                any(InputChannel.class))).thenReturn(true);
+        when(mWm.mInputManager.startDragAndDrop(any(IBinder.class),
+                any(IBinder.class))).thenReturn(true);
 
         mWm.mWindowMap.put(mWindow.mClient.asBinder(), mWindow);
     }
@@ -707,8 +706,7 @@
                     .setFormat(PixelFormat.TRANSLUCENT)
                     .build();
 
-            assertTrue(mWm.mInputManager.startDragAndDrop(new InputChannel(),
-                    new InputChannel()));
+            assertTrue(mWm.mInputManager.startDragAndDrop(new Binder(), new Binder()));
             mToken = mTarget.performDrag(TEST_PID, 0, mWindow.mClient,
                     flag, surface, 0, 0, 0, 0, 0, 0, 0, data);
             assertNotNull(mToken);
diff --git a/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java
index a0a2904..61a6f31 100644
--- a/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java
@@ -16,17 +16,9 @@
 
 package com.android.server.wm;
 
-import static android.content.pm.ActivityInfo.FORCE_NON_RESIZE_APP;
-import static android.content.pm.ActivityInfo.FORCE_RESIZE_APP;
-import static android.content.pm.ActivityInfo.OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION;
-import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
-import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
 import static android.view.InsetsSource.FLAG_INSETS_ROUNDED_CORNER;
-import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE;
-import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.eq;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
@@ -36,15 +28,12 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.annotation.Nullable;
 import android.compat.testing.PlatformCompatChangeRule;
 import android.content.ComponentName;
-import android.content.pm.PackageManager;
-import android.content.pm.PackageManager.Property;
 import android.content.res.Resources;
 import android.graphics.Rect;
 import android.platform.test.annotations.DisableFlags;
@@ -62,9 +51,6 @@
 import com.android.internal.R;
 import com.android.window.flags.Flags;
 
-import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
-import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
-
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -309,203 +295,6 @@
         return mainWindow;
     }
 
-    // shouldUseDisplayLandscapeNaturalOrientation
-
-    @Test
-    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
-    public void testShouldUseDisplayLandscapeNaturalOrientation_override_returnsTrue() {
-        prepareActivityThatShouldUseDisplayLandscapeNaturalOrientation();
-        assertTrue(mController.shouldUseDisplayLandscapeNaturalOrientation());
-    }
-
-    @Test
-    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
-    public void testShouldUseDisplayLandscapeNaturalOrientation_overrideAndFalseProperty_returnsFalse()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_DISPLAY_ORIENTATION_OVERRIDE, /* value */ false);
-
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        prepareActivityThatShouldUseDisplayLandscapeNaturalOrientation();
-        assertFalse(mController.shouldUseDisplayLandscapeNaturalOrientation());
-    }
-
-    @Test
-    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
-    public void testShouldUseDisplayLandscapeNaturalOrientation_portraitNaturalOrientation_returnsFalse() {
-        prepareActivityThatShouldUseDisplayLandscapeNaturalOrientation();
-        doReturn(ORIENTATION_PORTRAIT).when(mDisplayContent).getNaturalOrientation();
-
-        assertFalse(mController.shouldUseDisplayLandscapeNaturalOrientation());
-    }
-
-    @Test
-    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
-    public void testShouldUseDisplayLandscapeNaturalOrientation_disabledIgnoreOrientationRequest_returnsFalse() {
-        prepareActivityThatShouldUseDisplayLandscapeNaturalOrientation();
-        mDisplayContent.setIgnoreOrientationRequest(false);
-
-        assertFalse(mController.shouldUseDisplayLandscapeNaturalOrientation());
-    }
-
-    @Test
-    @EnableCompatChanges({OVERRIDE_USE_DISPLAY_LANDSCAPE_NATURAL_ORIENTATION})
-    public void testShouldUseDisplayLandscapeNaturalOrientation_inMultiWindowMode_returnsFalse() {
-        prepareActivityThatShouldUseDisplayLandscapeNaturalOrientation();
-
-        spyOn(mTask);
-        doReturn(true).when(mTask).inMultiWindowMode();
-
-        assertFalse(mController.shouldUseDisplayLandscapeNaturalOrientation());
-    }
-
-    @Test
-    @EnableCompatChanges({FORCE_RESIZE_APP})
-    public void testshouldOverrideForceResizeApp_overrideEnabled_returnsTrue() {
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertTrue(mController.shouldOverrideForceResizeApp());
-    }
-
-    @Test
-    @EnableCompatChanges({FORCE_RESIZE_APP})
-    public void testshouldOverrideForceResizeApp_propertyTrue_overrideEnabled_returnsTrue()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ true);
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertTrue(mController.shouldOverrideForceResizeApp());
-    }
-
-    @Test
-    @DisableCompatChanges({FORCE_RESIZE_APP})
-    public void testshouldOverrideForceResizeApp_propertyTrue_overrideDisabled_returnsFalse()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ true);
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceResizeApp());
-    }
-
-    @Test
-    @DisableCompatChanges({FORCE_RESIZE_APP})
-    public void testshouldOverrideForceResizeApp_overrideDisabled_returnsFalse() {
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceResizeApp());
-    }
-
-    @Test
-    @EnableCompatChanges({FORCE_RESIZE_APP})
-    public void testshouldOverrideForceResizeApp_propertyFalse_overrideEnabled_returnsFalse()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ false);
-
-        mActivity = setUpActivityWithComponent();
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceResizeApp());
-    }
-
-    @Test
-    @DisableCompatChanges({FORCE_RESIZE_APP})
-    public void testshouldOverrideForceResizeApp_propertyFalse_noOverride_returnsFalse()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ false);
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceResizeApp());
-    }
-
-    @Test
-    @EnableCompatChanges({FORCE_NON_RESIZE_APP})
-    public void testshouldOverrideForceNonResizeApp_overrideEnabled_returnsTrue() {
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertTrue(mController.shouldOverrideForceNonResizeApp());
-    }
-
-    @Test
-    @EnableCompatChanges({FORCE_NON_RESIZE_APP})
-    public void testshouldOverrideForceNonResizeApp_propertyTrue_overrideEnabled_returnsTrue()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ true);
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertTrue(mController.shouldOverrideForceNonResizeApp());
-    }
-
-    @Test
-    @DisableCompatChanges({FORCE_NON_RESIZE_APP})
-    public void testshouldOverrideForceNonResizeApp_propertyTrue_overrideDisabled_returnsFalse()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ true);
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceNonResizeApp());
-    }
-
-    @Test
-    @DisableCompatChanges({FORCE_NON_RESIZE_APP})
-    public void testshouldOverrideForceNonResizeApp_overrideDisabled_returnsFalse() {
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceNonResizeApp());
-    }
-
-    @Test
-    @EnableCompatChanges({FORCE_NON_RESIZE_APP})
-    public void testshouldOverrideForceNonResizeApp_propertyFalse_overrideEnabled_returnsFalse()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ false);
-
-        mActivity = setUpActivityWithComponent();
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceNonResizeApp());
-    }
-
-    @Test
-    @DisableCompatChanges({FORCE_NON_RESIZE_APP})
-    public void testshouldOverrideForceNonResizeApp_propertyFalse_noOverride_returnsFalse()
-            throws Exception {
-        mockThatProperty(PROPERTY_COMPAT_ALLOW_RESIZEABLE_ACTIVITY_OVERRIDES, /* value */ false);
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertFalse(mController.shouldOverrideForceNonResizeApp());
-    }
-
-    @Test
-    public void testgetFixedOrientationLetterboxAspectRatio_splitScreenAspectEnabled() {
-        doReturn(true).when(mActivity.mWmService.mAppCompatConfiguration)
-                .isCameraCompatTreatmentEnabled();
-        doReturn(true).when(mActivity.mWmService.mAppCompatConfiguration)
-                .isCameraCompatTreatmentEnabledAtBuildTime();
-        doReturn(true).when(mActivity.mWmService.mAppCompatConfiguration)
-                .isCameraCompatSplitScreenAspectRatioEnabled();
-        doReturn(false).when(mActivity.mWmService.mAppCompatConfiguration)
-                .getIsDisplayAspectRatioEnabledForFixedOrientationLetterbox();
-        doReturn(1.5f).when(mActivity.mWmService.mAppCompatConfiguration)
-                .getFixedOrientationLetterboxAspectRatio();
-
-        // Recreate DisplayContent with DisplayRotationCompatPolicy
-        mActivity = setUpActivityWithComponent();
-        mController = new LetterboxUiController(mWm, mActivity);
-
-        assertEquals(1.5f, mController.getFixedOrientationLetterboxAspectRatio(
-                mActivity.getParent().getConfiguration()), /* delta */ 0.01);
-
-        spyOn(mDisplayContent.mAppCompatCameraPolicy);
-        doReturn(true).when(mDisplayContent.mAppCompatCameraPolicy)
-                .isTreatmentEnabledForActivity(eq(mActivity));
-
-        final AppCompatAspectRatioOverrides aspectRatioOverrides =
-                mActivity.mAppCompatController.getAppCompatAspectRatioOverrides();
-        assertEquals(aspectRatioOverrides.getSplitScreenAspectRatio(),
-                aspectRatioOverrides.getFixedOrientationLetterboxAspectRatio(
-                        mActivity.getParent().getConfiguration()), /* delta */  0.01);
-    }
-
     @Test
     public void testIsVerticalThinLetterboxed() {
         // Vertical thin letterbox disabled
@@ -604,20 +393,6 @@
         verify(mAppCompatConfiguration).getIsEducationEnabled();
     }
 
-    private void mockThatProperty(String propertyName, boolean value) throws Exception {
-        Property property = new Property(propertyName, /* value */ value, /* packageName */ "",
-                /* className */ "");
-        PackageManager pm = mWm.mContext.getPackageManager();
-        spyOn(pm);
-        doReturn(property).when(pm).getProperty(eq(propertyName), anyString());
-    }
-
-    private void prepareActivityThatShouldUseDisplayLandscapeNaturalOrientation() {
-        spyOn(mDisplayContent);
-        doReturn(ORIENTATION_LANDSCAPE).when(mDisplayContent).getNaturalOrientation();
-        mDisplayContent.setIgnoreOrientationRequest(true);
-    }
-
     private ActivityRecord setUpActivityWithComponent() {
         mDisplayContent = new TestDisplayContent
                 .Builder(mAtm, /* dw */ 1000, /* dh */ 2000).build();
diff --git a/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java b/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
index 72747c9..ed93a8c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SizeCompatTests.java
@@ -788,7 +788,7 @@
         // Change the fixed orientation.
         mActivity.setRequestedOrientation(SCREEN_ORIENTATION_LANDSCAPE);
         assertTrue(mActivity.isRelaunching());
-        assertTrue(mActivity.mLetterboxUiController
+        assertTrue(mActivity.mAppCompatController.getAppCompatOrientationOverrides()
                 .getIsRelaunchingAfterRequestedOrientationChanged());
 
         assertFitted();
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskTests.java b/services/tests/wmtests/src/com/android/server/wm/TaskTests.java
index a1ac02a..a232ff0 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskTests.java
@@ -46,6 +46,7 @@
 import static com.android.server.wm.Task.FLAG_FORCE_HIDDEN_FOR_TASK_ORG;
 import static com.android.server.wm.TaskFragment.EMBEDDED_DIM_AREA_PARENT_TASK;
 import static com.android.server.wm.TaskFragment.TASK_FRAGMENT_VISIBILITY_VISIBLE_BEHIND_TRANSLUCENT;
+import static com.android.server.wm.WindowContainer.POSITION_TOP;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -78,6 +79,7 @@
 import android.content.pm.ActivityInfo;
 import android.content.pm.ApplicationInfo;
 import android.content.res.Configuration;
+import android.graphics.Color;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.IBinder;
@@ -2031,6 +2033,47 @@
                 task.getTaskInfo().appCompatTaskInfo.cameraCompatTaskInfo.freeformCameraCompatMode);
     }
 
+    @Test
+    public void testUpdateTaskDescriptionOnReparent() {
+        final Task rootTask1 = createTask(mDisplayContent);
+        final Task rootTask2 = createTask(mDisplayContent);
+        final Task childTask = createTaskInRootTask(rootTask1, 0 /* userId */);
+        final ActivityRecord activity = createActivityRecord(mDisplayContent, childTask);
+        final String testLabel = "test_task_description_label";
+        final ActivityManager.TaskDescription td = new ActivityManager.TaskDescription(testLabel);
+        activity.setTaskDescription(td);
+
+        // Ensure the td is set for the original root task
+        assertEquals(testLabel, rootTask1.getTaskDescription().getLabel());
+        assertNull(rootTask2.getTaskDescription().getLabel());
+
+        childTask.reparent(rootTask2, POSITION_TOP, false /* moveParents */, "reparent");
+
+        // Ensure the td is set for the new root task
+        assertEquals(testLabel, rootTask2.getTaskDescription().getLabel());
+    }
+
+    @Test
+    public void testUpdateTaskDescriptionOnReorder() {
+        final Task task = createTask(mDisplayContent);
+        final ActivityRecord activity1 = createActivityRecord(mDisplayContent, task);
+        final ActivityRecord activity2 = createActivityRecord(mDisplayContent, task);
+        final ActivityManager.TaskDescription td1 = new ActivityManager.TaskDescription();
+        td1.setBackgroundColor(Color.RED);
+        activity1.setTaskDescription(td1);
+        final ActivityManager.TaskDescription td2 = new ActivityManager.TaskDescription();
+        td2.setBackgroundColor(Color.BLUE);
+        activity2.setTaskDescription(td2);
+
+        // Ensure the td is set for the original root task
+        assertEquals(Color.BLUE, task.getTaskDescription().getBackgroundColor());
+
+        task.positionChildAt(POSITION_TOP, activity1, false /* includeParents */);
+
+        // Ensure the td is set for the original root task
+        assertEquals(Color.RED, task.getTaskDescription().getBackgroundColor());
+    }
+
     private Task getTestTask() {
         return new TaskBuilder(mSupervisor).setCreateActivity(true).build();
     }
diff --git a/telephony/java/android/telephony/CarrierConfigManager.java b/telephony/java/android/telephony/CarrierConfigManager.java
index 61698db..0468f48 100644
--- a/telephony/java/android/telephony/CarrierConfigManager.java
+++ b/telephony/java/android/telephony/CarrierConfigManager.java
@@ -9992,6 +9992,51 @@
     @FlaggedApi(Flags.FLAG_CARRIER_ROAMING_NB_IOT_NTN)
     public static final String KEY_SATELLITE_ESOS_SUPPORTED_BOOL = "satellite_esos_supported_bool";
 
+    /** @hide */
+    @IntDef({
+            CARRIER_ROAMING_NTN_CONNECT_AUTOMATIC,
+            CARRIER_ROAMING_NTN_CONNECT_MANUAL,
+    })
+    public @interface CARRIER_ROAMING_NTN_CONNECT_TYPE {}
+
+    /**
+     * Device can connect to carrier roaming non-terrestrial network automatically.
+     * @hide
+     */
+    public static final int CARRIER_ROAMING_NTN_CONNECT_AUTOMATIC = 0;
+    /**
+     * Device can connect to carrier roaming non-terrestrial network only if user manually triggers
+     * satellite connection.
+     * @hide
+     */
+    public static final int CARRIER_ROAMING_NTN_CONNECT_MANUAL = 1;
+    /**
+     * Indicates carrier roaming non-terrestrial network connect type that the device can use to
+     * perform satellite communication.
+     * If this key is set to CARRIER_ROAMING_NTN_CONNECT_MANUAL then connect button will be
+     * displayed to user when the device is eligible to use carrier roaming
+     * non-terrestrial network.
+     * @hide
+     */
+    public static final String KEY_CARRIER_ROAMING_NTN_CONNECT_TYPE_INT =
+            "carrier_roaming_ntn_connect_type_int";
+
+    /**
+     * The carrier roaming non-terrestrial network hysteresis time in seconds.
+     *
+     * If the device supports P2P satellite messaging which is defined by
+     * {@link CarrierConfigManager#KEY_CARRIER_SUPPORTED_SATELLITE_SERVICES_PER_PROVIDER_BUNDLE}
+     * and the device is in {@link ServiceState#STATE_OUT_OF_SERVICE}, not connected to Wi-Fi,
+     * then hysteresis timer defined by this key will start.
+     * After the timer is expired, device is marked as eligible for satellite communication.
+     *
+     * The default value is 180 seconds.
+     *
+     * @hide
+     */
+    public static final String KEY_CARRIER_SUPPORTED_SATELLITE_NOTIFICATION_HYSTERESIS_SEC_INT =
+            "carrier_supported_satellite_notification_hysteresis_sec_int";
+
     /**
      * Indicating whether DUN APN should be disabled when the device is roaming. In that case,
      * the default APN (i.e. internet) will be used for tethering.
@@ -11150,6 +11195,8 @@
         sDefaults.putInt(KEY_EMERGENCY_CALL_TO_SATELLITE_T911_HANDOVER_TIMEOUT_MILLIS_INT,
                 (int) TimeUnit.SECONDS.toMillis(30));
         sDefaults.putBoolean(KEY_SATELLITE_ESOS_SUPPORTED_BOOL, false);
+        sDefaults.putInt(KEY_CARRIER_ROAMING_NTN_CONNECT_TYPE_INT, 0);
+        sDefaults.putInt(KEY_CARRIER_SUPPORTED_SATELLITE_NOTIFICATION_HYSTERESIS_SEC_INT, 180);
         sDefaults.putString(KEY_DEFAULT_PREFERRED_APN_NAME_STRING, "");
         sDefaults.putBoolean(KEY_SUPPORTS_CALL_COMPOSER_BOOL, false);
         sDefaults.putBoolean(KEY_SUPPORTS_BUSINESS_CALL_COMPOSER_BOOL, false);
diff --git a/wifi/wifi.aconfig b/wifi/wifi.aconfig
index f7162f6..5a214b7 100644
--- a/wifi/wifi.aconfig
+++ b/wifi/wifi.aconfig
@@ -35,3 +35,14 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "hotspot_network_connecting_state_for_details_page"
+    namespace: "wifi"
+    description: "Update getConnectedState in HotspotNetworkEntry so that details page displays correctly."
+    bug: "321096462"
+    is_fixed_read_only: true
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}