Merge "Update boot image and system server profiles [M80C35P56S0PP]" into main
diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 1ae9ada..deb6f13 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -789,6 +789,7 @@
     min_sdk_version: "30",
     apex_available: [
         "//apex_available:platform",
+        "com.android.healthfitness",
         "com.android.permission",
         "com.android.nfcservices",
     ],
diff --git a/MULTIUSER_OWNERS b/MULTIUSER_OWNERS
index b8857ec..1738a35 100644
--- a/MULTIUSER_OWNERS
+++ b/MULTIUSER_OWNERS
@@ -3,7 +3,5 @@
 bookatz@google.com
 nykkumar@google.com
 olilan@google.com
-omakoto@google.com
 tetianameronyk@google.com
 tyk@google.com
-yamasani@google.com
diff --git a/OWNERS b/OWNERS
index afa60be..eb2bfcf 100644
--- a/OWNERS
+++ b/OWNERS
@@ -7,8 +7,6 @@
 hackbod@android.com #{LAST_RESORT_SUGGESTION}
 hackbod@google.com #{LAST_RESORT_SUGGESTION}
 jjaggi@google.com #{LAST_RESORT_SUGGESTION}
-jsharkey@android.com #{LAST_RESORT_SUGGESTION}
-jsharkey@google.com #{LAST_RESORT_SUGGESTION}
 lorenzo@google.com #{LAST_RESORT_SUGGESTION}
 michaelwr@google.com #{LAST_RESORT_SUGGESTION}
 nandana@google.com #{LAST_RESORT_SUGGESTION}
@@ -33,19 +31,19 @@
 per-file TestProtoLibraries.bp = file:platform/platform_testing:/libraries/health/OWNERS
 per-file TestProtoLibraries.bp = file:platform/tools/tradefederation:/OWNERS
 
-per-file INPUT_OWNERS = file:/INPUT_OWNERS
-per-file ZYGOTE_OWNERS = file:/ZYGOTE_OWNERS
-per-file SQLITE_OWNERS = file:/SQLITE_OWNERS
-
 per-file *ravenwood* = file:ravenwood/OWNERS
 per-file *Ravenwood* = file:ravenwood/OWNERS
 
+per-file INPUT_OWNERS = file:/INPUT_OWNERS
+per-file ZYGOTE_OWNERS = file:/ZYGOTE_OWNERS
+per-file SQLITE_OWNERS = file:/SQLITE_OWNERS
 per-file PERFORMANCE_OWNERS = file:/PERFORMANCE_OWNERS
-
 per-file PACKAGE_MANAGER_OWNERS = file:/PACKAGE_MANAGER_OWNERS
-
 per-file WEAR_OWNERS = file:/WEAR_OWNERS
-
+per-file ACTIVITY_MANAGER_OWNERS = file:/ACTIVITY_MANAGER_OWNERS
+per-file BATTERY_STATS_OWNERS = file:/BATTERY_STATS_OWNERS
+per-file OOM_ADJUSTER_OWNERS = file:/OOM_ADJUSTER_OWNERS
+per-file MULTIUSER_OWNERS = file:/MULTIUSER_OWNERS
+per-file BROADCASTS_OWNERS = file:/BROADCASTS_OWNERS
 per-file ADPF_OWNERS = file:/ADPF_OWNERS
-
 per-file GAME_MANAGER_OWNERS = file:/GAME_MANAGER_OWNERS
diff --git a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
index c1894f0..a37779e 100644
--- a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
+++ b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
@@ -3568,7 +3568,7 @@
             Slog.i(TAG, "becomeActiveLocked, reason=" + activeReason
                     + ", changeLightIdle=" + changeLightIdle);
         }
-        if (mState != STATE_ACTIVE || mLightState != STATE_ACTIVE) {
+        if (mState != STATE_ACTIVE || mLightState != LIGHT_STATE_ACTIVE) {
             moveToStateLocked(STATE_ACTIVE, activeReason);
             mInactiveTimeout = newInactiveTimeout;
             resetIdleManagementLocked();
diff --git a/core/api/current.txt b/core/api/current.txt
index 012a2e6..b95295c 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -10864,6 +10864,7 @@
     field public static final String IPSEC_SERVICE = "ipsec";
     field public static final String JOB_SCHEDULER_SERVICE = "jobscheduler";
     field public static final String KEYGUARD_SERVICE = "keyguard";
+    field @FlaggedApi("android.security.keystore_grant_api") public static final String KEYSTORE_SERVICE = "keystore";
     field public static final String LAUNCHER_APPS_SERVICE = "launcherapps";
     field @UiContext public static final String LAYOUT_INFLATER_SERVICE = "layout_inflater";
     field public static final String LOCALE_SERVICE = "locale";
@@ -16103,6 +16104,7 @@
     enum_constant public static final android.graphics.ColorSpace.Named CIE_LAB;
     enum_constant public static final android.graphics.ColorSpace.Named CIE_XYZ;
     enum_constant public static final android.graphics.ColorSpace.Named DCI_P3;
+    enum_constant @FlaggedApi("com.android.graphics.flags.display_bt2020_colorspace") public static final android.graphics.ColorSpace.Named DISPLAY_BT2020;
     enum_constant public static final android.graphics.ColorSpace.Named DISPLAY_P3;
     enum_constant public static final android.graphics.ColorSpace.Named EXTENDED_SRGB;
     enum_constant public static final android.graphics.ColorSpace.Named LINEAR_EXTENDED_SRGB;
@@ -18619,6 +18621,7 @@
     field public static final int DATASPACE_BT709 = 281083904; // 0x10c10000
     field public static final int DATASPACE_DCI_P3 = 155844608; // 0x94a0000
     field public static final int DATASPACE_DEPTH = 4096; // 0x1000
+    field @FlaggedApi("com.android.graphics.flags.display_bt2020_colorspace") public static final int DATASPACE_DISPLAY_BT2020 = 142999552; // 0x8860000
     field public static final int DATASPACE_DISPLAY_P3 = 143261696; // 0x88a0000
     field public static final int DATASPACE_DYNAMIC_DEPTH = 4098; // 0x1002
     field public static final int DATASPACE_HEIF = 4100; // 0x1004
@@ -40164,6 +40167,14 @@
     method @NonNull public android.security.keystore.KeyProtection.Builder setUserPresenceRequired(boolean);
   }
 
+  @FlaggedApi("android.security.keystore_grant_api") public class KeyStoreManager {
+    method @NonNull public java.util.List<java.security.cert.X509Certificate> getGrantedCertificateChainFromId(long) throws android.security.keystore.KeyPermanentlyInvalidatedException, java.security.UnrecoverableKeyException;
+    method @NonNull public java.security.Key getGrantedKeyFromId(long) throws android.security.keystore.KeyPermanentlyInvalidatedException, java.security.UnrecoverableKeyException;
+    method @NonNull public java.security.KeyPair getGrantedKeyPairFromId(long) throws android.security.keystore.KeyPermanentlyInvalidatedException, java.security.UnrecoverableKeyException;
+    method public long grantKeyAccess(@NonNull String, int) throws android.security.KeyStoreException, java.security.UnrecoverableKeyException;
+    method public void revokeKeyAccess(@NonNull String, int) throws android.security.KeyStoreException, java.security.UnrecoverableKeyException;
+  }
+
   public class SecureKeyImportUnavailableException extends java.security.ProviderException {
     ctor public SecureKeyImportUnavailableException();
     ctor public SecureKeyImportUnavailableException(String);
@@ -52646,6 +52657,8 @@
     ctor public SurfaceView(android.content.Context, android.util.AttributeSet, int);
     ctor public SurfaceView(android.content.Context, android.util.AttributeSet, int, int);
     method public void applyTransactionToFrame(@NonNull android.view.SurfaceControl.Transaction);
+    method @FlaggedApi("android.view.flags.surface_view_get_surface_package") public void clearChildSurfacePackage();
+    method @FlaggedApi("android.view.flags.surface_view_get_surface_package") @Nullable public android.view.SurfaceControlViewHost.SurfacePackage getChildSurfacePackage();
     method @FlaggedApi("android.view.flags.surface_view_set_composition_order") public int getCompositionOrder();
     method public android.view.SurfaceHolder getHolder();
     method @Deprecated @Nullable public android.os.IBinder getHostToken();
diff --git a/core/java/android/animation/AnimationHandler.java b/core/java/android/animation/AnimationHandler.java
index 4fc90ae..d84a4c1 100644
--- a/core/java/android/animation/AnimationHandler.java
+++ b/core/java/android/animation/AnimationHandler.java
@@ -81,6 +81,25 @@
      */
     private final ArrayList<WeakReference<Object>> mAnimatorRequestors = new ArrayList<>();
 
+    /**
+     * The callbacks which will invoke {@link Animator#notifyEndListeners(boolean)} on next frame.
+     * It is only used if {@link Animator#setPostNotifyEndListenerEnabled(boolean)} sets true.
+     */
+    private ArrayList<Runnable> mPendingEndAnimationListeners;
+
+    /**
+     * The value of {@link Choreographer#getVsyncId()} at the last animation frame.
+     * It is only used if {@link Animator#setPostNotifyEndListenerEnabled(boolean)} sets true.
+     */
+    private long mLastAnimationFrameVsyncId;
+
+    /**
+     * The value of {@link Choreographer#getVsyncId()} when calling
+     * {@link Animator#notifyEndListeners(boolean)}.
+     * It is only used if {@link Animator#setPostNotifyEndListenerEnabled(boolean)} sets true.
+     */
+    private long mEndAnimationFrameVsyncId;
+
     private final Choreographer.FrameCallback mFrameCallback = new Choreographer.FrameCallback() {
         @Override
         public void doFrame(long frameTimeNanos) {
@@ -332,6 +351,39 @@
         }
     }
 
+    /**
+     * Returns the vsyncId of last animation frame if the given {@param currentVsyncId} matches
+     * the vsyncId from the end callback of animation. Otherwise it returns the given vsyncId.
+     * It only takes effect if {@link #postEndAnimationCallback(Runnable)} is called.
+     */
+    public long getLastAnimationFrameVsyncId(long currentVsyncId) {
+        return currentVsyncId == mEndAnimationFrameVsyncId && mLastAnimationFrameVsyncId != 0
+                ? mLastAnimationFrameVsyncId : currentVsyncId;
+    }
+
+    /** Runs the given callback on next frame to notify the end of the animation. */
+    public void postEndAnimationCallback(Runnable notifyEndAnimation) {
+        if (mPendingEndAnimationListeners == null) {
+            mPendingEndAnimationListeners = new ArrayList<>();
+        }
+        mPendingEndAnimationListeners.add(notifyEndAnimation);
+        if (mPendingEndAnimationListeners.size() > 1) {
+            return;
+        }
+        final Choreographer choreographer = Choreographer.getInstance();
+        mLastAnimationFrameVsyncId = choreographer.getVsyncId();
+        getProvider().postFrameCallback(frame -> {
+            mEndAnimationFrameVsyncId = choreographer.getVsyncId();
+            // The animation listeners can only get vsyncId of last animation frame in this frame
+            // by getLastAnimationFrameVsyncId(currentVsyncId).
+            while (mPendingEndAnimationListeners.size() > 0) {
+                mPendingEndAnimationListeners.remove(0).run();
+            }
+            mEndAnimationFrameVsyncId = 0;
+            mLastAnimationFrameVsyncId = 0;
+        });
+    }
+
     private void doAnimationFrame(long frameTime) {
         long currentTime = SystemClock.uptimeMillis();
         final int size = mAnimationCallbacks.size();
diff --git a/core/java/android/animation/Animator.java b/core/java/android/animation/Animator.java
index c58624e..d1eb8e8 100644
--- a/core/java/android/animation/Animator.java
+++ b/core/java/android/animation/Animator.java
@@ -16,6 +16,7 @@
 
 package android.animation;
 
+import android.annotation.CallSuper;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.TestApi;
@@ -23,6 +24,7 @@
 import android.content.pm.ActivityInfo.Config;
 import android.content.res.ConstantState;
 import android.os.Build;
+import android.os.Trace;
 import android.util.LongArray;
 
 import java.util.ArrayList;
@@ -73,6 +75,13 @@
     private static long sBackgroundPauseDelay = 1000;
 
     /**
+     * If true, when the animation plays normally to the end, the callback
+     * {@link AnimatorListener#onAnimationEnd(Animator)} will be scheduled on the next frame.
+     * It is to avoid the last animation frame being delayed by the implementation of listeners.
+     */
+    static boolean sPostNotifyEndListenerEnabled;
+
+    /**
      * A cache of the values in a list. Used so that when calling the list, we have a copy
      * of it in case the list is modified while iterating. The array can be reused to avoid
      * allocation on every notification.
@@ -124,6 +133,14 @@
     }
 
     /**
+     * @see #sPostNotifyEndListenerEnabled
+     * @hide
+     */
+    public static void setPostNotifyEndListenerEnabled(boolean enable) {
+        sPostNotifyEndListenerEnabled = enable;
+    }
+
+    /**
      * Starts this animation. If the animation has a nonzero startDelay, the animation will start
      * running after that delay elapses. A non-delayed animation will have its initial
      * value(s) set immediately, followed by calls to
@@ -635,6 +652,28 @@
         }
     }
 
+    void notifyEndListenersFromEndAnimation(boolean isReversing, boolean postNotifyEndListener) {
+        if (postNotifyEndListener) {
+            AnimationHandler.getInstance().postEndAnimationCallback(
+                    () -> completeEndAnimation(isReversing, "postNotifyAnimEnd"));
+        } else {
+            completeEndAnimation(isReversing, "notifyAnimEnd");
+        }
+    }
+
+    @CallSuper
+    void completeEndAnimation(boolean isReversing, String notifyListenerTraceName) {
+        final boolean useTrace = mListeners != null && Trace.isTagEnabled(Trace.TRACE_TAG_VIEW);
+        if (useTrace) {
+            Trace.traceBegin(Trace.TRACE_TAG_VIEW, notifyListenerTraceName
+                    + "-" + getClass().getSimpleName());
+        }
+        notifyEndListeners(isReversing);
+        if (useTrace) {
+            Trace.traceEnd(Trace.TRACE_TAG_VIEW);
+        }
+    }
+
     /**
      * Calls <code>call</code> for every item in <code>list</code> with <code>animator</code> and
      * <code>isReverse</code> as parameters.
diff --git a/core/java/android/animation/AnimatorSet.java b/core/java/android/animation/AnimatorSet.java
index ac37113..76098db 100644
--- a/core/java/android/animation/AnimatorSet.java
+++ b/core/java/android/animation/AnimatorSet.java
@@ -1442,6 +1442,8 @@
     }
 
     private void endAnimation() {
+        final boolean postNotifyEndListener = sPostNotifyEndListenerEnabled && mListeners != null
+                && mLastFrameTime > 0;
         mStarted = false;
         mLastFrameTime = -1;
         mFirstFrame = -1;
@@ -1453,7 +1455,12 @@
 
         // No longer receive callbacks
         removeAnimationCallback();
-        notifyEndListeners(mReversing);
+        notifyEndListenersFromEndAnimation(mReversing, postNotifyEndListener);
+    }
+
+    @Override
+    void completeEndAnimation(boolean isReversing, String notifyListenerTraceName) {
+        super.completeEndAnimation(isReversing, notifyListenerTraceName);
         removeAnimationEndListener();
         mSelfPulse = true;
         mReversing = false;
diff --git a/core/java/android/animation/ValueAnimator.java b/core/java/android/animation/ValueAnimator.java
index 5de7f38..e849aba 100644
--- a/core/java/android/animation/ValueAnimator.java
+++ b/core/java/android/animation/ValueAnimator.java
@@ -1289,6 +1289,8 @@
         if (mAnimationEndRequested) {
             return;
         }
+        final boolean postNotifyEndListener = sPostNotifyEndListenerEnabled && mListeners != null
+                && mLastFrameTime > 0;
         removeAnimationCallback();
 
         mAnimationEndRequested = true;
@@ -1303,15 +1305,20 @@
         mStartTime = -1;
         mRunning = false;
         mStarted = false;
-        notifyEndListeners(mReversing);
-        // mReversing needs to be reset *after* notifying the listeners for the end callbacks.
-        mReversing = false;
+        notifyEndListenersFromEndAnimation(mReversing, postNotifyEndListener);
         if (Trace.isTagEnabled(Trace.TRACE_TAG_VIEW)) {
             Trace.asyncTraceEnd(Trace.TRACE_TAG_VIEW, getNameForTrace(),
                     System.identityHashCode(this));
         }
     }
 
+    @Override
+    void completeEndAnimation(boolean isReversing, String notifyListenerTraceName) {
+        super.completeEndAnimation(isReversing, notifyListenerTraceName);
+        // mReversing needs to be reset *after* notifying the listeners for the end callbacks.
+        mReversing = false;
+    }
+
     /**
      * Called internally to start an animation by adding it to the active animations list. Must be
      * called on the UI thread.
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index 4ef5b51..64aa705 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -11176,7 +11176,7 @@
     }
 
     /**
-     * A Notification Style used to to define a notification whose expanded state includes
+     * A Notification Style used to define a notification whose expanded state includes
      * a highly customizable progress bar with segments, points, a custom tracker icon,
      * and custom icons at the start and end of the progress bar.
      *
diff --git a/core/java/android/app/NotificationChannel.java b/core/java/android/app/NotificationChannel.java
index 4a2b016..ebe7b3a 100644
--- a/core/java/android/app/NotificationChannel.java
+++ b/core/java/android/app/NotificationChannel.java
@@ -38,6 +38,7 @@
 import android.service.notification.NotificationListenerService;
 import android.text.TextUtils;
 import android.util.Log;
+import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
 
 import com.android.internal.util.Preconditions;
@@ -1369,12 +1370,17 @@
         if (sound == null || Uri.EMPTY.equals(sound)) {
             return null;
         }
-        Uri canonicalSound = getCanonicalizedSoundUri(context.getContentResolver(), sound);
-        if (canonicalSound == null) {
-            // The content provider does not support canonical uris so we backup the default
+        try {
+            Uri canonicalSound = getCanonicalizedSoundUri(context.getContentResolver(), sound);
+            if (canonicalSound == null) {
+                // The content provider does not support canonical uris so we backup the default
+                return Settings.System.DEFAULT_NOTIFICATION_URI;
+            }
+            return canonicalSound;
+        } catch (Exception e) {
+            Slog.e(TAG, "Cannot find file for sound " + sound + " using default");
             return Settings.System.DEFAULT_NOTIFICATION_URI;
         }
-        return canonicalSound;
     }
 
     /**
diff --git a/core/java/android/app/OWNERS b/core/java/android/app/OWNERS
index ba71afb..6e4c28f 100644
--- a/core/java/android/app/OWNERS
+++ b/core/java/android/app/OWNERS
@@ -3,50 +3,54 @@
 per-file ContextImpl.java = *
 
 # ActivityManager
-per-file ActivityManager* = file:/services/core/java/com/android/server/am/OWNERS
-per-file *ApplicationStartInfo* = file:/services/core/java/com/android/server/am/OWNERS
-per-file ApplicationErrorReport* = file:/services/core/java/com/android/server/am/OWNERS
-per-file ApplicationExitInfo* = file:/services/core/java/com/android/server/am/OWNERS
-per-file Application.java = file:/services/core/java/com/android/server/am/OWNERS
-per-file ApplicationLoaders.java = file:/services/core/java/com/android/server/am/OWNERS
-per-file ApplicationThreadConstants.java = file:/services/core/java/com/android/server/am/OWNERS
-per-file ContentProviderHolder* = file:/services/core/java/com/android/server/am/OWNERS
-per-file *ForegroundService* = file:/services/core/java/com/android/server/am/OWNERS
-per-file IActivityController.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IActivityManager.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IApplicationThread.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IAppTraceRetriever.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IForegroundServiceObserver.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IInstrumentationWatcher.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IntentService.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IServiceConnection.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IStopUserCallback.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file IUidObserver.aidl = file:/services/core/java/com/android/server/am/OWNERS
-per-file LoadedApk.java = file:/services/core/java/com/android/server/am/OWNERS
-per-file LocalActivityManager.java = file:/services/core/java/com/android/server/am/OWNERS
-per-file PendingIntent* = file:/services/core/java/com/android/server/am/OWNERS
-per-file *Process* = file:/services/core/java/com/android/server/am/OWNERS
-per-file ProfilerInfo* = file:/services/core/java/com/android/server/am/OWNERS
-per-file Service* = file:/services/core/java/com/android/server/am/OWNERS
-per-file SystemServiceRegistry.java = file:/services/core/java/com/android/server/am/OWNERS
-per-file *UserSwitchObserver* = file:/services/core/java/com/android/server/am/OWNERS
+per-file ActivityManager* = file:/ACTIVITY_MANAGER_OWNERS
+per-file Application.java = file:/ACTIVITY_MANAGER_OWNERS
+per-file ApplicationErrorReport* = file:/ACTIVITY_MANAGER_OWNERS
+per-file ApplicationLoaders.java = file:/ACTIVITY_MANAGER_OWNERS
+per-file ApplicationThreadConstants.java = file:/ACTIVITY_MANAGER_OWNERS
+per-file ContentProviderHolder* = file:/ACTIVITY_MANAGER_OWNERS
+per-file *ForegroundService* = file:/ACTIVITY_MANAGER_OWNERS
+per-file IActivityController.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IActivityManager.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IApplicationThread.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IAppTraceRetriever.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IForegroundServiceObserver.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IInstrumentationWatcher.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IntentService.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IServiceConnection.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IStopUserCallback.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file IUidObserver.aidl = file:/ACTIVITY_MANAGER_OWNERS
+per-file LoadedApk.java = file:/ACTIVITY_MANAGER_OWNERS
+per-file LocalActivityManager.java = file:/ACTIVITY_MANAGER_OWNERS
+per-file PendingIntent* = file:/ACTIVITY_MANAGER_OWNERS
+per-file *Process* = file:/ACTIVITY_MANAGER_OWNERS
+per-file ProfilerInfo* = file:/ACTIVITY_MANAGER_OWNERS
+per-file Service* = file:/ACTIVITY_MANAGER_OWNERS
+per-file SystemServiceRegistry.java = file:/ACTIVITY_MANAGER_OWNERS
+per-file *UserSwitchObserver* = file:/ACTIVITY_MANAGER_OWNERS
+
+# UI Automation
 per-file *UiAutomation* = file:/services/accessibility/OWNERS
 per-file *UiAutomation* = file:/core/java/android/permission/OWNERS
+
+# Game Manager
 per-file GameManager* = file:/GAME_MANAGER_OWNERS
 per-file GameMode* = file:/GAME_MANAGER_OWNERS
 per-file GameState* = file:/GAME_MANAGER_OWNERS
 per-file IGameManager* = file:/GAME_MANAGER_OWNERS
 per-file IGameMode* = file:/GAME_MANAGER_OWNERS
+
+# Background Starts
 per-file BackgroundStartPrivileges.java = file:/BAL_OWNERS
 per-file activity_manager.aconfig = file:/ACTIVITY_MANAGER_OWNERS
 
 # ActivityThread
-per-file ActivityThread.java = file:/services/core/java/com/android/server/am/OWNERS
+per-file ActivityThread.java = file:/ACTIVITY_MANAGER_OWNERS
 per-file ActivityThread.java = file:/services/core/java/com/android/server/wm/OWNERS
 per-file ActivityThread.java = file:RESOURCES_OWNERS
 
 # Alarm
-per-file *Alarm* = file:/apex/jobscheduler/OWNERS
+per-file *Alarm* = file:/apex/jobscheduler/ALARM_OWNERS
 
 # AppOps
 per-file *AppOp* = file:/core/java/android/permission/OWNERS
@@ -97,6 +101,8 @@
 
 # Performance
 per-file PropertyInvalidatedCache.java = file:/PERFORMANCE_OWNERS
+per-file *ApplicationStartInfo* = file:/PERFORMANCE_OWNERS
+per-file ApplicationExitInfo* = file:/PERFORMANCE_OWNERS
 per-file performance.aconfig = file:/PERFORMANCE_OWNERS
 
 # Pinner
diff --git a/core/java/android/app/PropertyInvalidatedCache.java b/core/java/android/app/PropertyInvalidatedCache.java
index 7481764..e4d3baa 100644
--- a/core/java/android/app/PropertyInvalidatedCache.java
+++ b/core/java/android/app/PropertyInvalidatedCache.java
@@ -21,13 +21,11 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.TestApi;
-import android.os.Binder;
 import android.os.Build;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
 import android.os.ParcelFileDescriptor;
-import android.os.Process;
 import android.os.SystemClock;
 import android.os.SystemProperties;
 import android.text.TextUtils;
@@ -80,15 +78,10 @@
         public abstract @Nullable R apply(@NonNull Q query);
 
         /**
-         * Return true if a query should not use the cache. The default implementation returns true
-         * if the process UID differs from the calling UID. This is to prevent a binder caller from
-         * reading a cached value created due to a different binder caller, when processes are
-         * caching on behalf of other processes.
+         * Return true if a query should not use the cache.  The default implementation
+         * always uses the cache.
          */
         public boolean shouldBypassCache(@NonNull Q query) {
-            if(android.multiuser.Flags.propertyInvalidatedCacheBypassMismatchedUids()) {
-                return Binder.getCallingUid() != Process.myUid();
-            }
             return false;
         }
     };
@@ -240,9 +233,9 @@
     // This is the initial value of all cache keys.  It is changed when a cache is invalidated.
     private static final int NONCE_UNSET = 0;
     // This value is used in two ways.  First, it is used internally to indicate that the cache is
-    // disabled for the current query.  Secondly, it is used to global disable the cache across the
-    // entire system.  Once a cache is disabled, there is no way to enable it again.  The global
-    // behavior is unused and will likely be removed in the future.
+    // disabled for the current query.  Secondly, it is used to globally disable the cache across
+    // the entire system.  Once a cache is disabled, there is no way to enable it again.  The
+    // global behavior is unused and will likely be removed in the future.
     private static final int NONCE_DISABLED = 1;
     // The cache is corked, which means that clients must act as though the cache is always
     // invalid.  This is used when the server is processing updates that continuously invalidate
@@ -264,10 +257,22 @@
     private static final String[] sNonceName =
             new String[]{ "unset", "disabled", "corked", "bypass" };
 
+    // The standard tag for logging.
     private static final String TAG = "PropertyInvalidatedCache";
+
+    // Set this true to enable very chatty logging.  Never commit this true.
     private static final boolean DEBUG = false;
+
+    // Set this true to enable cache verification.  On every cache hit, the cache will compare the
+    // cached value to a value pulled directly from the source.  This completely negates any
+    // performance advantage of the cache.  Enable it only to test if a particular cache is not
+    // being properly invalidated.
     private static final boolean VERIFY = false;
 
+    // The test mode. This is only used to ensure that the test functions setTestMode() and
+    // testPropertyName() are used correctly.
+    private static boolean sTestMode = false;
+
     /**
      * The object-private lock.
      */
@@ -410,15 +415,56 @@
         @GuardedBy("mLock")
         private int mCorks = 0;
 
-        // The methods to get and set a nonce from whatever storage is being used.
-        abstract long getNonce();
-        abstract void setNonce(long value);
+        // True if this handler is in test mode.  If it is in test mode, then nonces are stored
+        // and retrieved from mTestNonce.
+        @GuardedBy("mLock")
+        private boolean mTestMode = false;
+
+        /**
+         * The local value of the handler, used during testing but also used directly by the
+         * NonceLocal handler.
+         */
+        @GuardedBy("mLock")
+        protected long mTestNonce = NONCE_UNSET;
+
+        /**
+         * The methods to get and set a nonce from whatever storage is being used.  mLock may be
+         * held when these methods are called.  Implementations that take locks must behave as
+         * though mLock could be held.
+         */
+        abstract long getNonceInternal();
+        abstract void setNonceInternal(long value);
 
         NonceHandler(@NonNull String name) {
             mName = name;
         }
 
         /**
+         * Get a nonce from storage.  If the handler is in test mode, the nonce is returned from
+         * the local mTestNonce.
+         */
+        long getNonce() {
+            synchronized (mLock) {
+                if (mTestMode) return mTestNonce;
+            }
+            return getNonceInternal();
+        }
+
+        /**
+         * Write a nonce to storage.  If the handler is in test mode, the nonce is written to the
+         * local mTestNonce and storage is not affected.
+         */
+        void setNonce(long val) {
+            synchronized (mLock) {
+                if (mTestMode) {
+                    mTestNonce = val;
+                    return;
+                }
+            }
+            setNonceInternal(val);
+        }
+
+        /**
          * Write the invalidation nonce for the property.
          */
         void invalidate() {
@@ -528,6 +574,10 @@
             }
         }
 
+        /**
+         * Globally (that is, system-wide) disable all caches that use this key.  There is no way
+         * to re-enable these caches.
+         */
         void disable() {
             if (!sEnabled) {
                 return;
@@ -537,6 +587,21 @@
             }
         }
 
+        /**
+         * Put this handler in or out of test mode.  Regardless of the current and next mode, the
+         * test nonce variable is reset to UNSET.
+         */
+        void setTestMode(boolean mode) {
+            synchronized (mLock) {
+                mTestMode = mode;
+                mTestNonce = NONCE_UNSET;
+            }
+        }
+
+        /**
+         * Return the statistics associated with the key.  These statistics are not associated
+         * with any individual cache.
+         */
         record Stats(int invalidated, int corkedInvalidates) {}
         Stats getStats() {
             synchronized (mLock) {
@@ -556,21 +621,32 @@
             super(name);
         }
 
+        /**
+         * Retrieve the nonce from the system property.  If the handle is null, this method
+         * attempts to create a handle.  If handle creation fails, the method returns UNSET.  If
+         * the handle is not null, the method returns a value read via the handle.  This read
+         * occurs outside any lock.
+         */
         @Override
-        long getNonce() {
+        long getNonceInternal() {
             if (mHandle == null) {
                 synchronized (mLock) {
-                    mHandle = SystemProperties.find(mName);
                     if (mHandle == null) {
-                        return NONCE_UNSET;
+                        mHandle = SystemProperties.find(mName);
+                        if (mHandle == null) {
+                            return NONCE_UNSET;
+                        }
                     }
                 }
             }
             return mHandle.getLong(NONCE_UNSET);
         }
 
+        /**
+         * Write a nonce to a system property.
+         */
         @Override
-        void setNonce(long value) {
+        void setNonceInternal(long value) {
             // Failing to set the nonce is a fatal error.  Failures setting a system property have
             // been reported; given that the failure is probably transient, this function includes
             // a retry.
@@ -607,44 +683,34 @@
 
     /**
      * SystemProperties and shared storage are protected and cannot be written by random
-     * processes.  So, for testing purposes, the NonceTest handler stores the nonce locally.
+     * processes.  So, for testing purposes, the NonceLocal handler stores the nonce locally.  The
+     * NonceLocal uses the mTestNonce in the superclass, regardless of test mode.
      */
-    private static class NonceTest extends NonceHandler {
+    private static class NonceLocal extends NonceHandler {
         // The saved nonce.
         private long mValue;
 
-        // If this flag is false, the handler has been shutdown during a test.  Access to the
-        // handler in this state is an error.
-        private boolean mIsActive = true;
-
-        NonceTest(@NonNull String name) {
+        NonceLocal(@NonNull String name) {
             super(name);
         }
 
-        void shutdown() {
-            // The handler has been discarded as part of test cleanup.  Further access is an
-            // error.
-            mIsActive = false;
+        @Override
+        long getNonceInternal() {
+            return mTestNonce;
         }
 
         @Override
-        long getNonce() {
-            if (!mIsActive) {
-                throw new IllegalStateException("handler " + mName + " is shutdown");
-            }
-            return mValue;
-        }
-
-        @Override
-        void setNonce(long value) {
-            if (!mIsActive) {
-                throw new IllegalStateException("handler " + mName + " is shutdown");
-            }
-            mValue = value;
+        void setNonceInternal(long value) {
+            mTestNonce = value;
         }
     }
 
     /**
+     * Complete key prefixes.
+     */
+    private static final String PREFIX_TEST = CACHE_KEY_PREFIX + "." + MODULE_TEST + ".";
+
+    /**
      * A static list of nonce handlers, indexed by name.  NonceHandlers can be safely shared by
      * multiple threads, and can therefore be shared by multiple instances of the same cache, and
      * with static calls (see {@link #invalidateCache}.  Addition and removal are guarded by the
@@ -662,8 +728,8 @@
             synchronized (sGlobalLock) {
                 h = sHandlers.get(name);
                 if (h == null) {
-                    if (name.startsWith("cache_key.test.")) {
-                        h = new NonceTest(name);
+                    if (name.startsWith(PREFIX_TEST)) {
+                        h = new NonceLocal(name);
                     } else {
                         h = new NonceSysprop(name);
                     }
@@ -776,44 +842,55 @@
     }
 
     /**
-     * Enable or disable testing.  At this time, no action is taken when testing begins.
+     * Enable or disable testing.  The protocol requires that the mode toggle: for instance, it is
+     * illegal to clear the test mode if the test mode is already off.  The purpose is solely to
+     * ensure that test clients do not forget to use the test mode properly, even though the
+     * current logic does not care.
      * @hide
      */
     @TestApi
     public static void setTestMode(boolean mode) {
-        if (mode) {
-            // No action when testing begins.
-        } else {
-            resetAfterTest();
+        synchronized (sGlobalLock) {
+            if (sTestMode == mode) {
+                throw new IllegalStateException("cannot set test mode redundantly: mode=" + mode);
+            }
+            sTestMode = mode;
+            if (mode) {
+                // No action when testing begins.
+            } else {
+                resetAfterTestLocked();
+            }
         }
     }
 
     /**
-     * Enable testing the specific cache key.  This is a legacy API that will be removed as part of
-     * b/360897450.
+     * Clean up when testing ends. All handlers are reset out of test mode.  NonceLocal handlers
+     * (MODULE_TEST) are reset to the NONCE_UNSET state.  This has no effect on any other handlers
+     * that were not originally in test mode.
+     */
+    @GuardedBy("sGlobalLock")
+    private static void resetAfterTestLocked() {
+        for (Iterator<String> e = sHandlers.keys().asIterator(); e.hasNext(); ) {
+            String s = e.next();
+            final NonceHandler h = sHandlers.get(s);
+            h.setTestMode(false);
+        }
+    }
+
+    /**
+     * Enable testing the specific cache key.  This API allows a test process to invalidate caches
+     * for which it would not otherwise have permission.  Caches in test mode do NOT write their
+     * values to the system properties.  The effect is local to the current process.  Test mode
+     * must be true when this method is called.
      * @hide
      */
     @TestApi
     public void testPropertyName() {
-    }
-
-    /**
-     * Clean up when testing ends. All NonceTest handlers are erased from the global list and are
-     * poisoned, just in case the test program has retained a handle to one of the associated
-     * caches.
-     * @hide
-     */
-    @VisibleForTesting
-    public static void resetAfterTest() {
         synchronized (sGlobalLock) {
-            for (Iterator<String> e = sHandlers.keys().asIterator(); e.hasNext(); ) {
-                String s = e.next();
-                final NonceHandler h = sHandlers.get(s);
-                if (h instanceof NonceTest t) {
-                    t.shutdown();
-                    sHandlers.remove(s);
-                }
+            if (sTestMode == false) {
+                throw new IllegalStateException("cannot test property name with test mode off");
             }
+            mNonce.setTestMode(true);
         }
     }
 
@@ -1119,8 +1196,9 @@
     }
 
     /**
-     * Non-static convenience version of invalidateCache() for situations in which only a single
-     * PropertyInvalidatedCache is keyed on a particular property value.
+     * Non-static version of invalidateCache() for situations in which a cache instance is
+     * available.  This is slightly faster than than the static versions because it does not have
+     * to look up the NonceHandler for a given property name.
      * @hide
      */
     @TestApi
diff --git a/core/java/android/app/SystemServiceRegistry.java b/core/java/android/app/SystemServiceRegistry.java
index cad96e3..bd26db5 100644
--- a/core/java/android/app/SystemServiceRegistry.java
+++ b/core/java/android/app/SystemServiceRegistry.java
@@ -76,7 +76,6 @@
 import android.compat.Compatibility;
 import android.compat.annotation.ChangeId;
 import android.compat.annotation.EnabledAfter;
-import android.compat.annotation.EnabledSince;
 import android.content.ClipboardManager;
 import android.content.ContentCaptureOptions;
 import android.content.Context;
@@ -172,8 +171,7 @@
 import android.net.PacProxyManager;
 import android.net.TetheringManager;
 import android.net.VpnManager;
-import android.net.vcn.IVcnManagementService;
-import android.net.vcn.VcnManager;
+import android.net.vcn.VcnFrameworkInitializer;
 import android.net.wifi.WifiFrameworkInitializer;
 import android.net.wifi.nl80211.WifiNl80211Manager;
 import android.net.wifi.sharedconnectivity.app.SharedConnectivityManager;
@@ -208,7 +206,6 @@
 import android.os.ServiceManager.ServiceNotFoundException;
 import android.os.StatsFrameworkInitializer;
 import android.os.SystemConfigManager;
-import android.os.SystemProperties;
 import android.os.SystemUpdateManager;
 import android.os.SystemVibrator;
 import android.os.SystemVibratorManager;
@@ -239,6 +236,7 @@
 import android.security.advancedprotection.IAdvancedProtectionService;
 import android.security.attestationverification.AttestationVerificationManager;
 import android.security.attestationverification.IAttestationVerificationManagerService;
+import android.security.keystore.KeyStoreManager;
 import android.service.oemlock.IOemLockService;
 import android.service.oemlock.OemLockManager;
 import android.service.persistentdata.IPersistentDataBlockService;
@@ -302,18 +300,6 @@
     public static boolean sEnableServiceNotFoundWtf = false;
 
     /**
-     * Starting with {@link VANILLA_ICE_CREAM}, Telephony feature flags
-     * (e.g. {@link PackageManager#FEATURE_TELEPHONY_SUBSCRIPTION}) are being checked before
-     * returning managers that depend on them. If the feature is missing,
-     * {@link Context#getSystemService} will return null.
-     *
-     * This change is specific to VcnManager.
-     */
-    @ChangeId
-    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.VANILLA_ICE_CREAM)
-    static final long ENABLE_CHECKING_TELEPHONY_FEATURES_FOR_VCN = 330902016;
-
-    /**
      * After {@link Build.VERSION_CODES.VANILLA_ICE_CREAM}, Wear devices will be allowed to publish
      * no {@link GameManager} instance. This is because the respective system service is no longer
      * started for Wear devices given that the applications of the service do not currently apply to
@@ -323,16 +309,6 @@
     @EnabledAfter(targetSdkVersion = Build.VERSION_CODES.VANILLA_ICE_CREAM)
     static final long NULL_GAME_MANAGER_IN_WEAR = 340929737;
 
-    /**
-     * The corresponding vendor API for Android V
-     *
-     * <p>Starting with Android V, the vendor API format has switched to YYYYMM.
-     *
-     * @see <a href="https://preview.source.android.com/docs/core/architecture/api-flags">Vendor API
-     *     level</a>
-     */
-    private static final int VENDOR_API_FOR_ANDROID_V = 202404;
-
     // Service registry information.
     // This information is never changed once static initialization has completed.
     private static final Map<Class<?>, String> SYSTEM_SERVICE_NAMES =
@@ -499,22 +475,6 @@
                 return new VpnManager(ctx, service);
             }});
 
-        registerService(Context.VCN_MANAGEMENT_SERVICE, VcnManager.class,
-                new CachedServiceFetcher<VcnManager>() {
-            @Override
-            public VcnManager createService(ContextImpl ctx) throws ServiceNotFoundException {
-                final String telephonyFeatureToCheck = getVcnFeatureDependency();
-
-                if (telephonyFeatureToCheck != null
-                    && !ctx.getPackageManager().hasSystemFeature(telephonyFeatureToCheck)) {
-                    return null;
-                }
-
-                IBinder b = ServiceManager.getService(Context.VCN_MANAGEMENT_SERVICE);
-                IVcnManagementService service = IVcnManagementService.Stub.asInterface(b);
-                return new VcnManager(ctx, service);
-            }});
-
         registerService(Context.COUNTRY_DETECTOR, CountryDetector.class,
                 new StaticServiceFetcher<CountryDetector>() {
             @Override
@@ -1746,6 +1706,17 @@
                     }
                 });
 
+        registerService(Context.KEYSTORE_SERVICE, KeyStoreManager.class,
+                new StaticServiceFetcher<KeyStoreManager>() {
+                    @Override
+                    public KeyStoreManager createService()
+                            throws ServiceNotFoundException {
+                        if (!android.security.Flags.keystoreGrantApi()) {
+                            throw new ServiceNotFoundException("KeyStoreManager is not supported");
+                        }
+                        return KeyStoreManager.getInstance();
+                    }});
+
         registerService(Context.CONTACT_KEYS_SERVICE, E2eeContactKeysManager.class,
                 new CachedServiceFetcher<E2eeContactKeysManager>() {
                     @Override
@@ -1829,6 +1800,8 @@
             OnDevicePersonalizationFrameworkInitializer.registerServiceWrappers();
             DeviceLockFrameworkInitializer.registerServiceWrappers();
             VirtualizationFrameworkInitializer.registerServiceWrappers();
+            VcnFrameworkInitializer.registerServiceWrappers();
+
             if (com.android.server.telecom.flags.Flags.telecomMainlineBlockedNumbersManager()) {
                 ProviderFrameworkInitializer.registerServiceWrappers();
             }
@@ -1890,30 +1863,6 @@
         return manager.hasSystemFeature(featureName);
     }
 
-    // Suppressing AndroidFrameworkCompatChange because we're querying vendor
-    // partition SDK level, not application's target SDK version (which BTW we
-    // also check through Compatibility framework a few lines below).
-    @SuppressWarnings("AndroidFrameworkCompatChange")
-    @Nullable
-    private static String getVcnFeatureDependency() {
-        // Check SDK version of the client app. Apps targeting pre-V SDK might
-        // have not checked for existence of these features.
-        if (!Compatibility.isChangeEnabled(ENABLE_CHECKING_TELEPHONY_FEATURES_FOR_VCN)) {
-            return null;
-        }
-
-        // Check SDK version of the vendor partition. Pre-V devices might have
-        // incorrectly under-declared telephony features.
-        final int vendorApiLevel = SystemProperties.getInt(
-                "ro.vendor.api_level", Build.VERSION.DEVICE_INITIAL_SDK_INT);
-        if (vendorApiLevel < VENDOR_API_FOR_ANDROID_V) {
-            return PackageManager.FEATURE_TELEPHONY;
-        } else {
-            return PackageManager.FEATURE_TELEPHONY_SUBSCRIPTION;
-        }
-
-    }
-
     /**
      * Gets a system service from a given context.
      * @hide
diff --git a/core/java/android/content/Context.java b/core/java/android/content/Context.java
index ffa3375..07106e8 100644
--- a/core/java/android/content/Context.java
+++ b/core/java/android/content/Context.java
@@ -4770,6 +4770,18 @@
 
     /**
      * Use with {@link #getSystemService(String)} to retrieve a {@link
+     * android.security.keystore.KeyStoreManager} for accessing
+     * <a href="/privacy-and-security/keystore">Android Keystore</a>
+     * functions.
+     *
+     * @see #getSystemService(String)
+     * @see android.security.keystore.KeyStoreManager
+     */
+    @FlaggedApi(android.security.Flags.FLAG_KEYSTORE_GRANT_API)
+    public static final String KEYSTORE_SERVICE = "keystore";
+
+    /**
+     * Use with {@link #getSystemService(String)} to retrieve a {@link
      * android.os.storage.StorageManager} for accessing system storage
      * functions.
      *
diff --git a/core/java/android/content/OWNERS b/core/java/android/content/OWNERS
index 6d9dc45..392f62a 100644
--- a/core/java/android/content/OWNERS
+++ b/core/java/android/content/OWNERS
@@ -1,11 +1,11 @@
 # Remain no owner because multiple modules may touch this file.
 per-file Context.java = *
 per-file ContextWrapper.java = *
-per-file *Content* = file:/services/core/java/com/android/server/am/OWNERS
-per-file *Sync* = file:/services/core/java/com/android/server/am/OWNERS
+per-file *Content* = varunshah@google.com, yamasani@google.com
+per-file *Sync* = file:/apex/jobscheduler/JOB_OWNERS
 per-file IntentFilter.java = file:/PACKAGE_MANAGER_OWNERS
 per-file UriRelativeFilter* = file:/PACKAGE_MANAGER_OWNERS
-per-file IntentFilter.java = file:/services/core/java/com/android/server/am/OWNERS
+per-file IntentFilter.java = file:/INTENT_OWNERS
 per-file Intent.java = file:/INTENT_OWNERS
 per-file AutofillOptions* = file:/core/java/android/service/autofill/OWNERS
 per-file ContentCaptureOptions* = file:/core/java/android/service/contentcapture/OWNERS
diff --git a/core/java/android/content/pm/multiuser.aconfig b/core/java/android/content/pm/multiuser.aconfig
index fa3bc9e..35f9cff 100644
--- a/core/java/android/content/pm/multiuser.aconfig
+++ b/core/java/android/content/pm/multiuser.aconfig
@@ -533,3 +533,13 @@
     purpose: PURPOSE_BUGFIX
   }
 }
+
+flag {
+  name: "ignore_restrictions_when_deleting_private_profile"
+  namespace: "multiuser"
+  description: "Ignore any user restrictions when deleting private profiles."
+  bug: "350953833"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/core/java/android/hardware/DataSpace.java b/core/java/android/hardware/DataSpace.java
index 312bfdf..6117384 100644
--- a/core/java/android/hardware/DataSpace.java
+++ b/core/java/android/hardware/DataSpace.java
@@ -15,9 +15,12 @@
  */
 package android.hardware;
 
+import android.annotation.FlaggedApi;
 import android.annotation.IntDef;
 import android.view.SurfaceControl;
 
+import com.android.graphics.flags.Flags;
+
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 
@@ -639,6 +642,18 @@
      */
     public static final int DATASPACE_SRGB_LINEAR = 138477568;
 
+    /**
+     * Display BT. 2020 encoding.
+     *
+     * <p>Composed of the following -</p>
+     * <pre>
+     *   Primaries: STANDARD_BT2020
+     *   Transfer: TRANSFER_SRGB
+     *   Range: RANGE_FULL</pre>
+     */
+    @FlaggedApi(Flags.FLAG_DISPLAY_BT2020_COLORSPACE)
+    public static final int DATASPACE_DISPLAY_BT2020 = 142999552;
+
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
     @IntDef(flag = true, value = {
@@ -660,7 +675,8 @@
         DATASPACE_BT2020,
         DATASPACE_BT709,
         DATASPACE_DCI_P3,
-        DATASPACE_SRGB_LINEAR
+        DATASPACE_SRGB_LINEAR,
+        DATASPACE_DISPLAY_BT2020
     })
     public @interface NamedDataSpace {};
 
diff --git a/core/java/android/hardware/display/DisplayManager.java b/core/java/android/hardware/display/DisplayManager.java
index 97f6899..b0ea92d 100644
--- a/core/java/android/hardware/display/DisplayManager.java
+++ b/core/java/android/hardware/display/DisplayManager.java
@@ -18,6 +18,7 @@
 
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.HdrCapabilities.HdrType;
+import static android.view.Display.INVALID_DISPLAY;
 
 import android.Manifest;
 import android.annotation.FlaggedApi;
@@ -47,6 +48,7 @@
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.ServiceManager;
+import android.os.UserManager;
 import android.util.Log;
 import android.util.Pair;
 import android.util.Slog;
@@ -96,6 +98,8 @@
     @GuardedBy("mLock")
     private final WeakDisplayCache mDisplayCache = new WeakDisplayCache();
 
+    private int mDisplayIdToMirror = INVALID_DISPLAY;
+
     /**
      * Broadcast receiver that indicates when the Wifi display status changes.
      * <p>
@@ -1086,6 +1090,7 @@
         if (surface != null) {
             builder.setSurface(surface);
         }
+        builder.setDisplayIdToMirror(getDisplayIdToMirror());
         return createVirtualDisplay(builder.build(), handler, callback);
     }
 
@@ -1163,6 +1168,7 @@
         if (surface != null) {
             builder.setSurface(surface);
         }
+        builder.setDisplayIdToMirror(getDisplayIdToMirror());
         return createVirtualDisplay(projection, builder.build(), callback, handler);
     }
 
@@ -1708,6 +1714,16 @@
         return mGlobal.getDefaultDozeBrightness(displayId);
     }
 
+    private int getDisplayIdToMirror() {
+        if (mDisplayIdToMirror == INVALID_DISPLAY) {
+            final UserManager userManager = mContext.getSystemService(UserManager.class);
+            mDisplayIdToMirror = userManager.isVisibleBackgroundUsersSupported()
+                    ? userManager.getMainDisplayIdAssignedToUser()
+                    : DEFAULT_DISPLAY;
+        }
+        return mDisplayIdToMirror;
+    }
+
     /**
      * Listens for changes in available display devices.
      */
diff --git a/core/java/android/net/vcn/VcnFrameworkInitializer.java b/core/java/android/net/vcn/VcnFrameworkInitializer.java
new file mode 100644
index 0000000..8cb213b
--- /dev/null
+++ b/core/java/android/net/vcn/VcnFrameworkInitializer.java
@@ -0,0 +1,106 @@
+/*
+ * 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.net.vcn;
+
+import android.annotation.Nullable;
+import android.app.SystemServiceRegistry;
+import android.compat.Compatibility;
+import android.compat.annotation.ChangeId;
+import android.compat.annotation.EnabledSince;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.os.Build;
+import android.os.SystemProperties;
+
+/**
+ * Class for performing registration for VCN service.
+ *
+ * @hide
+ */
+// TODO: Expose it as @SystemApi(client = MODULE_LIBRARIES)
+public final class VcnFrameworkInitializer {
+    /**
+     * Starting with {@link VANILLA_ICE_CREAM}, Telephony feature flags (e.g. {@link
+     * PackageManager#FEATURE_TELEPHONY_SUBSCRIPTION}) are being checked before returning managers
+     * that depend on them. If the feature is missing, {@link Context#getSystemService} will return
+     * null.
+     *
+     * <p>This change is specific to VcnManager.
+     */
+    @ChangeId
+    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.VANILLA_ICE_CREAM)
+    private static final long ENABLE_CHECKING_TELEPHONY_FEATURES_FOR_VCN = 330902016;
+
+    /**
+     * The corresponding vendor API for Android V
+     *
+     * <p>Starting with Android V, the vendor API format has switched to YYYYMM.
+     *
+     * @see <a href="https://preview.source.android.com/docs/core/architecture/api-flags">Vendor API
+     *     level</a>
+     */
+    private static final int VENDOR_API_FOR_ANDROID_V = 202404;
+
+    private VcnFrameworkInitializer() {}
+
+    // Suppressing AndroidFrameworkCompatChange because we're querying vendor
+    // partition SDK level, not application's target SDK version (which BTW we
+    // also check through Compatibility framework a few lines below).
+    @Nullable
+    private static String getVcnFeatureDependency() {
+        // Check SDK version of the client app. Apps targeting pre-V SDK might
+        // have not checked for existence of these features.
+        if (!Compatibility.isChangeEnabled(ENABLE_CHECKING_TELEPHONY_FEATURES_FOR_VCN)) {
+            return null;
+        }
+
+        // Check SDK version of the vendor partition. Pre-V devices might have
+        // incorrectly under-declared telephony features.
+        final int vendorApiLevel =
+                SystemProperties.getInt(
+                        "ro.vendor.api_level", Build.VERSION.DEVICE_INITIAL_SDK_INT);
+        if (vendorApiLevel < VENDOR_API_FOR_ANDROID_V) {
+            return PackageManager.FEATURE_TELEPHONY;
+        } else {
+            return PackageManager.FEATURE_TELEPHONY_SUBSCRIPTION;
+        }
+    }
+
+    /**
+     * Register VCN service to {@link Context}, so that {@link Context#getSystemService} can return
+     * a VcnManager.
+     *
+     * @throws IllegalStateException if this is called anywhere besides {@link
+     *     SystemServiceRegistry}.
+     */
+    public static void registerServiceWrappers() {
+        SystemServiceRegistry.registerContextAwareService(
+                VcnManager.VCN_MANAGEMENT_SERVICE_STRING,
+                VcnManager.class,
+                (context, serviceBinder) -> {
+                    final String telephonyFeatureToCheck = getVcnFeatureDependency();
+                    if (telephonyFeatureToCheck != null
+                            && !context.getPackageManager()
+                                    .hasSystemFeature(telephonyFeatureToCheck)) {
+                        return null;
+                    }
+                    IVcnManagementService service =
+                            IVcnManagementService.Stub.asInterface(serviceBinder);
+                    return new VcnManager(context, service);
+                });
+    }
+}
diff --git a/core/java/android/net/vcn/flags.aconfig b/core/java/android/net/vcn/flags.aconfig
index efddd1f..5b30624 100644
--- a/core/java/android/net/vcn/flags.aconfig
+++ b/core/java/android/net/vcn/flags.aconfig
@@ -14,11 +14,4 @@
     namespace: "vcn"
     description: "Feature flag for adjustable safe mode timeout"
     bug: "317406085"
-}
-
-flag{
-    name: "network_metric_monitor"
-    namespace: "vcn"
-    description: "Feature flag for enabling network metric monitor"
-    bug: "282996138"
 }
\ No newline at end of file
diff --git a/core/java/android/os/BatteryUsageStats.java b/core/java/android/os/BatteryUsageStats.java
index a698b9d..ddcb5c6 100644
--- a/core/java/android/os/BatteryUsageStats.java
+++ b/core/java/android/os/BatteryUsageStats.java
@@ -1031,7 +1031,10 @@
             return this;
         }
 
-        private long getStatsDuration() {
+        /**
+         * Returns the duration of the battery session reflected by these stats.
+         */
+        public long getStatsDuration() {
             if (mStatsDurationMs != -1) {
                 return mStatsDurationMs;
             } else {
diff --git a/core/java/android/os/BatteryUsageStatsQuery.java b/core/java/android/os/BatteryUsageStatsQuery.java
index b533225..e68c4ca 100644
--- a/core/java/android/os/BatteryUsageStatsQuery.java
+++ b/core/java/android/os/BatteryUsageStatsQuery.java
@@ -20,6 +20,8 @@
 import android.annotation.NonNull;
 import android.util.IntArray;
 
+import com.android.internal.os.MonotonicClock;
+
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 
@@ -85,8 +87,11 @@
     @NonNull
     private final int[] mUserIds;
     private final long mMaxStatsAgeMs;
-    private final long mFromTimestamp;
-    private final long mToTimestamp;
+
+    private final long mAggregatedFromTimestamp;
+    private final long mAggregatedToTimestamp;
+    private long mMonotonicStartTime;
+    private long mMonotonicEndTime;
     private final double mMinConsumedPowerThreshold;
     private final @BatteryConsumer.PowerComponentId int[] mPowerComponents;
 
@@ -96,8 +101,10 @@
                 : new int[]{UserHandle.USER_ALL};
         mMaxStatsAgeMs = builder.mMaxStatsAgeMs;
         mMinConsumedPowerThreshold = builder.mMinConsumedPowerThreshold;
-        mFromTimestamp = builder.mFromTimestamp;
-        mToTimestamp = builder.mToTimestamp;
+        mAggregatedFromTimestamp = builder.mAggregateFromTimestamp;
+        mAggregatedToTimestamp = builder.mAggregateToTimestamp;
+        mMonotonicStartTime = builder.mMonotonicStartTime;
+        mMonotonicEndTime = builder.mMonotonicEndTime;
         mPowerComponents = builder.mPowerComponents;
     }
 
@@ -163,11 +170,27 @@
     }
 
     /**
-     * Returns the exclusive lower bound of the stored snapshot timestamps that should be included
-     * in the aggregation.  Ignored if {@link #getToTimestamp()} is zero.
+     * Returns the exclusive lower bound of the battery history that should be included in
+     * the aggregated battery usage stats.
      */
-    public long getFromTimestamp() {
-        return mFromTimestamp;
+    public long getMonotonicStartTime() {
+        return mMonotonicStartTime;
+    }
+
+    /**
+     * Returns the inclusive upper bound of the battery history that should be included in
+     * the aggregated battery usage stats.
+     */
+    public long getMonotonicEndTime() {
+        return mMonotonicEndTime;
+    }
+
+    /**
+     * Returns the exclusive lower bound of the stored snapshot timestamps that should be included
+     * in the aggregation.  Ignored if {@link #getAggregatedToTimestamp()} is zero.
+     */
+    public long getAggregatedFromTimestamp() {
+        return mAggregatedFromTimestamp;
     }
 
     /**
@@ -175,8 +198,8 @@
      * be included in the aggregation.  The default is to include only the current stats
      * accumulated since the latest battery reset.
      */
-    public long getToTimestamp() {
-        return mToTimestamp;
+    public long getAggregatedToTimestamp() {
+        return mAggregatedToTimestamp;
     }
 
     private BatteryUsageStatsQuery(Parcel in) {
@@ -185,8 +208,8 @@
         in.readIntArray(mUserIds);
         mMaxStatsAgeMs = in.readLong();
         mMinConsumedPowerThreshold = in.readDouble();
-        mFromTimestamp = in.readLong();
-        mToTimestamp = in.readLong();
+        mAggregatedFromTimestamp = in.readLong();
+        mAggregatedToTimestamp = in.readLong();
         mPowerComponents = in.createIntArray();
     }
 
@@ -197,8 +220,8 @@
         dest.writeIntArray(mUserIds);
         dest.writeLong(mMaxStatsAgeMs);
         dest.writeDouble(mMinConsumedPowerThreshold);
-        dest.writeLong(mFromTimestamp);
-        dest.writeLong(mToTimestamp);
+        dest.writeLong(mAggregatedFromTimestamp);
+        dest.writeLong(mAggregatedToTimestamp);
         dest.writeIntArray(mPowerComponents);
     }
 
@@ -228,8 +251,10 @@
         private int mFlags;
         private IntArray mUserIds;
         private long mMaxStatsAgeMs = DEFAULT_MAX_STATS_AGE_MS;
-        private long mFromTimestamp;
-        private long mToTimestamp;
+        private long mMonotonicStartTime = MonotonicClock.UNDEFINED;
+        private long mMonotonicEndTime = MonotonicClock.UNDEFINED;
+        private long mAggregateFromTimestamp;
+        private long mAggregateToTimestamp;
         private double mMinConsumedPowerThreshold = 0;
         private @BatteryConsumer.PowerComponentId int[] mPowerComponents;
 
@@ -241,6 +266,17 @@
         }
 
         /**
+         * Specifies the time range for the requested stats, in terms of MonotonicClock
+         * @param monotonicStartTime Inclusive starting monotonic timestamp
+         * @param monotonicEndTime Exclusive ending timestamp. Can be MonotonicClock.UNDEFINED
+         */
+        public Builder monotonicTimeRange(long monotonicStartTime, long monotonicEndTime) {
+            mMonotonicStartTime = monotonicStartTime;
+            mMonotonicEndTime = monotonicEndTime;
+            return this;
+        }
+
+        /**
          * Add a user whose battery stats should be included in the battery usage stats.
          * {@link UserHandle#USER_ALL} will be used by default if no users are added explicitly.
          */
@@ -345,8 +381,8 @@
          */
         // TODO(b/298459065): switch to monotonic clock
         public Builder aggregateSnapshots(long fromTimestamp, long toTimestamp) {
-            mFromTimestamp = fromTimestamp;
-            mToTimestamp = toTimestamp;
+            mAggregateFromTimestamp = fromTimestamp;
+            mAggregateToTimestamp = toTimestamp;
             return this;
         }
 
diff --git a/core/java/android/os/ConcurrentMessageQueue/MessageQueue.java b/core/java/android/os/ConcurrentMessageQueue/MessageQueue.java
index b2d9260..6afb8e0 100644
--- a/core/java/android/os/ConcurrentMessageQueue/MessageQueue.java
+++ b/core/java/android/os/ConcurrentMessageQueue/MessageQueue.java
@@ -754,7 +754,7 @@
                 // Idle handles only run if the queue is empty or if the first message
                 // in the queue (possibly a barrier) is due to be handled in the future.
                 if (pendingIdleHandlerCount < 0
-                        && mNextPollTimeoutMillis != 0) {
+                        && isIdle()) {
                     pendingIdleHandlerCount = mIdleHandlers.size();
                 }
                 if (pendingIdleHandlerCount <= 0) {
diff --git a/core/java/android/os/OWNERS b/core/java/android/os/OWNERS
index a1b75034..590ddb4 100644
--- a/core/java/android/os/OWNERS
+++ b/core/java/android/os/OWNERS
@@ -122,3 +122,6 @@
 per-file StatsBootstrapAtomValue.aidl = file:/services/core/java/com/android/server/stats/OWNERS
 per-file StatsBootstrapAtomService.java = file:/services/core/java/com/android/server/stats/OWNERS
 per-file StatsServiceManager.java = file:/services/core/java/com/android/server/stats/OWNERS
+
+# Dropbox
+per-file DropBoxManager* = mwachens@google.com
diff --git a/core/java/android/os/SemiConcurrentMessageQueue/MessageQueue.java b/core/java/android/os/SemiConcurrentMessageQueue/MessageQueue.java
index 80c24a9..02335972 100644
--- a/core/java/android/os/SemiConcurrentMessageQueue/MessageQueue.java
+++ b/core/java/android/os/SemiConcurrentMessageQueue/MessageQueue.java
@@ -712,7 +712,7 @@
                 // Idle handles only run if the queue is empty or if the first message
                 // in the queue (possibly a barrier) is due to be handled in the future.
                 if (pendingIdleHandlerCount < 0
-                        && mNextPollTimeoutMillis != 0) {
+                        && isIdle()) {
                     pendingIdleHandlerCount = mIdleHandlers.size();
                 }
                 if (pendingIdleHandlerCount <= 0) {
diff --git a/core/java/android/os/StatsBootstrapAtomValue.aidl b/core/java/android/os/StatsBootstrapAtomValue.aidl
index b59bc06..b31eb6f 100644
--- a/core/java/android/os/StatsBootstrapAtomValue.aidl
+++ b/core/java/android/os/StatsBootstrapAtomValue.aidl
@@ -19,12 +19,36 @@
  *
  * @hide
  */
-union StatsBootstrapAtomValue {
-    boolean boolValue;
-    int intValue;
-    long longValue;
-    float floatValue;
-    String stringValue;
-    byte[] bytesValue;
-    String[] stringArrayValue;
-}
\ No newline at end of file
+parcelable StatsBootstrapAtomValue {
+    union Primitive {
+        boolean boolValue;
+        int intValue;
+        long longValue;
+        float floatValue;
+        String stringValue;
+        byte[] bytesValue;
+	String[] stringArrayValue;
+    }
+
+    Primitive value;
+
+    parcelable Annotation {
+        // Match the definitions in
+        // packages/modules/StatsD/framework/java/android/util/StatsLog.java
+        // Only supports UIDs for now.
+        @Backing(type="byte")
+        enum Id {
+            NONE,
+            IS_UID,
+        }
+        Id id;
+
+        union Primitive {
+            boolean boolValue;
+            int intValue;
+        }
+        Primitive value;
+    }
+
+    Annotation[] annotations;
+}
diff --git a/core/java/android/os/UidBatteryConsumer.java b/core/java/android/os/UidBatteryConsumer.java
index 7f7ef04..f893739 100644
--- a/core/java/android/os/UidBatteryConsumer.java
+++ b/core/java/android/os/UidBatteryConsumer.java
@@ -109,9 +109,11 @@
      * Returns the amount of time in milliseconds this UID spent in the specified process state.
      */
     public long getTimeInProcessStateMs(@ProcessState int state) {
-        Key key = getKey(POWER_COMPONENT_BASE, state);
-        if (key != null) {
-            return getUsageDurationMillis(key);
+        if (state != BatteryConsumer.PROCESS_STATE_UNSPECIFIED) {
+            Key key = getKey(POWER_COMPONENT_BASE, state);
+            if (key != null) {
+                return getUsageDurationMillis(key);
+            }
         }
         return 0;
     }
diff --git a/core/java/android/os/UserManager.java b/core/java/android/os/UserManager.java
index bd3da0d..fa99f35 100644
--- a/core/java/android/os/UserManager.java
+++ b/core/java/android/os/UserManager.java
@@ -772,9 +772,10 @@
     public static final String DISALLOW_CONFIG_CREDENTIALS = "no_config_credentials";
 
     /**
-     * When set on the admin user this specifies if the user can remove users.
+     * When set on the admin user this specifies if the user can remove secondary users. Managed
+     * profiles and private profiles can still be removed even if this is set on the admin user.
      * When set on a non-admin secondary user, this specifies if the user can remove itself.
-     * This restriction has no effect on managed profiles.
+     * This restriction has no effect when set on managed profiles.
      * The default value is <code>false</code>.
      *
      * <p>Holders of the permission
@@ -793,7 +794,8 @@
      * Specifies if managed profiles of this user can be removed, other than by its profile owner.
      * The default value is <code>false</code>.
      * <p>
-     * This restriction has no effect on managed profiles.
+     * This restriction has no effect on managed profiles, and this restriction does not block the
+     * removal of private profiles of this user.
      *
      * <p>Key for user restrictions.
      * <p>Type: Boolean
@@ -5600,14 +5602,30 @@
             android.Manifest.permission.MANAGE_USERS,
             android.Manifest.permission.INTERACT_ACROSS_USERS
     })
+    @CachedProperty(api = "user_manager_users")
     public @Nullable UserHandle getProfileParent(@NonNull UserHandle user) {
-        UserInfo info = getProfileParent(user.getIdentifier());
-
-        if (info == null) {
-            return null;
+        if (android.multiuser.Flags.cacheProfileParentReadOnly()) {
+            final UserHandle userHandle = UserManagerCache.getProfileParent(
+                    (UserHandle query) -> {
+                        UserInfo info = getProfileParent(query.getIdentifier());
+                        // TODO: Remove when b/372923336 is fixed
+                        if (info == null) {
+                            return UserHandle.of(UserHandle.USER_NULL);
+                        }
+                        return UserHandle.of(info.id);
+                    },
+                    user);
+            if (userHandle.getIdentifier() == UserHandle.USER_NULL) {
+                return null;
+            }
+            return userHandle;
+        } else {
+            UserInfo info = getProfileParent(user.getIdentifier());
+            if (info == null) {
+                return null;
+            }
+            return UserHandle.of(info.id);
         }
-
-        return UserHandle.of(info.id);
     }
 
     /**
@@ -6422,6 +6440,9 @@
      */
     public static final void invalidateCacheOnUserListChange() {
         UserManagerCache.invalidateUserSerialNumber();
+        if (android.multiuser.Flags.cacheProfileParentReadOnly()) {
+            UserManagerCache.invalidateProfileParent();
+        }
     }
 
     /**
diff --git a/core/java/android/security/flags.aconfig b/core/java/android/security/flags.aconfig
index aedf8e0..1d35344 100644
--- a/core/java/android/security/flags.aconfig
+++ b/core/java/android/security/flags.aconfig
@@ -113,3 +113,10 @@
     description: "AFL feature"
     bug: "365994454"
 }
+
+flag {
+    name: "keystore_grant_api"
+    namespace: "hardware_backed_security"
+    description: "Feature flag for exposing KeyStore grant APIs"
+    bug: "351158708"
+}
diff --git a/core/java/android/telephony/TelephonyCallback.java b/core/java/android/telephony/TelephonyCallback.java
index 14d5800..5295b60 100644
--- a/core/java/android/telephony/TelephonyCallback.java
+++ b/core/java/android/telephony/TelephonyCallback.java
@@ -33,6 +33,7 @@
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.annotations.WeaklyReferencedCallback;
 import com.android.internal.telephony.IPhoneStateListener;
 import com.android.internal.telephony.flags.Flags;
 
@@ -70,6 +71,7 @@
  * its manifest file. Where permissions apply, they are noted in the
  * appropriate sub-interfaces.
  */
+@WeaklyReferencedCallback
 public class TelephonyCallback {
     private static final String LOG_TAG = "TelephonyCallback";
     /**
diff --git a/core/java/android/view/DisplayEventReceiver.java b/core/java/android/view/DisplayEventReceiver.java
index 18080e4..fc7a65d 100644
--- a/core/java/android/view/DisplayEventReceiver.java
+++ b/core/java/android/view/DisplayEventReceiver.java
@@ -24,6 +24,7 @@
 import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.annotations.WeaklyReferencedCallback;
 
 import dalvik.annotation.optimization.FastNative;
 
@@ -40,6 +41,7 @@
  *
  * @hide
  */
+@WeaklyReferencedCallback
 public abstract class DisplayEventReceiver {
 
     /**
diff --git a/core/java/android/view/LetterboxScrollProcessor.java b/core/java/android/view/LetterboxScrollProcessor.java
index dc736d6..1364a82 100644
--- a/core/java/android/view/LetterboxScrollProcessor.java
+++ b/core/java/android/view/LetterboxScrollProcessor.java
@@ -16,6 +16,8 @@
 
 package android.view;
 
+import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
+
 import android.annotation.Nullable;
 import android.content.Context;
 import android.graphics.Rect;
@@ -23,6 +25,8 @@
 
 import androidx.annotation.NonNull;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -35,6 +39,7 @@
  *
  * @hide
  */
+@VisibleForTesting(visibility = PACKAGE)
 public class LetterboxScrollProcessor {
 
     private enum LetterboxScrollState {
@@ -53,6 +58,7 @@
     /** IDs of events generated from this class */
     private final Set<Integer> mGeneratedEventIds = new HashSet<>();
 
+    @VisibleForTesting(visibility = PACKAGE)
     public LetterboxScrollProcessor(@NonNull Context context, @Nullable Handler handler) {
         mContext = context;
         mScrollDetector = new GestureDetector(context, new ScrollListener(), handler);
@@ -69,7 +75,9 @@
      * @return The list of adjusted events, or null if no adjustments are needed. The list is empty
      * if the event should be ignored. Do not keep a reference to the output as the list is reused.
      */
-    public List<MotionEvent> processMotionEvent(MotionEvent motionEvent) {
+    @Nullable
+    @VisibleForTesting(visibility = PACKAGE)
+    public List<MotionEvent> processMotionEvent(@NonNull MotionEvent motionEvent) {
         mProcessedEvents.clear();
         final Rect appBounds = getAppBounds();
 
@@ -124,11 +132,9 @@
             mState = LetterboxScrollState.AWAITING_GESTURE_START;
         }
 
-        if (makeNoAdjustments) return null;
-        return mProcessedEvents;
+        return makeNoAdjustments ? null : mProcessedEvents;
     }
 
-
     /**
      * Processes the InputEvent for compatibility before it is finished by calling
      * InputEventReceiver#finishInputEvent().
@@ -136,21 +142,33 @@
      * @param motionEvent The MotionEvent to process.
      * @return The motionEvent to finish, or null if it should not be finished.
      */
-    public InputEvent processMotionEventBeforeFinish(MotionEvent motionEvent) {
-        if (mGeneratedEventIds.remove(motionEvent.getId())) return null;
-        return motionEvent;
+    @Nullable
+    @VisibleForTesting(visibility = PACKAGE)
+    public InputEvent processMotionEventBeforeFinish(@NonNull MotionEvent motionEvent) {
+        return mGeneratedEventIds.remove(motionEvent.getId()) ? null : motionEvent;
     }
 
+    @NonNull
     private Rect getAppBounds() {
         return mContext.getResources().getConfiguration().windowConfiguration.getBounds();
     }
 
-    private boolean isOutsideAppBounds(MotionEvent motionEvent, Rect appBounds) {
-        return motionEvent.getX() < 0 || motionEvent.getX() >= appBounds.width()
-                || motionEvent.getY() < 0 || motionEvent.getY() >= appBounds.height();
+    /** Checks whether the gesture is located on the letterbox area. */
+    private boolean isOutsideAppBounds(@NonNull MotionEvent motionEvent, @NonNull Rect appBounds) {
+        // The events are in the coordinate system of the ViewRootImpl (window). The window might
+        // not have the same dimensions as the app bounds - for example in case of Dialogs - thus
+        // `getRawX()` and `getRawY()` are used, with the absolute bounds (left, top, etc) instead
+        // of width and height.
+        // The event should be passed to the app if it has happened anywhere in the app area,
+        // irrespective of the current window size, therefore the app bounds are used instead of the
+        // current window.
+        return motionEvent.getRawX() < appBounds.left
+                || motionEvent.getRawX() >= appBounds.right
+                || motionEvent.getRawY() < appBounds.top
+                || motionEvent.getRawY() >= appBounds.bottom;
     }
 
-    private void applyOffset(MotionEvent event, Rect appBounds) {
+    private void applyOffset(@NonNull MotionEvent event, @NonNull Rect appBounds) {
         float horizontalOffset = calculateOffset(event.getX(), appBounds.width());
         float verticalOffset = calculateOffset(event.getY(), appBounds.height());
         // Apply the offset to the motion event so it is over the app's view.
diff --git a/core/java/android/view/OWNERS b/core/java/android/view/OWNERS
index 1ea58bc..80484a6 100644
--- a/core/java/android/view/OWNERS
+++ b/core/java/android/view/OWNERS
@@ -88,6 +88,7 @@
 per-file OnReceiveContentListener.java = file:/core/java/android/widget/OWNERS
 per-file ContentInfo.java = file:/core/java/android/service/autofill/OWNERS
 per-file ContentInfo.java = file:/core/java/android/widget/OWNERS
+per-file view_flags.aconfig = file:/services/core/java/com/android/server/wm/OWNERS
 
 # WindowManager
 per-file ContentRecordingSession.aidl = file:/services/core/java/com/android/server/wm/OWNERS
diff --git a/core/java/android/view/RoundScrollbarRenderer.java b/core/java/android/view/RoundScrollbarRenderer.java
index 5f6d5e2..59c2598 100644
--- a/core/java/android/view/RoundScrollbarRenderer.java
+++ b/core/java/android/view/RoundScrollbarRenderer.java
@@ -16,17 +16,26 @@
 
 package android.view;
 
+import static android.util.MathUtils.acos;
+
+import static java.lang.Math.sin;
+
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.Paint;
 import android.graphics.Rect;
 import android.graphics.RectF;
+import android.os.SystemProperties;
 import android.util.DisplayMetrics;
+import android.view.flags.Flags;
 
 /**
  * Helper class for drawing round scroll bars on round Wear devices.
+ *
+ * @hide
  */
-class RoundScrollbarRenderer {
+public class RoundScrollbarRenderer {
+    private static final String BLUECHIP_ENABLED_SYSPROP = "persist.cw_build.bluechip.enabled";
     // The range of the scrollbar position represented as an angle in degrees.
     private static final float SCROLLBAR_ANGLE_RANGE = 28.8f;
     private static final float MAX_SCROLLBAR_ANGLE_SWIPE = 26.3f; // 90%
@@ -45,12 +54,15 @@
     private final Paint mTrackPaint = new Paint();
     private final RectF mRect = new RectF();
     private final View mParent;
-    private final int mMaskThickness;
+    private final float mInset;
 
     private float mPreviousMaxScroll = 0;
     private float mMaxScrollDiff = 0;
     private float mPreviousCurrentScroll = 0;
     private float mCurrentScrollDiff = 0;
+    private float mThumbStrokeWidthAsDegrees = 0;
+    private boolean mDrawToLeft;
+    private boolean mUseRefactoredRoundScrollbar;
 
     public RoundScrollbarRenderer(View parent) {
         // Paints for the round scrollbar.
@@ -69,29 +81,36 @@
         // Fetch the resource indicating the thickness of CircularDisplayMask, rounding in the same
         // way WindowManagerService.showCircularMask does. The scroll bar is inset by this amount so
         // that it doesn't get clipped.
-        mMaskThickness = parent.getContext().getResources().getDimensionPixelSize(
-                com.android.internal.R.dimen.circular_display_mask_thickness);
+        int maskThickness =
+                parent.getContext()
+                        .getResources()
+                        .getDimensionPixelSize(
+                                com.android.internal.R.dimen.circular_display_mask_thickness);
+
+        float thumbWidth = dpToPx(THUMB_WIDTH_DP);
+        mThumbPaint.setStrokeWidth(thumbWidth);
+        mTrackPaint.setStrokeWidth(thumbWidth);
+        mInset = thumbWidth / 2 + maskThickness;
+
+        mUseRefactoredRoundScrollbar =
+                Flags.useRefactoredRoundScrollbar()
+                        && SystemProperties.getBoolean(BLUECHIP_ENABLED_SYSPROP, false);
     }
 
-    public void drawRoundScrollbars(Canvas canvas, float alpha, Rect bounds, boolean drawToLeft) {
-        if (alpha == 0) {
-            return;
-        }
-        // Get information about the current scroll state of the parent view.
-        float maxScroll = mParent.computeVerticalScrollRange();
-        float scrollExtent = mParent.computeVerticalScrollExtent();
-        float newScroll = mParent.computeVerticalScrollOffset();
-
+    private float computeScrollExtent(float scrollExtent, float maxScroll) {
         if (scrollExtent <= 0) {
             if (!mParent.canScrollVertically(1) && !mParent.canScrollVertically(-1)) {
-                return;
+                return -1f;
             } else {
-                scrollExtent = 0;
+                return 0f;
             }
         } else if (maxScroll <= scrollExtent) {
-            return;
+            return -1f;
         }
+        return scrollExtent;
+    }
 
+    private void resizeGradually(float maxScroll, float newScroll) {
         // Make changes to the VerticalScrollRange happen gradually
         if (Math.abs(maxScroll - mPreviousMaxScroll) > RESIZING_THRESHOLD_PX
                 && mPreviousMaxScroll != 0) {
@@ -106,49 +125,79 @@
                 || Math.abs(mCurrentScrollDiff) > RESIZING_THRESHOLD_PX) {
             mMaxScrollDiff *= RESIZING_RATE;
             mCurrentScrollDiff *= RESIZING_RATE;
-
-            maxScroll -= mMaxScrollDiff;
-            newScroll -= mCurrentScrollDiff;
         } else {
             mMaxScrollDiff = 0;
             mCurrentScrollDiff = 0;
         }
+    }
 
-        float currentScroll = Math.max(0, newScroll);
-        float linearThumbLength = scrollExtent;
-        float thumbWidth = dpToPx(THUMB_WIDTH_DP);
-        mThumbPaint.setStrokeWidth(thumbWidth);
-        mTrackPaint.setStrokeWidth(thumbWidth);
-
-        setThumbColor(applyAlpha(DEFAULT_THUMB_COLOR, alpha));
-        setTrackColor(applyAlpha(DEFAULT_TRACK_COLOR, alpha));
-
-        // Normalize the sweep angle for the scroll bar.
-        float sweepAngle = (linearThumbLength / maxScroll) * SCROLLBAR_ANGLE_RANGE;
-        sweepAngle = clamp(sweepAngle, MIN_SCROLLBAR_ANGLE_SWIPE, MAX_SCROLLBAR_ANGLE_SWIPE);
-        // Normalize the start angle so that it falls on the track.
-        float startAngle = (currentScroll * (SCROLLBAR_ANGLE_RANGE - sweepAngle))
-                / (maxScroll - linearThumbLength) - SCROLLBAR_ANGLE_RANGE / 2f;
-        startAngle = clamp(startAngle, -SCROLLBAR_ANGLE_RANGE / 2f,
-                SCROLLBAR_ANGLE_RANGE / 2f - sweepAngle);
-
-        // Draw the track and the thumb.
-        float inset = thumbWidth / 2 + mMaskThickness;
-        mRect.set(
-                bounds.left + inset,
-                bounds.top + inset,
-                bounds.right - inset,
-                bounds.bottom - inset);
-
-        if (drawToLeft) {
-            canvas.drawArc(mRect, 180 + SCROLLBAR_ANGLE_RANGE / 2f, -SCROLLBAR_ANGLE_RANGE, false,
-                    mTrackPaint);
-            canvas.drawArc(mRect, 180 - startAngle, -sweepAngle, false, mThumbPaint);
-        } else {
-            canvas.drawArc(mRect, -SCROLLBAR_ANGLE_RANGE / 2f, SCROLLBAR_ANGLE_RANGE, false,
-                    mTrackPaint);
-            canvas.drawArc(mRect, startAngle, sweepAngle, false, mThumbPaint);
+    public void drawRoundScrollbars(Canvas canvas, float alpha, Rect bounds, boolean drawToLeft) {
+        if (alpha == 0) {
+            return;
         }
+        // Get information about the current scroll state of the parent view.
+        float maxScroll = mParent.computeVerticalScrollRange();
+        float scrollExtent = mParent.computeVerticalScrollExtent();
+        float newScroll = mParent.computeVerticalScrollOffset();
+
+        scrollExtent = computeScrollExtent(scrollExtent, maxScroll);
+        if (scrollExtent < 0f) {
+            return;
+        }
+
+        // Make changes to the VerticalScrollRange happen gradually
+        resizeGradually(maxScroll, newScroll);
+        maxScroll -= mMaxScrollDiff;
+        newScroll -= mCurrentScrollDiff;
+
+        applyThumbColor(alpha);
+
+        float sweepAngle = computeSweepAngle(scrollExtent, maxScroll);
+        float startAngle =
+                computeStartAngle(Math.max(0, newScroll), sweepAngle, maxScroll, scrollExtent);
+
+        updateBounds(bounds);
+
+        mDrawToLeft = drawToLeft;
+        drawRoundScrollbars(canvas, startAngle, sweepAngle, alpha);
+    }
+
+    private void drawRoundScrollbars(
+            Canvas canvas, float startAngle, float sweepAngle, float alpha) {
+        if (mUseRefactoredRoundScrollbar) {
+            draw(canvas, startAngle, sweepAngle, alpha);
+        } else {
+            applyTrackColor(alpha);
+            drawArc(canvas, -SCROLLBAR_ANGLE_RANGE / 2f, SCROLLBAR_ANGLE_RANGE, mTrackPaint);
+            drawArc(canvas, startAngle, sweepAngle, mThumbPaint);
+        }
+    }
+
+    /** Returns true if horizontal bounds are updated */
+    private void updateBounds(Rect bounds) {
+        mRect.set(
+                bounds.left + mInset,
+                bounds.top + mInset,
+                bounds.right - mInset,
+                bounds.bottom - mInset);
+        mThumbStrokeWidthAsDegrees =
+                getVertexAngle((mRect.right - mRect.left) / 2f, mThumbPaint.getStrokeWidth() / 2f);
+    }
+
+    private float computeSweepAngle(float scrollExtent, float maxScroll) {
+        // Normalize the sweep angle for the scroll bar.
+        float sweepAngle = (scrollExtent / maxScroll) * SCROLLBAR_ANGLE_RANGE;
+        return clamp(sweepAngle, MIN_SCROLLBAR_ANGLE_SWIPE, MAX_SCROLLBAR_ANGLE_SWIPE);
+    }
+
+    private float computeStartAngle(
+            float currentScroll, float sweepAngle, float maxScroll, float scrollExtent) {
+        // Normalize the start angle so that it falls on the track.
+        float startAngle =
+                (currentScroll * (SCROLLBAR_ANGLE_RANGE - sweepAngle)) / (maxScroll - scrollExtent)
+                        - SCROLLBAR_ANGLE_RANGE / 2f;
+        return clamp(
+                startAngle, -SCROLLBAR_ANGLE_RANGE / 2f, SCROLLBAR_ANGLE_RANGE / 2f - sweepAngle);
     }
 
     void getRoundVerticalScrollBarBounds(Rect bounds) {
@@ -164,10 +213,8 @@
     private static float clamp(float val, float min, float max) {
         if (val < min) {
             return min;
-        } else if (val > max) {
-            return max;
         } else {
-            return val;
+            return Math.min(val, max);
         }
     }
 
@@ -176,15 +223,17 @@
         return Color.argb(alphaByte, Color.red(color), Color.green(color), Color.blue(color));
     }
 
-    private void setThumbColor(int thumbColor) {
-        if (mThumbPaint.getColor() != thumbColor) {
-            mThumbPaint.setColor(thumbColor);
+    private void applyThumbColor(float alpha) {
+        int color = applyAlpha(DEFAULT_THUMB_COLOR, alpha);
+        if (mThumbPaint.getColor() != color) {
+            mThumbPaint.setColor(color);
         }
     }
 
-    private void setTrackColor(int trackColor) {
-        if (mTrackPaint.getColor() != trackColor) {
-            mTrackPaint.setColor(trackColor);
+    private void applyTrackColor(float alpha) {
+        int color = applyAlpha(DEFAULT_TRACK_COLOR, alpha);
+        if (mTrackPaint.getColor() != color) {
+            mTrackPaint.setColor(color);
         }
     }
 
@@ -192,4 +241,88 @@
         return dp * ((float) mParent.getContext().getResources().getDisplayMetrics().densityDpi)
                 / DisplayMetrics.DENSITY_DEFAULT;
     }
+
+    private static float getVertexAngle(float edge, float base) {
+        float edgeSquare = edge * edge * 2;
+        float baseSquare = base * base;
+        float gapInRadians = acos(((edgeSquare - baseSquare) / edgeSquare));
+        return (float) Math.toDegrees(gapInRadians);
+    }
+
+    private static float getKiteEdge(float knownEdge, float angleBetweenKnownEdgesInDegrees) {
+        return (float) (2 * knownEdge * sin(Math.toRadians(angleBetweenKnownEdgesInDegrees / 2)));
+    }
+
+    private void draw(Canvas canvas, float thumbStartAngle, float thumbSweepAngle, float alpha) {
+        // Draws the top arc
+        drawTrack(
+                canvas,
+                // The highest point of the top track on a vertical scale. Here the thumb width is
+                // reduced to account for the arc formed by ROUND stroke style
+                -SCROLLBAR_ANGLE_RANGE / 2f - mThumbStrokeWidthAsDegrees,
+                // The lowest point of the top track on a vertical scale. Here the thumb width is
+                // reduced twice to (a) account for the arc formed by ROUND stroke style (b) gap
+                // between thumb and top track
+                thumbStartAngle - mThumbStrokeWidthAsDegrees * 2,
+                alpha);
+        // Draws the thumb
+        drawArc(canvas, thumbStartAngle, thumbSweepAngle, mThumbPaint);
+        // Draws the bottom arc
+        drawTrack(
+                canvas,
+                // The highest point of the bottom track on a vertical scale. Here the thumb width
+                // is added twice to (a) account for the arc formed by ROUND stroke style (b) gap
+                // between thumb and bottom track
+                (thumbStartAngle + thumbSweepAngle) + mThumbStrokeWidthAsDegrees * 2,
+                // The lowest point of the top track on a vertical scale. Here the thumb width is
+                // added to account for the arc formed by ROUND stroke style
+                SCROLLBAR_ANGLE_RANGE / 2f + mThumbStrokeWidthAsDegrees,
+                alpha);
+    }
+
+    private void drawTrack(Canvas canvas, float beginAngle, float endAngle, float alpha) {
+        // Angular distance between end and begin
+        float angleBetweenEndAndBegin = endAngle - beginAngle;
+        // The sweep angle for the track is the angular distance between end and begin less the
+        // thumb width twice to account for top and bottom arc formed by the ROUND stroke style
+        float sweepAngle = angleBetweenEndAndBegin - 2 * mThumbStrokeWidthAsDegrees;
+
+        float startAngle = -1f;
+        float strokeWidth = -1f;
+        if (sweepAngle > 0f) {
+            // The angle is greater than 0 which means a normal arc should be drawn with stroke
+            // width same as the thumb. The ROUND stroke style will cover the top/bottom arc of the
+            // track
+            startAngle = beginAngle + mThumbStrokeWidthAsDegrees;
+            strokeWidth = mThumbPaint.getStrokeWidth();
+        } else if (Math.abs(sweepAngle) < 2 * mThumbStrokeWidthAsDegrees) {
+            // The sweep angle is less than 0 but is still relevant in creating a circle for the
+            // top/bottom track. The start angle is adjusted to account for being the mid point of
+            // begin / end angle.
+            startAngle = beginAngle + angleBetweenEndAndBegin / 2;
+            // The radius of this circle forms a kite with the radius of the arc drawn for the rect
+            // with the given angular difference between the arc radius which is used to compute the
+            // new stroke width.
+            strokeWidth = getKiteEdge(((mRect.right - mRect.left) / 2), angleBetweenEndAndBegin);
+            // The opacity is decreased proportionally, if the stroke width of the track is 50% or
+            // less that that of the thumb
+            alpha = alpha * Math.min(1f, 2 * strokeWidth / mThumbPaint.getStrokeWidth());
+            // As we desire a circle to be drawn, the sweep angle is set to a minimal value
+            sweepAngle = Float.MIN_NORMAL;
+        } else {
+            return;
+        }
+
+        applyTrackColor(alpha);
+        mTrackPaint.setStrokeWidth(strokeWidth);
+        drawArc(canvas, startAngle, sweepAngle, mTrackPaint);
+    }
+
+    private void drawArc(Canvas canvas, float startAngle, float sweepAngle, Paint paint) {
+        if (mDrawToLeft) {
+            canvas.drawArc(mRect, /* startAngle= */ 180 - startAngle, -sweepAngle, false, paint);
+        } else {
+            canvas.drawArc(mRect, startAngle, sweepAngle, /* useCenter= */ false, paint);
+        }
+    }
 }
diff --git a/core/java/android/view/SurfaceView.java b/core/java/android/view/SurfaceView.java
index 82235d2..9cad3e5 100644
--- a/core/java/android/view/SurfaceView.java
+++ b/core/java/android/view/SurfaceView.java
@@ -16,6 +16,7 @@
 
 package android.view;
 
+import static android.view.flags.Flags.FLAG_SURFACE_VIEW_GET_SURFACE_PACKAGE;
 import static android.view.flags.Flags.FLAG_SURFACE_VIEW_SET_COMPOSITION_ORDER;
 import static android.view.WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON;
 import static android.view.WindowManagerPolicyConstants.APPLICATION_MEDIA_OVERLAY_SUBLAYER;
@@ -27,6 +28,7 @@
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.SuppressLint;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.content.res.CompatibilityInfo.Translator;
@@ -2112,7 +2114,7 @@
     }
 
     /**
-     * Display the view-hierarchy embedded within a {@link SurfaceControlViewHost.SurfacePackage}
+     * Displays the view-hierarchy embedded within a {@link SurfaceControlViewHost.SurfacePackage}
      * within this SurfaceView.
      *
      * This can be called independently of the SurfaceView lifetime callbacks. SurfaceView
@@ -2132,6 +2134,8 @@
      * SurfaceView the underlying {@link SurfaceControlViewHost} remains managed by it's original
      * remote-owner.
      *
+     * Users can call {@link SurfaceView#clearChildSurfacePackage} to clear the package.
+     *
      * @param p The SurfacePackage to embed.
      */
     public void setChildSurfacePackage(@NonNull SurfaceControlViewHost.SurfacePackage p) {
@@ -2155,6 +2159,46 @@
         invalidate();
     }
 
+    /**
+     * Returns the {@link SurfaceControlViewHost.SurfacePackage} that was set on this SurfaceView.
+     *
+     * Note: This method will return {@code null} if
+     * {@link #setChildSurfacePackage(SurfaceControlViewHost.SurfacePackage)}
+     * has not been called or if {@link #clearChildSurfacePackage()} has been called.
+     *
+     * @see #setChildSurfacePackage(SurfaceControlViewHost.SurfacePackage)
+     */
+    @SuppressLint("GetterSetterNullability")
+    @FlaggedApi(FLAG_SURFACE_VIEW_GET_SURFACE_PACKAGE)
+    public @Nullable SurfaceControlViewHost.SurfacePackage getChildSurfacePackage() {
+        return mSurfacePackage;
+    }
+
+    /**
+     * Clears the {@link SurfaceControlViewHost.SurfacePackage} that was set on this SurfaceView.
+     * This hides any content rendered by the provided
+     * {@link SurfaceControlViewHost.SurfacePackage}.
+     *
+     * @see #setChildSurfacePackage(SurfaceControlViewHost.SurfacePackage)
+     */
+    @FlaggedApi(FLAG_SURFACE_VIEW_GET_SURFACE_PACKAGE)
+    public void clearChildSurfacePackage() {
+        if (mSurfacePackage != null) {
+            mSurfaceControlViewHostParent.detach();
+            mEmbeddedWindowParams.clear();
+
+            // Reparent the SurfaceControl to remove the content on screen.
+            final SurfaceControl sc = mSurfacePackage.getSurfaceControl();
+            final SurfaceControl.Transaction transaction = new Transaction();
+            transaction.reparent(sc, null);
+            mSurfacePackage.release();
+            applyTransactionOnVriDraw(transaction);
+
+            mSurfacePackage = null;
+            invalidate();
+        }
+    }
+
     private void reparentSurfacePackage(SurfaceControl.Transaction t,
             SurfaceControlViewHost.SurfacePackage p) {
         final SurfaceControl sc = p.getSurfaceControl();
diff --git a/core/java/android/view/flags/view_flags.aconfig b/core/java/android/view/flags/view_flags.aconfig
index 1cf26ab..1b86f96 100644
--- a/core/java/android/view/flags/view_flags.aconfig
+++ b/core/java/android/view/flags/view_flags.aconfig
@@ -116,4 +116,20 @@
     description: "Add a SurfaceView composition order control API."
     bug: "341021569"
     is_fixed_read_only: true
+}
+
+flag {
+    name: "surface_view_get_surface_package"
+    namespace: "window_surfaces"
+    description: "Add APIs to manage SurfacePackage of the parent SurfaceView."
+    bug: "341021569"
+    is_fixed_read_only: true
+}
+
+flag {
+    name: "use_refactored_round_scrollbar"
+    namespace: "wear_frameworks"
+    description: "Use refactored round scrollbar."
+    bug: "333417898"
+    is_fixed_read_only: true
 }
\ No newline at end of file
diff --git a/core/java/android/window/BackNavigationInfo.java b/core/java/android/window/BackNavigationInfo.java
index 59639d0..6cefc4d 100644
--- a/core/java/android/window/BackNavigationInfo.java
+++ b/core/java/android/window/BackNavigationInfo.java
@@ -16,6 +16,8 @@
 
 package android.window;
 
+import static android.app.ActivityTaskManager.INVALID_TASK_ID;
+
 import android.annotation.AnimRes;
 import android.annotation.ColorInt;
 import android.annotation.IntDef;
@@ -118,6 +120,7 @@
     private final Rect mTouchableRegion;
 
     private final boolean mAppProgressGenerationAllowed;
+    private final int mFocusedTaskId;
 
     /**
      * Create a new {@link BackNavigationInfo} instance.
@@ -135,7 +138,8 @@
             @Nullable CustomAnimationInfo customAnimationInfo,
             int letterboxColor,
             @Nullable Rect touchableRegion,
-            boolean appProgressGenerationAllowed) {
+            boolean appProgressGenerationAllowed,
+            int focusedTaskId) {
         mType = type;
         mOnBackNavigationDone = onBackNavigationDone;
         mOnBackInvokedCallback = onBackInvokedCallback;
@@ -145,6 +149,7 @@
         mLetterboxColor = letterboxColor;
         mTouchableRegion = new Rect(touchableRegion);
         mAppProgressGenerationAllowed = appProgressGenerationAllowed;
+        mFocusedTaskId = focusedTaskId;
     }
 
     private BackNavigationInfo(@NonNull Parcel in) {
@@ -157,6 +162,7 @@
         mLetterboxColor = in.readInt();
         mTouchableRegion = in.readTypedObject(Rect.CREATOR);
         mAppProgressGenerationAllowed = in.readBoolean();
+        mFocusedTaskId = in.readInt();
     }
 
     /** @hide */
@@ -171,6 +177,7 @@
         dest.writeInt(mLetterboxColor);
         dest.writeTypedObject(mTouchableRegion, flags);
         dest.writeBoolean(mAppProgressGenerationAllowed);
+        dest.writeInt(mFocusedTaskId);
     }
 
     /**
@@ -238,6 +245,14 @@
     }
 
     /**
+     * @return The focused task id when back gesture start.
+     * @hide
+     */
+    public int getFocusedTaskId() {
+        return mFocusedTaskId;
+    }
+
+    /**
      * Callback to be called when the back preview is finished in order to notify the server that
      * it can clean up the resources created for the animation.
      * @hide
@@ -435,6 +450,7 @@
         private int mLetterboxColor = Color.TRANSPARENT;
         private Rect mTouchableRegion;
         private boolean mAppProgressGenerationAllowed;
+        private int mFocusedTaskId = INVALID_TASK_ID;
 
         /**
          * @see BackNavigationInfo#getType()
@@ -527,6 +543,14 @@
         }
 
         /**
+         * @param focusedTaskId The current focused taskId when back gesture start.
+         */
+        public Builder setFocusedTaskId(int focusedTaskId) {
+            mFocusedTaskId = focusedTaskId;
+            return this;
+        }
+
+        /**
          * Builds and returns an instance of {@link BackNavigationInfo}
          */
         public BackNavigationInfo build() {
@@ -537,7 +561,8 @@
                     mCustomAnimationInfo,
                     mLetterboxColor,
                     mTouchableRegion,
-                    mAppProgressGenerationAllowed);
+                    mAppProgressGenerationAllowed,
+                    mFocusedTaskId);
         }
     }
 }
diff --git a/core/java/android/window/SnapshotDrawerUtils.java b/core/java/android/window/SnapshotDrawerUtils.java
index 9a7bce0..5397da1 100644
--- a/core/java/android/window/SnapshotDrawerUtils.java
+++ b/core/java/android/window/SnapshotDrawerUtils.java
@@ -151,7 +151,9 @@
         @VisibleForTesting
         public void setFrames(Rect frame, Rect systemBarInsets) {
             mFrame.set(frame);
-            mSizeMismatch = (mFrame.width() != mSnapshotW || mFrame.height() != mSnapshotH);
+            final Rect letterboxInsets = mSnapshot.getLetterboxInsets();
+            mSizeMismatch = (mFrame.width() != mSnapshotW || mFrame.height() != mSnapshotH)
+                    || letterboxInsets.left != 0 || letterboxInsets.top != 0;
             if (!Flags.drawSnapshotAspectRatioMatch() && systemBarInsets != null) {
                 mSystemBarInsets.set(systemBarInsets);
                 mSystemBarBackgroundPainter.setInsets(systemBarInsets);
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index 0d235ff..d15f52c 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -9,6 +9,16 @@
 }
 
 flag {
+  name: "reset_draw_state_on_client_invisible"
+  namespace: "windowing_frontend"
+  description: "Reset draw state if the client is notified to be invisible"
+  bug: "373023636"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
+
+flag {
     name: "wait_for_transition_on_display_switch"
     namespace: "windowing_frontend"
     description: "Waits for Shell transition to start before unblocking the screen after display switch"
diff --git a/core/java/com/android/internal/jank/FrameTracker.java b/core/java/com/android/internal/jank/FrameTracker.java
index d474c6d..6448f10 100644
--- a/core/java/com/android/internal/jank/FrameTracker.java
+++ b/core/java/com/android/internal/jank/FrameTracker.java
@@ -30,6 +30,7 @@
 import static com.android.internal.jank.InteractionJankMonitor.ACTION_SESSION_END;
 import static com.android.internal.jank.InteractionJankMonitor.EXECUTOR_TASK_TIMEOUT;
 
+import android.animation.AnimationHandler;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -344,7 +345,8 @@
     @UiThread
     public boolean end(@Reasons int reason) {
         if (mCancelled || mEndVsyncId != INVALID_ID) return false;
-        mEndVsyncId = mChoreographer.getVsyncId();
+        mEndVsyncId = AnimationHandler.getInstance().getLastAnimationFrameVsyncId(
+                mChoreographer.getVsyncId());
         // Cancel the session if:
         // 1. The session begins and ends at the same vsync id.
         // 2. The session never begun.
diff --git a/core/java/com/android/internal/os/BatteryStatsHistory.java b/core/java/com/android/internal/os/BatteryStatsHistory.java
index d61785e..07aa720 100644
--- a/core/java/com/android/internal/os/BatteryStatsHistory.java
+++ b/core/java/com/android/internal/os/BatteryStatsHistory.java
@@ -83,7 +83,7 @@
     private static final String TAG = "BatteryStatsHistory";
 
     // Current on-disk Parcel version. Must be updated when the format of the parcelable changes
-    private static final int VERSION = 210;
+    private static final int VERSION = 211;
 
     private static final String HISTORY_DIR = "battery-history";
     private static final String FILE_SUFFIX = ".bh";
@@ -210,6 +210,8 @@
     private final MonotonicClock mMonotonicClock;
     // Monotonic time when we started writing to the history buffer
     private long mHistoryBufferStartTime;
+    // Monotonically increasing size of written history
+    private long mMonotonicHistorySize;
     private final ArraySet<PowerStats.Descriptor> mWrittenPowerStatsDescriptors = new ArraySet<>();
     private byte mLastHistoryStepLevel = 0;
     private boolean mMutable = true;
@@ -909,6 +911,8 @@
                 }
                 // skip monotonic time field.
                 p.readLong();
+                // skip monotonic size field
+                p.readLong();
 
                 final int bufSize = p.readInt();
                 final int curPos = p.dataPosition();
@@ -964,6 +968,8 @@
         }
         // skip monotonic time field.
         out.readLong();
+        // skip monotonic size field
+        out.readLong();
         return true;
     }
 
@@ -987,6 +993,7 @@
         p.setDataPosition(0);
         p.readInt();        // Skip the version field
         long monotonicTime = p.readLong();
+        p.readLong();       // Skip monotonic size field
         p.setDataPosition(pos);
         return monotonicTime;
     }
@@ -1819,6 +1826,7 @@
             // as long as no bit has changed both between now and the last entry, as
             // well as the last entry and the one before it (so we capture any toggles).
             if (DEBUG) Slog.i(TAG, "ADD: rewinding back to " + mHistoryBufferLastPos);
+            mMonotonicHistorySize -= (mHistoryBuffer.dataSize() - mHistoryBufferLastPos);
             mHistoryBuffer.setDataSize(mHistoryBufferLastPos);
             mHistoryBuffer.setDataPosition(mHistoryBufferLastPos);
             mHistoryBufferLastPos = -1;
@@ -1934,6 +1942,7 @@
         }
         mHistoryLastWritten.tagsFirstOccurrence = hasTags;
         writeHistoryDelta(mHistoryBuffer, mHistoryLastWritten, mHistoryLastLastWritten);
+        mMonotonicHistorySize += (mHistoryBuffer.dataSize() - mHistoryBufferLastPos);
         cur.wakelockTag = null;
         cur.wakeReasonTag = null;
         cur.eventCode = HistoryItem.EVENT_NONE;
@@ -2344,6 +2353,8 @@
             }
 
             mHistoryBufferStartTime = in.readLong();
+            mMonotonicHistorySize = in.readLong();
+
             mHistoryBuffer.setDataSize(0);
             mHistoryBuffer.setDataPosition(0);
 
@@ -2370,6 +2381,7 @@
     private void writeHistoryBuffer(Parcel out) {
         out.writeInt(BatteryStatsHistory.VERSION);
         out.writeLong(mHistoryBufferStartTime);
+        out.writeLong(mMonotonicHistorySize);
         out.writeInt(mHistoryBuffer.dataSize());
         if (DEBUG) {
             Slog.i(TAG, "***************** WRITING HISTORY: "
@@ -2457,6 +2469,14 @@
     }
 
     /**
+     * Returns the monotonically increasing size of written history, including the buffers
+     * that have already been discarded.
+     */
+    public long getMonotonicHistorySize() {
+        return mMonotonicHistorySize;
+    }
+
+    /**
      * Prints battery stats history for debugging.
      */
     public void dump(PrintWriter pw, long startTimeMs, long endTimeMs) {
diff --git a/core/jni/android_view_InputDevice.cpp b/core/jni/android_view_InputDevice.cpp
index f5992d9..ce40e51 100644
--- a/core/jni/android_view_InputDevice.cpp
+++ b/core/jni/android_view_InputDevice.cpp
@@ -60,7 +60,7 @@
                                                                   ? layoutInfo->layoutType.c_str()
                                                                   : NULL));
 
-    std::shared_ptr<KeyCharacterMap> map = deviceInfo.getKeyCharacterMap();
+    const KeyCharacterMap* map = deviceInfo.getKeyCharacterMap();
     std::unique_ptr<KeyCharacterMap> mapCopy;
     if (map != nullptr) {
         mapCopy = std::make_unique<KeyCharacterMap>(*map);
diff --git a/core/proto/android/app/OWNERS b/core/proto/android/app/OWNERS
index a137ea9..519bf9a 100644
--- a/core/proto/android/app/OWNERS
+++ b/core/proto/android/app/OWNERS
@@ -1,3 +1,3 @@
-per-file appstartinfo.proto = file:/services/core/java/com/android/server/am/OWNERS
+per-file appstartinfo.proto = file:/PERFORMANCE_OWNERS
 per-file location_time_zone_manager.proto = file:platform/frameworks/base:/services/core/java/com/android/server/timezonedetector/OWNERS
 per-file time_zone_detector.proto = file:platform/frameworks/base:/services/core/java/com/android/server/timezonedetector/OWNERS
diff --git a/core/res/OWNERS b/core/res/OWNERS
index d109cee..faed4d8 100644
--- a/core/res/OWNERS
+++ b/core/res/OWNERS
@@ -53,7 +53,7 @@
 per-file res/values/dimens_car.xml = file:/platform/packages/services/Car:/OWNERS
 
 # Device Idle
-per-file res/values/config_device_idle.xml = file:/apex/jobscheduler/OWNERS
+per-file res/values/config_device_idle.xml = file:/apex/jobscheduler/DEVICE_IDLE_OWNERS
 
 # Display Manager
 per-file res/values/config_display.xml = file:/services/core/java/com/android/server/display/OWNERS
diff --git a/core/res/res/values/config_battery_stats.xml b/core/res/res/values/config_battery_stats.xml
index 80cf088..9498273 100644
--- a/core/res/res/values/config_battery_stats.xml
+++ b/core/res/res/values/config_battery_stats.xml
@@ -45,7 +45,12 @@
     <integer name="config_powerStatsAggregationPeriod">14400000</integer>
 
     <!-- PowerStats aggregation span duration in milliseconds. This is the length of battery
-    history time for every aggregated power stats span that is stored stored in PowerStatsStore.
+    history time for every aggregated power stats span that is stored in PowerStatsStore.
     It should not be larger than config_powerStatsAggregationPeriod (but it can be the same) -->
     <integer name="config_aggregatedPowerStatsSpanDuration">3600000</integer>
+
+    <!-- BatteryUsageStats accumulation period as determined by the size of accumulated
+    battery history, in bytes. -->
+    <integer name="config_accumulatedBatteryUsageStatsSpanSize">32768</integer>
+
 </resources>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index b424955..0b2b345 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -5301,6 +5301,7 @@
   <java-symbol type="string" name="config_powerStatsThrottlePeriods" />
   <java-symbol type="integer" name="config_powerStatsAggregationPeriod" />
   <java-symbol type="integer" name="config_aggregatedPowerStatsSpanDuration" />
+  <java-symbol type="integer" name="config_accumulatedBatteryUsageStatsSpanSize" />
 
   <java-symbol name="materialColorOnSecondaryFixedVariant" type="attr"/>
   <java-symbol name="materialColorOnTertiaryFixedVariant" type="attr"/>
diff --git a/core/tests/coretests/src/android/animation/ValueAnimatorTests.java b/core/tests/coretests/src/android/animation/ValueAnimatorTests.java
index eb463fd..0469846 100644
--- a/core/tests/coretests/src/android/animation/ValueAnimatorTests.java
+++ b/core/tests/coretests/src/android/animation/ValueAnimatorTests.java
@@ -18,6 +18,8 @@
 
 import static android.test.MoreAsserts.assertNotEqual;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertNotNull;
@@ -892,6 +894,34 @@
     }
 
     @Test
+    public void testPostNotifyEndListener() throws Throwable {
+        ValueAnimator.setPostNotifyEndListenerEnabled(true);
+        final CountDownLatch latch = new CountDownLatch(1);
+        final long[] lastAnimFrameId = new long[1];
+        final long[] endAnimCallbackId = new long[1];
+        try {
+            a1.addUpdateListener(animator -> {
+                if (animator.getAnimatedFraction() == 1f) {
+                    lastAnimFrameId[0] = Choreographer.getInstance().getVsyncId();
+                }
+            });
+            a1.addListener(new MyListener() {
+                @Override
+                public void onAnimationEnd(Animator animation) {
+                    super.onAnimationEnd(animation);
+                    endAnimCallbackId[0] = Choreographer.getInstance().getVsyncId();
+                    latch.countDown();
+                }
+            });
+            mActivityRule.runOnUiThread(() -> a1.start());
+            assertTrue(latch.await(1, TimeUnit.SECONDS));
+            assertThat(endAnimCallbackId[0]).isGreaterThan(lastAnimFrameId[0]);
+        } finally {
+            ValueAnimator.setPostNotifyEndListenerEnabled(false);
+        }
+    }
+
+    @Test
     public void testZeroDuration() throws Throwable {
         // Run two animators with zero duration, with one running forward and the other one
         // backward. Check that the animations start and finish with the correct end fractions.
diff --git a/core/tests/coretests/src/android/app/NotificationChannelTest.java b/core/tests/coretests/src/android/app/NotificationChannelTest.java
index e19f887..e4b5407 100644
--- a/core/tests/coretests/src/android/app/NotificationChannelTest.java
+++ b/core/tests/coretests/src/android/app/NotificationChannelTest.java
@@ -27,6 +27,7 @@
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -51,6 +52,7 @@
 import android.platform.test.flag.junit.FlagsParameterization;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.MediaStore.Audio.AudioColumns;
+import android.provider.Settings;
 import android.test.mock.MockContentResolver;
 import android.util.Xml;
 
@@ -399,6 +401,29 @@
     }
 
     @Test
+    public void testWriteXmlForBackup_noAccessToFile() throws Exception {
+        Uri uri = Uri.parse("content://media/1");
+
+        AudioAttributes mAudioAttributes =
+                new AudioAttributes.Builder()
+                        .setContentType(AudioAttributes.CONTENT_TYPE_UNKNOWN)
+                        .setUsage(AudioAttributes.USAGE_NOTIFICATION)
+                        .setFlags(AudioAttributes.FLAG_AUDIBILITY_ENFORCED)
+                        .build();
+
+        NotificationChannel channel = new NotificationChannel("id", "name", 3);
+        channel.setSound(uri, mAudioAttributes);
+
+        when(mIContentProvider.canonicalize(any(), any())).thenThrow(new SecurityException(""));
+        doThrow(new SecurityException("")).when(mIContentProvider)
+                .canonicalizeAsync(any(), any(), any());
+
+        NotificationChannel restoredChannel = backUpAndRestore(channel);
+        assertThat(restoredChannel.getSound())
+                .isEqualTo(Settings.System.DEFAULT_NOTIFICATION_URI);
+    }
+
+    @Test
     public void testVibrationGetters_nonPatternBasedVibrationEffect_waveform() throws Exception {
         mSetFlagsRule.enableFlags(Flags.FLAG_NOTIFICATION_CHANNEL_VIBRATION_EFFECT_API);
         NotificationChannel channel = new NotificationChannel("id", "name", 3);
diff --git a/core/tests/coretests/src/android/app/PropertyInvalidatedCacheTests.java b/core/tests/coretests/src/android/app/PropertyInvalidatedCacheTests.java
index b5ee130..dcea5b2 100644
--- a/core/tests/coretests/src/android/app/PropertyInvalidatedCacheTests.java
+++ b/core/tests/coretests/src/android/app/PropertyInvalidatedCacheTests.java
@@ -19,6 +19,8 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 import android.platform.test.annotations.IgnoreUnderRavenwood;
 import android.platform.test.ravenwood.RavenwoodRule;
@@ -26,6 +28,7 @@
 import androidx.test.filters.SmallTest;
 
 import org.junit.After;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 
@@ -84,14 +87,20 @@
         public Boolean apply(Integer x) {
             return mServer.query(x);
         }
+
         @Override
         public boolean shouldBypassCache(Integer x) {
             return x % 13 == 0;
         }
     }
 
-    // Clear the test mode after every test, in case this process is used for other
-    // tests. This also resets the test property map.
+    // Prepare for testing.
+    @Before
+    public void setUp() throws Exception {
+        PropertyInvalidatedCache.setTestMode(true);
+    }
+
+    // Ensure all test configurations are cleared.
     @After
     public void tearDown() throws Exception {
         PropertyInvalidatedCache.setTestMode(false);
@@ -111,9 +120,6 @@
                 new PropertyInvalidatedCache<>(4, MODULE, API, "cache1",
                         new ServerQuery(tester));
 
-        PropertyInvalidatedCache.setTestMode(true);
-        testCache.testPropertyName();
-
         tester.verify(0);
         assertEquals(tester.value(3), testCache.query(3));
         tester.verify(1);
@@ -223,22 +229,16 @@
 
         TestCache(String module, String api) {
             this(module, api, new TestQuery());
-            setTestMode(true);
-            testPropertyName();
         }
 
         TestCache(String module, String api, TestQuery query) {
             super(4, module, api, api, query);
             mQuery = query;
-            setTestMode(true);
-            testPropertyName();
         }
 
         public int getRecomputeCount() {
             return mQuery.getRecomputeCount();
         }
-
-
     }
 
     @Test
@@ -375,4 +375,52 @@
             PropertyInvalidatedCache.MODULE_BLUETOOTH, "getState");
         assertEquals(n1, "cache_key.bluetooth.get_state");
     }
+
+    // Verify that test mode works properly.
+    @Test
+    public void testTestMode() {
+        // Create a cache that will write a system nonce.
+        TestCache sysCache = new TestCache(PropertyInvalidatedCache.MODULE_SYSTEM, "mode1");
+        try {
+            // Invalidate the cache, which writes the system property.  There must be a permission
+            // failure.
+            sysCache.invalidateCache();
+            fail("expected permission failure");
+        } catch (RuntimeException e) {
+            // The expected exception is a bare RuntimeException.  The test does not attempt to
+            // validate the text of the exception message.
+        }
+
+        sysCache.testPropertyName();
+        // Invalidate the cache.  This must succeed because the property has been marked for
+        // testing.
+        sysCache.invalidateCache();
+
+        // Create a cache that uses MODULE_TEST.  Invalidation succeeds whether or not the
+        // property is tagged as being tested.
+        TestCache testCache = new TestCache(PropertyInvalidatedCache.MODULE_TEST, "mode2");
+        testCache.invalidateCache();
+        testCache.testPropertyName();
+        testCache.invalidateCache();
+
+        // Clear test mode.  This fails if test mode is not enabled.
+        PropertyInvalidatedCache.setTestMode(false);
+        try {
+            PropertyInvalidatedCache.setTestMode(false);
+            fail("expected an IllegalStateException");
+        } catch (IllegalStateException e) {
+            // The expected exception.
+        }
+        // Configuring a property for testing must fail if test mode is false.
+        TestCache cache2 = new TestCache(PropertyInvalidatedCache.MODULE_SYSTEM, "mode3");
+        try {
+            cache2.testPropertyName();
+            fail("expected an IllegalStateException");
+        } catch (IllegalStateException e) {
+            // The expected exception.
+        }
+
+        // Re-enable test mode (so that the cleanup for the test does not throw).
+        PropertyInvalidatedCache.setTestMode(true);
+    }
 }
diff --git a/core/tests/coretests/src/android/os/IpcDataCacheTest.java b/core/tests/coretests/src/android/os/IpcDataCacheTest.java
index 64f77b3..5c56fdc 100644
--- a/core/tests/coretests/src/android/os/IpcDataCacheTest.java
+++ b/core/tests/coretests/src/android/os/IpcDataCacheTest.java
@@ -17,6 +17,7 @@
 package android.os;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
 
 import android.multiuser.Flags;
 import android.platform.test.annotations.IgnoreUnderRavenwood;
@@ -26,6 +27,7 @@
 import androidx.test.filters.SmallTest;
 
 import org.junit.After;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 
@@ -92,14 +94,20 @@
         public Boolean apply(Integer x) {
             return mServer.query(x);
         }
+
         @Override
         public boolean shouldBypassCache(Integer x) {
             return x % 13 == 0;
         }
     }
 
-    // Clear the test mode after every test, in case this process is used for other
-    // tests. This also resets the test property map.
+    // Prepare for testing.
+    @Before
+    public void setUp() throws Exception {
+        IpcDataCache.setTestMode(true);
+    }
+
+    // Ensure all test configurations are cleared.
     @After
     public void tearDown() throws Exception {
         IpcDataCache.setTestMode(false);
@@ -119,9 +127,6 @@
                 new IpcDataCache<>(4, MODULE, API, "testCache1",
                         new ServerQuery(tester));
 
-        IpcDataCache.setTestMode(true);
-        testCache.testPropertyName();
-
         tester.verify(0);
         assertEquals(tester.value(3), testCache.query(3));
         tester.verify(1);
@@ -165,9 +170,6 @@
         IpcDataCache<Integer, Boolean> testCache =
                 new IpcDataCache<>(config, (x) -> tester.query(x, x % 10 == 9));
 
-        IpcDataCache.setTestMode(true);
-        testCache.testPropertyName();
-
         tester.verify(0);
         assertEquals(tester.value(3), testCache.query(3));
         tester.verify(1);
@@ -205,9 +207,6 @@
         IpcDataCache<Integer, Boolean> testCache =
                 new IpcDataCache<>(config, (x) -> tester.query(x), (x) -> x % 9 == 0);
 
-        IpcDataCache.setTestMode(true);
-        testCache.testPropertyName();
-
         tester.verify(0);
         assertEquals(tester.value(3), testCache.query(3));
         tester.verify(1);
@@ -313,8 +312,6 @@
         TestCache(String module, String api, TestQuery query) {
             super(4, module, api, "testCache7", query);
             mQuery = query;
-            setTestMode(true);
-            testPropertyName();
         }
 
         TestCache(IpcDataCache.Config c) {
@@ -324,8 +321,6 @@
         TestCache(IpcDataCache.Config c, TestQuery query) {
             super(c, query);
             mQuery = query;
-            setTestMode(true);
-            testPropertyName();
         }
 
         int getRecomputeCount() {
@@ -456,4 +451,52 @@
         TestCache ec = new TestCache(e);
         assertEquals(ec.isDisabled(), true);
     }
+
+    // Verify that test mode works properly.
+    @Test
+    public void testTestMode() {
+        // Create a cache that will write a system nonce.
+        TestCache sysCache = new TestCache(IpcDataCache.MODULE_SYSTEM, "mode1");
+        try {
+            // Invalidate the cache, which writes the system property.  There must be a permission
+            // failure.
+            sysCache.invalidateCache();
+            fail("expected permission failure");
+        } catch (RuntimeException e) {
+            // The expected exception is a bare RuntimeException.  The test does not attempt to
+            // validate the text of the exception message.
+        }
+
+        sysCache.testPropertyName();
+        // Invalidate the cache.  This must succeed because the property has been marked for
+        // testing.
+        sysCache.invalidateCache();
+
+        // Create a cache that uses MODULE_TEST.  Invalidation succeeds whether or not the
+        // property is tagged as being tested.
+        TestCache testCache = new TestCache(IpcDataCache.MODULE_TEST, "mode2");
+        testCache.invalidateCache();
+        testCache.testPropertyName();
+        testCache.invalidateCache();
+
+        // Clear test mode.  This fails if test mode is not enabled.
+        IpcDataCache.setTestMode(false);
+        try {
+            IpcDataCache.setTestMode(false);
+            fail("expected an IllegalStateException");
+        } catch (IllegalStateException e) {
+            // The expected exception.
+        }
+        // Configuring a property for testing must fail if test mode is false.
+        TestCache cache2 = new TestCache(IpcDataCache.MODULE_SYSTEM, "mode3");
+        try {
+            cache2.testPropertyName();
+            fail("expected an IllegalStateException");
+        } catch (IllegalStateException e) {
+            // The expected exception.
+        }
+
+        // Re-enable test mode (so that the cleanup for the test does not throw).
+        IpcDataCache.setTestMode(true);
+    }
 }
diff --git a/core/tests/coretests/src/android/view/LetterboxScrollProcessorTest.java b/core/tests/coretests/src/android/view/LetterboxScrollProcessorTest.java
index f8ec9f4..235625d 100644
--- a/core/tests/coretests/src/android/view/LetterboxScrollProcessorTest.java
+++ b/core/tests/coretests/src/android/view/LetterboxScrollProcessorTest.java
@@ -31,13 +31,13 @@
 import android.os.Looper;
 import android.platform.test.annotations.Presubmit;
 
+import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
 
 import org.junit.Before;
 import org.junit.Test;
 
-
 import java.util.ArrayList;
 import java.util.List;
 
@@ -54,30 +54,32 @@
     private LetterboxScrollProcessor mLetterboxScrollProcessor;
     private Context mContext;
 
-    // Constant delta used when comparing coordinates (floats)
+    // Constant delta used when comparing coordinates (floats).
     private static final float EPSILON = 0.1f;
 
+    private static final Rect APP_BOUNDS =
+            new Rect(/* left= */ 200, /* top= */ 200, /* right= */ 600, /* bottom= */ 1000);
+
     @Before
     public void setUp() {
         mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
 
         // Set app bounds as if it was letterboxed.
-        mContext.getResources().getConfiguration().windowConfiguration
-                .setBounds(new Rect(200, 200, 600, 1000));
-
-        Handler handler = new Handler(Looper.getMainLooper());
+        mContext.getResources().getConfiguration().windowConfiguration.setBounds(APP_BOUNDS);
 
         // Recreate to reset LetterboxScrollProcessor state.
-        mLetterboxScrollProcessor = new LetterboxScrollProcessor(mContext, handler);
+        mLetterboxScrollProcessor = new LetterboxScrollProcessor(mContext,
+                new Handler(Looper.getMainLooper()));
     }
 
     @Test
     public void testGestureInBoundsHasNoAdjustments() {
         // Tap-like gesture in bounds (non-scroll).
-        List<MotionEvent> tapGestureEvents = createTapGestureEvents(0f, 0f);
+        final List<MotionEvent> tapGestureEvents = createTapGestureEvents(
+                /* startX= */ 0f, /* startY= */ 0f);
 
         // Get processed events from Letterbox Scroll Processor.
-        List<MotionEvent> processedEvents = processMotionEvents(tapGestureEvents);
+        final List<MotionEvent> processedEvents = processMotionEvents(tapGestureEvents);
 
         // Ensure no changes are made to events after processing - event locations should not be
         // adjusted because the gesture started in the app's bounds (for all gestures).
@@ -87,13 +89,32 @@
     }
 
     @Test
-    public void testGestureOutsideBoundsIsIgnored() {
-        // Tap-like gesture outside bounds (non-scroll).
-        List<MotionEvent> tapGestureEvents = createTapGestureEvents(-100f, -100f);
+    public void testGestureInAppBoundsButOutsideTopWindowAlsoForwardedToTheApp() {
+        final Rect dialogBounds =
+                new Rect(/* left= */ 300, /* top= */ 500, /* right= */ 500, /* bottom= */ 700);
+        // Tap-like gesture outside the dialog, but in app bounds.
+        List<MotionEvent> tapGestureEvents = createTapGestureEventsWithCoordinateSystem(0f, 0f,
+                dialogBounds);
 
         // Get processed events from Letterbox Scroll Processor.
         List<MotionEvent> processedEvents = processMotionEvents(tapGestureEvents);
 
+        // Ensure no changes are made to events after processing - the event should be forwarded as
+        // normal.
+        assertEventLocationsAreNotAdjusted(tapGestureEvents, processedEvents);
+        // Ensure all of these events should be finished (expect no generated events).
+        assertMotionEventsShouldBeFinished(processedEvents);
+    }
+
+    @Test
+    public void testGestureOutsideBoundsIsIgnored() {
+        // Tap-like gesture outside bounds (non-scroll).
+        final List<MotionEvent> tapGestureEvents = createTapGestureEvents(
+                /* startX= */ -100f, /* startY= */ -100f);
+
+        // Get processed events from Letterbox Scroll Processor.
+        final List<MotionEvent> processedEvents = processMotionEvents(tapGestureEvents);
+
         // All events should be ignored since it was a non-scroll gesture and out of bounds.
         assertEquals(0, processedEvents.size());
     }
@@ -101,10 +122,11 @@
     @Test
     public void testScrollGestureInBoundsHasNoAdjustments() {
         // Scroll gesture in bounds (non-scroll).
-        List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(0f, 0f);
+        final List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(
+                /* startX= */ 0f, /* startY= */ 0f);
 
         // Get processed events from Letterbox Scroll Processor.
-        List<MotionEvent> processedEvents = processMotionEvents(scrollGestureEvents);
+        final List<MotionEvent> processedEvents = processMotionEvents(scrollGestureEvents);
 
         // Ensure no changes are made to events after processing - event locations should not be
         // adjusted because the gesture started in the app's bounds (for all gestures).
@@ -116,10 +138,11 @@
     @Test
     public void testScrollGestureInBoundsThenLeavesBoundsHasNoAdjustments() {
         // Scroll gesture in bounds (non-scroll) that moves out of bounds.
-        List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(390f, 790f);
+        final List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(
+                /* startX= */ 390f, /* startY= */ 790f);
 
         // Get processed events from Letterbox Scroll Processor.
-        List<MotionEvent> processedEvents = processMotionEvents(scrollGestureEvents);
+        final List<MotionEvent> processedEvents = processMotionEvents(scrollGestureEvents);
 
         // Ensure no changes are made to events after processing - event locations should not be
         // adjusted because the gesture started in the app's bounds (for all gestures), even if it
@@ -132,7 +155,8 @@
     @Test
     public void testScrollGestureOutsideBoundsIsStartedInBounds() {
         // Scroll gesture outside bounds.
-        List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(-100f, 0f);
+        List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(
+                /* startX= */ -100f, /* startY= */ 0f);
 
         // Get processed events from Letterbox Scroll Processor.
         List<MotionEvent> processedEvents = processMotionEvents(scrollGestureEvents);
@@ -142,9 +166,9 @@
 
         // Ensure offset ACTION_DOWN is first event received.
         MotionEvent firstProcessedEvent = processedEvents.getFirst();
-        assertEquals(firstProcessedEvent.getAction(), ACTION_DOWN);
-        assertEquals(firstProcessedEvent.getX(), 0, EPSILON);
-        assertEquals(firstProcessedEvent.getY(), 0, EPSILON);
+        assertEquals(ACTION_DOWN, firstProcessedEvent.getAction());
+        assertEquals(0, firstProcessedEvent.getX(), EPSILON);
+        assertEquals(0, firstProcessedEvent.getY(), EPSILON);
         // Ensure this event is not finished (because it was generated by LetterboxScrollProcessor).
         assertNull(mLetterboxScrollProcessor.processMotionEventBeforeFinish(firstProcessedEvent));
     }
@@ -152,16 +176,17 @@
     @Test
     public void testScrollGestureOutsideBoundsIsMovedInBounds() {
         // Scroll gesture outside bounds.
-        List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(-100f, 0f);
+        final List<MotionEvent> scrollGestureEvents = createScrollGestureEvents(
+                /* startX= */ -100f, /* startY= */ 0f);
 
         // Get processed events from Letterbox Scroll Processor.
-        List<MotionEvent> processedEvents = processMotionEvents(scrollGestureEvents);
+        final List<MotionEvent> processedEvents = processMotionEvents(scrollGestureEvents);
 
         // When a scroll occurs outside bounds: once detected as a scroll, an offset ACTION_DOWN is
         // placed and then the rest of the gesture is offset also. Some ACTION_MOVE events may be
         // ignored until the gesture is 'detected as a scroll'.
         // For this test, we expect the first ACTION_MOVE event to be ignored:
-        scrollGestureEvents.remove(1);
+        scrollGestureEvents.remove(/* index= */ 1);
 
         // Ensure all processed events (that are not ignored) are offset over the app.
         assertXCoordinatesAdjustedToZero(scrollGestureEvents, processedEvents);
@@ -170,8 +195,9 @@
         assertMotionEventsShouldBeFinished(processedEvents.subList(1, processedEvents.size()));
     }
 
-    private List<MotionEvent> processMotionEvents(List<MotionEvent> motionEvents) {
-        List<MotionEvent> processedEvents = new ArrayList<>();
+    @NonNull
+    private List<MotionEvent> processMotionEvents(@NonNull List<MotionEvent> motionEvents) {
+        final List<MotionEvent> processedEvents = new ArrayList<>();
         for (MotionEvent motionEvent : motionEvents) {
             MotionEvent clonedEvent = MotionEvent.obtain(motionEvent);
             List<MotionEvent> letterboxScrollCompatEvents =
@@ -187,39 +213,76 @@
         return processedEvents;
     }
 
+    /**
+     * Creates and returns a tap gesture with X and Y in reference to the app bounds (top left
+     * corner is x=0, y=0).
+     */
+    @NonNull
     private List<MotionEvent> createTapGestureEvents(float startX, float startY) {
+        return createTapGestureEventsWithCoordinateSystem(startX, startY, APP_BOUNDS);
+    }
+
+    /**
+     * @param referenceWindowBounds the amount the event will be translated by.
+     */
+    private List<MotionEvent> createTapGestureEventsWithCoordinateSystem(float startX, float startY,
+            @NonNull Rect referenceWindowBounds) {
         // Events for tap-like gesture (non-scroll)
         List<MotionEvent> motionEvents = new ArrayList<>();
-        motionEvents.add(createBasicMotionEvent(0, ACTION_DOWN, startX, startY));
-        motionEvents.add(createBasicMotionEvent(10, ACTION_UP, startX , startY));
+        motionEvents.add(createBasicMotionEventWithCoordinateSystem(0, ACTION_DOWN,
+                startX, startY, referenceWindowBounds));
+        motionEvents.add(createBasicMotionEventWithCoordinateSystem(10, ACTION_UP,
+                startX , startY, referenceWindowBounds));
         return motionEvents;
     }
 
+    @NonNull
     private List<MotionEvent> createScrollGestureEvents(float startX, float startY) {
-        float touchSlop = (float) ViewConfiguration.get(mContext).getScaledTouchSlop();
+        final float touchSlop = (float) ViewConfiguration.get(mContext).getScaledTouchSlop();
 
-        // Events for scroll gesture (starts at (startX, startY) then moves down-right
-        List<MotionEvent> motionEvents = new ArrayList<>();
-        motionEvents.add(createBasicMotionEvent(0, ACTION_DOWN, startX, startY));
-        motionEvents.add(createBasicMotionEvent(10, ACTION_MOVE,
+        // Events for scroll gesture (starts at (startX, startY) then moves down-right.
+        final List<MotionEvent> motionEvents = new ArrayList<>();
+        motionEvents.add(createBasicMotionEvent(/* downTime= */ 0, ACTION_DOWN, startX, startY));
+        motionEvents.add(createBasicMotionEvent(/* downTime= */ 10, ACTION_MOVE,
                 startX + touchSlop / 2, startY + touchSlop / 2));
-        // Below event is first event in the scroll gesture where distance > touchSlop
-        motionEvents.add(createBasicMotionEvent(20, ACTION_MOVE,
+        // Below event is first event in the scroll gesture where distance > touchSlop.
+        motionEvents.add(createBasicMotionEvent(/* downTime= */ 20, ACTION_MOVE,
                 startX + touchSlop * 2, startY + touchSlop * 2));
-        motionEvents.add(createBasicMotionEvent(30, ACTION_MOVE,
+        motionEvents.add(createBasicMotionEvent(/* downTime= */ 30, ACTION_MOVE,
                 startX + touchSlop * 3, startY + touchSlop * 3));
-        motionEvents.add(createBasicMotionEvent(40, ACTION_UP,
+        motionEvents.add(createBasicMotionEvent(/* downTime= */ 40, ACTION_UP,
                 startX + touchSlop * 3, startY + touchSlop * 3));
         return motionEvents;
     }
 
-    private MotionEvent createBasicMotionEvent(int downTime, int action, float x, float y) {
-        return MotionEvent.obtain(0, downTime, action, x, y, 0);
+    /**
+     * Creates and returns an event with X and Y in reference to the app bounds (top left corner is
+     * x=0, y=0).
+     */
+    @NonNull
+    private MotionEvent createBasicMotionEvent(int eventTime, int action, float x, float y) {
+        return createBasicMotionEventWithCoordinateSystem(eventTime, action, x, y, APP_BOUNDS);
+    }
+
+    /**
+     * @param referenceWindowBounds the amount the event will be translated by.
+     */
+    @NonNull
+    private MotionEvent createBasicMotionEventWithCoordinateSystem(int eventTime, int action,
+            float x, float y, @NonNull Rect referenceWindowBounds) {
+        final float rawX = referenceWindowBounds.left + x;
+        final float rawY = referenceWindowBounds.top + y;
+        // RawX and RawY cannot be changed once the event is created. Therefore, pass rawX and rawY
+        // according to the app's bounds on the display, and then offset to make X and Y relative to
+        // the app's bounds.
+        final MotionEvent event = MotionEvent.obtain(0, eventTime, action, rawX, rawY, 0);
+        event.offsetLocation(-referenceWindowBounds.left, -referenceWindowBounds.top);
+        return event;
     }
 
     private void assertEventLocationsAreNotAdjusted(
-            List<MotionEvent> originalEvents,
-            List<MotionEvent> processedEvents) {
+            @NonNull List<MotionEvent> originalEvents,
+            @NonNull List<MotionEvent> processedEvents) {
         assertEquals("MotionEvent arrays are not the same size",
                 originalEvents.size(), processedEvents.size());
 
@@ -232,8 +295,8 @@
     }
 
     private void assertXCoordinatesAdjustedToZero(
-            List<MotionEvent> originalEvents,
-            List<MotionEvent> processedEvents) {
+            @NonNull List<MotionEvent> originalEvents,
+            @NonNull List<MotionEvent> processedEvents) {
         assertEquals("MotionEvent arrays are not the same size",
                 originalEvents.size(), processedEvents.size());
 
@@ -245,7 +308,7 @@
         }
     }
 
-    private void assertMotionEventsShouldBeFinished(List<MotionEvent> processedEvents) {
+    private void assertMotionEventsShouldBeFinished(@NonNull List<MotionEvent> processedEvents) {
         for (MotionEvent processedEvent : processedEvents) {
             assertNotNull(mLetterboxScrollProcessor.processMotionEventBeforeFinish(processedEvent));
         }
diff --git a/core/tests/coretests/src/android/view/RoundScrollbarRendererTest.java b/core/tests/coretests/src/android/view/RoundScrollbarRendererTest.java
new file mode 100644
index 0000000..262bd5c
--- /dev/null
+++ b/core/tests/coretests/src/android/view/RoundScrollbarRendererTest.java
@@ -0,0 +1,146 @@
+/*
+ * Copyright 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
+ *
+ *      https://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;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyFloat;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.graphics.Canvas;
+import android.graphics.Paint;
+import android.graphics.Rect;
+import android.platform.test.annotations.Presubmit;
+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.view.flags.Flags;
+
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for {@link RoundScrollbarRenderer}.
+ *
+ * <p>Build/Install/Run: atest FrameworksCoreTests:android.view.RoundScrollbarRendererTest
+ */
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+@Presubmit
+public class RoundScrollbarRendererTest {
+
+    private static final int DEFAULT_VERTICAL_SCROLL_RANGE = 100;
+    private static final int DEFAULT_VERTICAL_SCROLL_EXTENT = 20;
+    private static final int DEFAULT_VERTICAL_SCROLL_OFFSET = 40;
+    private static final float DEFAULT_ALPHA = 0.5f;
+    private static final Rect BOUNDS = new Rect(0, 0, 200, 200);
+
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
+    @Mock private Canvas mCanvas;
+    @Captor private ArgumentCaptor<Paint> mPaintCaptor;
+    private RoundScrollbarRenderer mScrollbar;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        MockView view = spy(new MockView(ApplicationProvider.getApplicationContext()));
+        when(view.canScrollVertically(anyInt())).thenReturn(true);
+        when(view.computeVerticalScrollRange()).thenReturn(DEFAULT_VERTICAL_SCROLL_RANGE);
+        when(view.computeVerticalScrollExtent()).thenReturn(DEFAULT_VERTICAL_SCROLL_EXTENT);
+        when(view.computeVerticalScrollOffset()).thenReturn(DEFAULT_VERTICAL_SCROLL_OFFSET);
+        mPaintCaptor = ArgumentCaptor.forClass(Paint.class);
+
+        mScrollbar = new RoundScrollbarRenderer(view);
+    }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_USE_REFACTORED_ROUND_SCROLLBAR)
+    public void testScrollbarDrawn_legacy() {
+        mScrollbar.drawRoundScrollbars(mCanvas, DEFAULT_ALPHA, BOUNDS, /* drawToLeft= */ false);
+
+        // The arc will be drawn twice, i.e. once for track and once for thumb
+        verify(mCanvas, times(2))
+                .drawArc(any(), anyFloat(), anyFloat(), eq(false), mPaintCaptor.capture());
+
+        Paint thumbPaint = mPaintCaptor.getAllValues().getFirst();
+        assertEquals(Paint.Cap.ROUND, thumbPaint.getStrokeCap());
+        assertEquals(Paint.Style.STROKE, thumbPaint.getStyle());
+        Paint trackPaint = mPaintCaptor.getAllValues().get(1);
+        assertEquals(Paint.Cap.ROUND, trackPaint.getStrokeCap());
+        assertEquals(Paint.Style.STROKE, trackPaint.getStyle());
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_USE_REFACTORED_ROUND_SCROLLBAR)
+    public void testScrollbarDrawn() {
+        mScrollbar.drawRoundScrollbars(mCanvas, DEFAULT_ALPHA, BOUNDS, /* drawToLeft= */ false);
+
+        // The arc will be drawn thrice, i.e. twice for track and once for thumb
+        verify(mCanvas, times(3))
+                .drawArc(any(), anyFloat(), anyFloat(), eq(false), mPaintCaptor.capture());
+
+        // Verify paint styles
+        Paint thumbPaint = mPaintCaptor.getAllValues().getFirst();
+        assertEquals(Paint.Cap.ROUND, thumbPaint.getStrokeCap());
+        assertEquals(Paint.Style.STROKE, thumbPaint.getStyle());
+        Paint trackPaint = mPaintCaptor.getAllValues().get(1);
+        assertEquals(Paint.Cap.ROUND, trackPaint.getStrokeCap());
+        assertEquals(Paint.Style.STROKE, trackPaint.getStyle());
+    }
+
+    public static class MockView extends View {
+
+        public MockView(Context context) {
+            super(context);
+        }
+
+        @Override
+        public int computeVerticalScrollRange() {
+            return super.getHeight();
+        }
+
+        @Override
+        public int computeVerticalScrollOffset() {
+            return super.computeVerticalScrollOffset();
+        }
+
+        @Override
+        public int computeVerticalScrollExtent() {
+            return super.computeVerticalScrollExtent();
+        }
+    }
+}
diff --git a/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java b/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java
index c3a5b19c94..1cbc7d6 100644
--- a/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java
+++ b/core/tests/coretests/src/com/android/internal/jank/FrameTrackerTest.java
@@ -40,7 +40,10 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.animation.AnimationHandler;
+import android.os.ConditionVariable;
 import android.os.Handler;
+import android.view.Choreographer;
 import android.view.FrameMetrics;
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.JankData;
@@ -576,6 +579,44 @@
     }
 
     @Test
+    public void testEndAnimationWithLastFrameSyncId() {
+        final long[] expectedLastAnimationFrameVsyncId = { 0 };
+        final long[] lastAnimationFrameVsyncId = { 0 };
+        final long[] endAnimationVsyncId = { 0 };
+        final ConditionVariable condition = new ConditionVariable();
+        final FrameTracker tracker = spyFrameTracker(/* surfaceOnly= */ false);
+        mActivity.runOnUiThread(() -> {
+            final AnimationHandler animationHandler = AnimationHandler.getInstance();
+            final Choreographer realChoreographer = Choreographer.getInstance();
+            // 3 v-sync ids:
+            //  1. Begin mocked vsyncId = 0
+            //  2. Current real vsyncId = last animation frame
+            //  3. Posted real vsyncId = end animation
+            when(mChoreographer.getVsyncId()).thenReturn(0L);
+            tracker.begin();
+            mRunnableArgumentCaptor.getValue().run();
+            when(mChoreographer.getVsyncId()).thenAnswer(a -> realChoreographer.getVsyncId());
+            expectedLastAnimationFrameVsyncId[0] = realChoreographer.getVsyncId();
+            // Simulate ending multiple animators. Their callbacks should run in a batch.
+            animationHandler.postEndAnimationCallback(() -> {
+                endAnimationVsyncId[0] = realChoreographer.getVsyncId();
+                lastAnimationFrameVsyncId[0] =
+                        animationHandler.getLastAnimationFrameVsyncId(endAnimationVsyncId[0]);
+            });
+            animationHandler.postEndAnimationCallback(() -> {
+                tracker.end(FrameTracker.REASON_END_NORMAL);
+                condition.open();
+            });
+        });
+
+        condition.block(1000L /* timeoutMs */);
+        assertThat(lastAnimationFrameVsyncId[0]).isEqualTo(expectedLastAnimationFrameVsyncId[0]);
+        assertThat(endAnimationVsyncId[0]).isGreaterThan(lastAnimationFrameVsyncId[0]);
+        // Verifies that FrameTracker#mEndVsyncId uses the vsyncId from AnimationHandler.
+        verify(mJankStatsRegistration).removeAfter(eq(lastAnimationFrameVsyncId[0]));
+    }
+
+    @Test
     public void testMaxSuccessiveMissedFramesCount() {
         FrameTracker tracker = spyFrameTracker(/* surfaceOnly= */ true);
         when(mChoreographer.getVsyncId()).thenReturn(100L);
diff --git a/core/xsd/vts/Android.bp b/core/xsd/vts/Android.bp
index 5d8407f..239eed0 100644
--- a/core/xsd/vts/Android.bp
+++ b/core/xsd/vts/Android.bp
@@ -15,6 +15,7 @@
 //
 
 package {
+    default_team: "trendy_team_android_kernel",
     // 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"
diff --git a/graphics/java/android/framework_graphics.aconfig b/graphics/java/android/framework_graphics.aconfig
index 5ad97e6..a63cbee 100644
--- a/graphics/java/android/framework_graphics.aconfig
+++ b/graphics/java/android/framework_graphics.aconfig
@@ -33,3 +33,12 @@
   description: "Add OkLab ColorSpace support"
   bug: "344038816"
 }
+
+flag {
+  name: "display_bt2020_colorspace"
+  is_exported: true
+  is_fixed_read_only: true
+  namespace: "core_graphics"
+  description: "Add DISPLAY_BT2020 ColorSpace support"
+  bug: "344038816"
+}
diff --git a/graphics/java/android/graphics/ColorSpace.java b/graphics/java/android/graphics/ColorSpace.java
index 3752257..4c47de0 100644
--- a/graphics/java/android/graphics/ColorSpace.java
+++ b/graphics/java/android/graphics/ColorSpace.java
@@ -768,7 +768,44 @@
          * </table>
          */
         @FlaggedApi(Flags.FLAG_OK_LAB_COLORSPACE)
-        OK_LAB
+        OK_LAB,
+
+        /**
+         * <p>{@link ColorSpace.Rgb RGB} color space BT.2020 based on Rec. ITU-R BT.2020-1 and IEC 61966-2.1:1999.</p></p>
+         * <table summary="Color space definition">
+         *     <tr>
+         *         <th>Chromaticity</th><th>Red</th><th>Green</th><th>Blue</th><th>White point</th>
+         *     </tr>
+         *     <tr><td>x</td><td>0.708</td><td>0.170</td><td>0.131</td><td>0.3127</td></tr>
+         *     <tr><td>y</td><td>0.292</td><td>0.797</td><td>0.046</td><td>0.3290</td></tr>
+         *     <tr><th>Property</th><th colspan="4">Value</th></tr>
+         *     <tr><td>Name</td><td colspan="4">Rec. ITU-R BT.2020-1</td></tr>
+         *     <tr><td>CIE standard illuminant</td><td colspan="4">D65</td></tr>
+         *     <tr>
+         *         <td>Opto-electronic transfer function (OETF)</td>
+         *         <td colspan="4">\(\begin{equation}
+         *             C_{DisplayBT2020} = \begin{cases} 12.92 \times C_{linear} & C_{linear} \lt 0.0030186 \\\
+         *             1.055 \times C_{linear}^{\frac{1}{2.4}} - 0.055 & C_{linear} \ge 0.0030186 \end{cases}
+         *             \end{equation}\)
+         *         </td>
+         *     </tr>
+         *     <tr>
+         *         <td>Electro-optical transfer function (EOTF)</td>
+         *         <td colspan="4">\(\begin{equation}
+         *             C_{linear} = \begin{cases}\frac{C_{DisplayBT2020}}{12.92} & C_{sRGB} \lt 0.04045 \\\
+         *             \left( \frac{C_{DisplayBT2020} + 0.055}{1.055} \right) ^{2.4} & C_{sRGB} \ge 0.04045 \end{cases}
+         *             \end{equation}\)
+         *         </td>
+         *     </tr>
+         *     <tr><td>Range</td><td colspan="4">\([0..1]\)</td></tr>
+         * </table>
+         * <p>
+         *     <img style="display: block; margin: 0 auto;" src="{@docRoot}reference/android/images/graphics/colorspace_bt2020.png" />
+         *     <figcaption style="text-align: center;">BT.2020 (orange) vs sRGB (white)</figcaption>
+         * </p>
+         */
+        @FlaggedApi(Flags.FLAG_DISPLAY_BT2020_COLORSPACE)
+        DISPLAY_BT2020
         // Update the initialization block next to #get(Named) when adding new values
     }
 
@@ -1721,6 +1758,19 @@
                     Named.OK_LAB.ordinal()
             ));
         }
+
+        if (Flags.displayBt2020Colorspace()) {
+            sNamedColorSpaceMap.put(Named.DISPLAY_BT2020.ordinal(), new ColorSpace.Rgb(
+                    "BT 2020",
+                    BT2020_PRIMARIES,
+                    ILLUMINANT_D65,
+                    null,
+                    SRGB_TRANSFER_PARAMETERS,
+                    Named.DISPLAY_BT2020.ordinal()
+            ));
+            sDataToColorSpaces.put(DataSpace.DATASPACE_DISPLAY_BT2020,
+                    Named.DISPLAY_BT2020.ordinal());
+        }
     }
 
     private static double transferHLGOETF(Rgb.TransferParameters params, double x) {
diff --git a/graphics/java/android/graphics/text/PositionedGlyphs.java b/graphics/java/android/graphics/text/PositionedGlyphs.java
index ed17fde..792e248 100644
--- a/graphics/java/android/graphics/text/PositionedGlyphs.java
+++ b/graphics/java/android/graphics/text/PositionedGlyphs.java
@@ -24,6 +24,7 @@
 import android.graphics.Paint;
 import android.graphics.Typeface;
 import android.graphics.fonts.Font;
+import android.os.Build;
 
 import com.android.internal.util.Preconditions;
 import com.android.text.flags.Flags;
@@ -53,6 +54,8 @@
                         Typeface.class.getClassLoader(), nReleaseFunc());
     }
 
+    private static boolean sIsRobolectric = Build.FINGERPRINT.equals("robolectric");
+
     private final long mLayoutPtr;
     private final float mXOffset;
     private final float mYOffset;
@@ -252,7 +255,7 @@
         mXOffset = xOffset;
         mYOffset = yOffset;
 
-        if (Flags.typefaceRedesign()) {
+        if (!sIsRobolectric && Flags.typefaceRedesign()) {
             int fontCount = nGetFontCount(layoutPtr);
             mFonts = new ArrayList<>(fontCount);
             for (int i = 0; i < fontCount; ++i) {
diff --git a/keystore/java/android/security/OWNERS b/keystore/java/android/security/OWNERS
index ed30587..32759b2 100644
--- a/keystore/java/android/security/OWNERS
+++ b/keystore/java/android/security/OWNERS
@@ -1 +1,2 @@
 per-file *.java,*.aidl = eranm@google.com,pgrafov@google.com,rubinxu@google.com
+per-file KeyStoreManager.java = mpgroover@google.com
diff --git a/keystore/java/android/security/keystore/KeyStoreManager.java b/keystore/java/android/security/keystore/KeyStoreManager.java
new file mode 100644
index 0000000..197aaba
--- /dev/null
+++ b/keystore/java/android/security/keystore/KeyStoreManager.java
@@ -0,0 +1,327 @@
+/*
+ * 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.security.keystore;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.SystemService;
+import android.content.Context;
+import android.security.KeyStore2;
+import android.security.KeyStoreException;
+import android.security.keystore2.AndroidKeyStoreProvider;
+import android.security.keystore2.AndroidKeyStorePublicKey;
+import android.system.keystore2.Domain;
+import android.system.keystore2.KeyDescriptor;
+import android.system.keystore2.KeyPermission;
+import android.util.Log;
+
+import com.android.internal.annotations.GuardedBy;
+
+import java.io.ByteArrayInputStream;
+import java.security.Key;
+import java.security.KeyPair;
+import java.security.PublicKey;
+import java.security.UnrecoverableKeyException;
+import java.security.cert.CertificateFactory;
+import java.security.cert.X509Certificate;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * This class provides methods for interacting with keys stored within the
+ * <a href="/privacy-and-security/keystore">Android Keystore</a>.
+ */
+@FlaggedApi(android.security.Flags.FLAG_KEYSTORE_GRANT_API)
+@SystemService(Context.KEYSTORE_SERVICE)
+public class KeyStoreManager {
+    private static final String TAG = "KeyStoreManager";
+
+    private static final Object sInstanceLock = new Object();
+    @GuardedBy("sInstanceLock")
+    private static KeyStoreManager sInstance;
+
+    private final KeyStore2 mKeyStore2;
+
+    /**
+     * Private constructor to ensure only a single instance is created.
+     */
+    private KeyStoreManager() {
+        mKeyStore2 = KeyStore2.getInstance();
+    }
+
+    /**
+     * Returns the single instance of the {@code KeyStoreManager}.
+     *
+     * @hide
+     */
+    public static KeyStoreManager getInstance() {
+        synchronized (sInstanceLock) {
+            if (sInstance == null) {
+                sInstance = new KeyStoreManager();
+            }
+            return sInstance;
+        }
+    }
+
+    /**
+     * Grants access to the key owned by the calling app stored under the specified {@code alias}
+     * to another app on the device with the provided {@code uid}.
+     *
+     * <p>This method supports granting access to instances of both {@link javax.crypto.SecretKey}
+     * and {@link java.security.PrivateKey}. The resulting ID will persist across reboots and can be
+     * used by the grantee app for the life of the key or until access is revoked with {@link
+     * #revokeKeyAccess(String, int)}.
+     *
+     * <p>If the provided {@code alias} does not correspond to a key in the Android KeyStore, then
+     * an {@link UnrecoverableKeyException} is thrown.
+     *
+     * @param alias the alias of the key to be granted to another app
+     * @param uid   the uid of the app to which the key should be granted
+     * @return the ID of the granted key; this can be shared with the specified app, and that
+     * app can use {@link #getGrantedKeyFromId(long)} to access the key
+     * @throws UnrecoverableKeyException if the specified key cannot be recovered
+     * @throws KeyStoreException if an error is encountered when attempting to grant access to
+     * the key
+     * @see #getGrantedKeyFromId(long)
+     */
+    public long grantKeyAccess(@NonNull String alias, int uid)
+            throws KeyStoreException, UnrecoverableKeyException {
+        KeyDescriptor keyDescriptor = createKeyDescriptorFromAlias(alias);
+        final int grantAccessVector = KeyPermission.USE | KeyPermission.GET_INFO;
+        // When a key is in the GRANT domain, the nspace field of the KeyDescriptor contains its ID.
+        KeyDescriptor result = null;
+        try {
+            result = mKeyStore2.grant(keyDescriptor, uid, grantAccessVector);
+        } catch (KeyStoreException e) {
+            // If the provided alias does not correspond to a valid key in the KeyStore, then throw
+            // an UnrecoverableKeyException to remain consistent with other APIs in this class.
+            if (e.getNumericErrorCode() == KeyStoreException.ERROR_KEY_DOES_NOT_EXIST) {
+                throw new UnrecoverableKeyException("No key found by the given alias");
+            }
+            throw e;
+        }
+        if (result == null) {
+            Log.e(TAG, "Received a null KeyDescriptor from grant");
+            throw new KeyStoreException(KeyStoreException.ERROR_INTERNAL_SYSTEM_ERROR,
+                    "No ID was returned for the grant request for alias " + alias + " to uid "
+                            + uid);
+        } else if (result.domain != Domain.GRANT) {
+            Log.e(TAG, "Received a result outside the grant domain: " + result.domain);
+            throw new KeyStoreException(KeyStoreException.ERROR_INTERNAL_SYSTEM_ERROR,
+                    "Unable to obtain a grant ID for alias " + alias + " to uid " + uid);
+        }
+        return result.nspace;
+    }
+
+    /**
+     * Revokes access to the key in the app's namespace stored under the specified {@code
+     * alias} that was previously granted to another app on the device with the provided
+     * {@code uid}.
+     *
+     * <p>If the provided {@code alias} does not correspond to a key in the Android KeyStore, then
+     * an {@link UnrecoverableKeyException} is thrown.
+     *
+     * @param alias the alias of the key to be revoked from another app
+     * @param uid   the uid of the app from which the key access should be revoked
+     * @throws UnrecoverableKeyException if the specified key cannot be recovered
+     * @throws KeyStoreException if an error is encountered when attempting to revoke access
+     * to the key
+     */
+    public void revokeKeyAccess(@NonNull String alias, int uid)
+            throws KeyStoreException, UnrecoverableKeyException {
+        KeyDescriptor keyDescriptor = createKeyDescriptorFromAlias(alias);
+        try {
+            mKeyStore2.ungrant(keyDescriptor, uid);
+        } catch (KeyStoreException e) {
+            // If the provided alias does not correspond to a valid key in the KeyStore, then throw
+            // an UnrecoverableKeyException to remain consistent with other APIs in this class.
+            if (e.getNumericErrorCode() == KeyStoreException.ERROR_KEY_DOES_NOT_EXIST) {
+                throw new UnrecoverableKeyException("No key found by the given alias");
+            }
+            throw e;
+        }
+    }
+
+    /**
+     * Returns the key with the specified {@code id} that was previously shared with the
+     * app.
+     *
+     * <p>This method can return instances of both {@link javax.crypto.SecretKey} and {@link
+     * java.security.PrivateKey}. If a key with the provide {@code id} has not been granted to the
+     * caller, then an {@link UnrecoverableKeyException} is thrown.
+     *
+     * @param id the ID of the key that was shared with the app
+     * @return the {@link Key} that was shared with the app
+     * @throws UnrecoverableKeyException          if the specified key cannot be recovered
+     * @throws KeyPermanentlyInvalidatedException if the specified key was authorized to only
+     *                                            be used if the user has been authenticated and a
+     *                                            change has been made to the users
+     *                                            lockscreen or biometric enrollment that
+     *                                            permanently invalidates the key
+     * @see #grantKeyAccess(String, int)
+     */
+    public @NonNull Key getGrantedKeyFromId(long id)
+            throws UnrecoverableKeyException, KeyPermanentlyInvalidatedException {
+        Key result = AndroidKeyStoreProvider.loadAndroidKeyStoreKeyFromKeystore(mKeyStore2, null,
+                id, Domain.GRANT);
+        if (result == null) {
+            throw new UnrecoverableKeyException("No key found by the given alias");
+        }
+        return result;
+    }
+
+    /**
+     * Returns a {@link KeyPair} containing the public and private key associated with
+     * the key that was previously shared with the app under the provided {@code id}.
+     *
+     * <p>If a {@link java.security.PrivateKey} has not been granted to the caller with the
+     * specified {@code id}, then an {@link UnrecoverableKeyException} is thrown.
+     *
+     * @param id the ID of the private key that was shared with the app
+     * @return a KeyPair containing the public and private key shared with the app
+     * @throws UnrecoverableKeyException          if the specified key cannot be recovered
+     * @throws KeyPermanentlyInvalidatedException if the specified key was authorized to only
+     *                                            be used if the user has been authenticated and a
+     *                                            change has been made to the users
+     *                                            lockscreen or biometric enrollment that
+     *                                            permanently invalidates the key
+     */
+    public @NonNull KeyPair getGrantedKeyPairFromId(long id)
+            throws UnrecoverableKeyException, KeyPermanentlyInvalidatedException {
+        KeyDescriptor keyDescriptor = createKeyDescriptorFromId(id, Domain.GRANT);
+        return AndroidKeyStoreProvider.loadAndroidKeyStoreKeyPairFromKeystore(mKeyStore2,
+                keyDescriptor);
+    }
+
+    /**
+     * Returns a {@link List} of {@link X509Certificate} instances representing the certificate
+     * chain for the key that was previously shared with the app under the provided {@code id}.
+     *
+     * <p>If a {@link java.security.PrivateKey} has not been granted to the caller with the
+     * specified {@code id}, then an {@link UnrecoverableKeyException} is thrown.
+     *
+     * @param id the ID of the asymmetric key that was shared with the app
+     * @return a List of X509Certificates with the certificate at index 0 corresponding to
+     * the private key shared with the app
+     * @throws UnrecoverableKeyException          if the specified key cannot be recovered
+     * @throws KeyPermanentlyInvalidatedException if the specified key was authorized to only
+     *                                            be used if the user has been authenticated and a
+     *                                            change has been made to the users
+     *                                            lockscreen or biometric enrollment that
+     *                                            permanently invalidates the key
+     * @see #grantKeyAccess(String, int)
+     */
+    // Java APIs should prefer mutable collection return types with the exception being
+    // Collection.empty return types.
+    @SuppressWarnings("MixedMutabilityReturnType")
+    public @NonNull List<X509Certificate> getGrantedCertificateChainFromId(long id)
+            throws UnrecoverableKeyException, KeyPermanentlyInvalidatedException {
+        KeyDescriptor keyDescriptor = createKeyDescriptorFromId(id, Domain.GRANT);
+        KeyPair keyPair = AndroidKeyStoreProvider.loadAndroidKeyStoreKeyPairFromKeystore(mKeyStore2,
+                keyDescriptor);
+        PublicKey keyStoreKey = keyPair.getPublic();
+        if (keyStoreKey instanceof AndroidKeyStorePublicKey) {
+            AndroidKeyStorePublicKey androidKeyStorePublicKey =
+                    (AndroidKeyStorePublicKey) keyStoreKey;
+            byte[] certBytes = androidKeyStorePublicKey.getCertificate();
+            X509Certificate cert = getCertificate(certBytes);
+            // If the leaf certificate is null, then a chain should not exist either
+            if (cert == null) {
+                return Collections.emptyList();
+            }
+            List<X509Certificate> result = new ArrayList<>();
+            result.add(cert);
+            byte[] certificateChain = androidKeyStorePublicKey.getCertificateChain();
+            Collection<X509Certificate> certificates = getCertificates(certificateChain);
+            result.addAll(certificates);
+            return result;
+        } else {
+            Log.e(TAG, "keyStoreKey is not of the expected type: " + keyStoreKey);
+        }
+        return Collections.emptyList();
+    }
+
+    /**
+     * Returns an {@link X509Certificate} instance from the provided {@code certificate} byte
+     * encoding of the certificate, or null if the provided encoding is null.
+     */
+    private static X509Certificate getCertificate(byte[] certificate) {
+        X509Certificate result = null;
+        if (certificate != null) {
+            try {
+                CertificateFactory certificateFactory = CertificateFactory.getInstance("X.509");
+                result = (X509Certificate) certificateFactory.generateCertificate(
+                        new ByteArrayInputStream(certificate));
+            } catch (Exception e) {
+                Log.e(TAG, "Caught an exception parsing the certificate: ", e);
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Returns a {@link Collection} of {@link X509Certificate} instances from the provided
+     * {@code certificateChain} byte encoding of the certificates, or null if the provided
+     * encoding is null.
+     */
+    private static Collection<X509Certificate> getCertificates(byte[] certificateChain) {
+        if (certificateChain != null) {
+            try {
+                CertificateFactory certificateFactory = CertificateFactory.getInstance("X.509");
+                Collection<X509Certificate> certificates =
+                        (Collection<X509Certificate>) certificateFactory.generateCertificates(
+                                new ByteArrayInputStream(certificateChain));
+                if (certificates == null) {
+                    Log.e(TAG, "Received null certificates from a non-null certificateChain");
+                    return Collections.emptyList();
+                }
+                return certificates;
+            } catch (Exception e) {
+                Log.e(TAG, "Caught an exception parsing the certs: ", e);
+            }
+        }
+        return Collections.emptyList();
+    }
+
+    /**
+     * Returns a new {@link KeyDescriptor} instance in the app domain / namespace with the {@code
+     * alias} set to the provided value.
+     */
+    private static KeyDescriptor createKeyDescriptorFromAlias(String alias) {
+        KeyDescriptor keyDescriptor = new KeyDescriptor();
+        keyDescriptor.domain = Domain.APP;
+        keyDescriptor.nspace = KeyProperties.NAMESPACE_APPLICATION;
+        keyDescriptor.alias = alias;
+        keyDescriptor.blob = null;
+        return keyDescriptor;
+    }
+
+    /**
+     * Returns a new {@link KeyDescriptor} instance in the provided {@code domain} with the nspace
+     * field set to the provided {@code id}.
+     */
+    private static KeyDescriptor createKeyDescriptorFromId(long id, int domain) {
+        KeyDescriptor keyDescriptor = new KeyDescriptor();
+        keyDescriptor.domain = domain;
+        keyDescriptor.nspace = id;
+        keyDescriptor.alias = null;
+        keyDescriptor.blob = null;
+        return keyDescriptor;
+    }
+}
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStoreProvider.java b/keystore/java/android/security/keystore2/AndroidKeyStoreProvider.java
index 99100de..dcc8844 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStoreProvider.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStoreProvider.java
@@ -17,6 +17,7 @@
 package android.security.keystore2;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.security.KeyStore2;
 import android.security.KeyStoreSecurityLevel;
 import android.security.keymaster.KeymasterDefs;
@@ -335,11 +336,11 @@
     }
 
     /**
-     * Loads an an AndroidKeyStoreKey from the AndroidKeyStore backend.
+     * Loads an AndroidKeyStoreKey from the AndroidKeyStore backend.
      *
      * @param keyStore The keystore2 backend.
      * @param alias The alias of the key in the Keystore database.
-     * @param namespace The a Keystore namespace. This is used by system api only to request
+     * @param namespace The Keystore namespace. This is used by system api only to request
      *         Android system specific keystore namespace, which can be configured
      *         in the device's SEPolicy. Third party apps and most system components
      *         set this parameter to -1 to indicate their application specific namespace.
@@ -351,14 +352,40 @@
     public static AndroidKeyStoreKey loadAndroidKeyStoreKeyFromKeystore(
             @NonNull KeyStore2 keyStore, @NonNull String alias, int namespace)
             throws UnrecoverableKeyException, KeyPermanentlyInvalidatedException {
-        KeyDescriptor descriptor = new KeyDescriptor();
+        int descriptorNamespace;
+        int descriptorDomain;
         if (namespace == KeyProperties.NAMESPACE_APPLICATION) {
-            descriptor.nspace = KeyProperties.NAMESPACE_APPLICATION; // ignored;
-            descriptor.domain = Domain.APP;
+            descriptorNamespace = KeyProperties.NAMESPACE_APPLICATION; // ignored;
+            descriptorDomain = Domain.APP;
         } else {
-            descriptor.nspace = namespace;
-            descriptor.domain = Domain.SELINUX;
+            descriptorNamespace = namespace;
+            descriptorDomain = Domain.SELINUX;
         }
+        return loadAndroidKeyStoreKeyFromKeystore(keyStore, alias, descriptorNamespace,
+                descriptorDomain);
+    }
+
+    /**
+     * Loads an AndroidKeyStoreKey from the AndroidKeyStore backend.
+     *
+     * @param keyStore The keystore2 backend
+     * @param alias The alias of the key in the Keystore database
+     * @param namespace The Keystore namespace. This is used by system api only to request
+     *         Android system specific keystore namespace, which can be configured
+     *         in the device's SEPolicy. Third party apps and most system components
+     *         set this parameter to -1 to indicate their application specific namespace.
+     *         See <a href="https://source.android.com/security/keystore#access-control">
+     *             Keystore 2.0 access control</a>
+     * @param domain The Keystore domain
+     * @return an AndroidKeyStoreKey corresponding to the provided values for the KeyDescriptor
+     * @hide
+     */
+    public static AndroidKeyStoreKey loadAndroidKeyStoreKeyFromKeystore(@NonNull KeyStore2 keyStore,
+            @Nullable String alias, long namespace, int domain)
+            throws UnrecoverableKeyException, KeyPermanentlyInvalidatedException {
+        KeyDescriptor descriptor = new KeyDescriptor();
+        descriptor.nspace = namespace;
+        descriptor.domain = domain;
         descriptor.alias = alias;
         descriptor.blob = null;
 
diff --git a/keystore/java/android/security/keystore2/AndroidKeyStorePublicKey.java b/keystore/java/android/security/keystore2/AndroidKeyStorePublicKey.java
index 0b3be32..bcf619b 100644
--- a/keystore/java/android/security/keystore2/AndroidKeyStorePublicKey.java
+++ b/keystore/java/android/security/keystore2/AndroidKeyStorePublicKey.java
@@ -44,6 +44,22 @@
         mEncoded = x509EncodedForm;
     }
 
+    /**
+     * Returns the byte array encoding of the certificate corresponding to this public key.
+     * @hide
+     */
+    public byte[] getCertificate() {
+        return mCertificate;
+    }
+
+    /**
+     * Returns the byte array encoding of the certificate chain for this public key.
+     * @hide
+     */
+    public byte[] getCertificateChain() {
+        return mCertificateChain;
+    }
+
     abstract AndroidKeyStorePrivateKey getPrivateKey();
 
     @Override
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/WindowExtensionsImpl.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/WindowExtensionsImpl.java
index 7f11fea..2ab0310 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/WindowExtensionsImpl.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/WindowExtensionsImpl.java
@@ -57,9 +57,9 @@
      */
     private static final int NO_LEVEL_OVERRIDE = -1;
 
-    private static final int EXTENSIONS_VERSION_V7 = 7;
+    private static final int EXTENSIONS_VERSION_V8 = 8;
 
-    private static final int EXTENSIONS_VERSION_V6 = 6;
+    private static final int EXTENSIONS_VERSION_V7 = 7;
 
     private final Object mLock = new Object();
     private volatile DeviceStateManagerFoldingFeatureProducer mFoldingFeatureProducer;
@@ -80,12 +80,10 @@
      */
     @VisibleForTesting
     static int getExtensionsVersionCurrentPlatform() {
-        if (Flags.activityEmbeddingAnimationCustomizationFlag()) {
-            // Activity Embedding animation customization is the only major feature for v7.
-            return EXTENSIONS_VERSION_V7;
-        } else {
-            return EXTENSIONS_VERSION_V6;
+        if (Flags.aeBackStackRestore()) {
+            return EXTENSIONS_VERSION_V8;
         }
+        return EXTENSIONS_VERSION_V7;
     }
 
     private String generateLogMessage() {
diff --git a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml
index aeb734e..5609663 100644
--- a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml
+++ b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml
@@ -184,6 +184,7 @@
             android:layout_height="20dp"
             android:layout_gravity="end|center_vertical"
             android:layout_marginEnd="16dp"
+            android:layout_marginStart="10dp"
             android:contentDescription="@string/open_by_default_settings_text"
             android:src="@drawable/desktop_mode_ic_handle_menu_open_by_default_settings"
             android:tint="?androidprv:attr/materialColorOnSurface"/>
diff --git a/libs/WindowManager/Shell/res/values/styles.xml b/libs/WindowManager/Shell/res/values/styles.xml
index d061ae1..55cda78 100644
--- a/libs/WindowManager/Shell/res/values/styles.xml
+++ b/libs/WindowManager/Shell/res/values/styles.xml
@@ -43,7 +43,8 @@
         <item name="android:layout_width">match_parent</item>
         <item name="android:layout_height">52dp</item>
         <item name="android:gravity">start|center_vertical</item>
-        <item name="android:padding">16dp</item>
+        <item name="android:paddingStart">16dp</item>
+        <item name="android:paddingEnd">0dp</item>
         <item name="android:textSize">14sp</item>
         <item name="android:textFontWeight">500</item>
         <item name="android:textColor">?androidprv:attr/materialColorOnSurface</item>
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
index 19b51f1..e4db7b6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
@@ -164,6 +164,9 @@
      */
     private BackTouchTracker mQueuedTracker = new BackTouchTracker();
 
+    private final BackTransitionObserver mBackTransitionObserver =
+            new BackTransitionObserver();
+
     private final Runnable mAnimationTimeoutRunnable = () -> {
         ProtoLog.w(WM_SHELL_BACK_PREVIEW, "Animation didn't finish in %d ms. Resetting...",
                 MAX_ANIMATION_DURATION);
@@ -268,6 +271,8 @@
         mBackTransitionHandler = new BackTransitionHandler();
         mTransitions.addHandler(mBackTransitionHandler);
         mHandler = handler;
+        mTransitions.registerObserver(mBackTransitionObserver);
+        mBackTransitionObserver.setBackTransitionHandler(mBackTransitionHandler);
         updateTouchableArea();
     }
 
@@ -729,6 +734,13 @@
     }
 
     /**
+     * @return Latest task id which back gesture has occurred on it.
+     */
+    public int getLatestTriggerBackTask() {
+        return mBackTransitionObserver.mFocusedTaskId;
+    }
+
+    /**
      * Sets to true when the back gesture has passed the triggering threshold, false otherwise.
      */
     public void setTriggerBack(boolean triggerBack) {
@@ -792,6 +804,11 @@
         boolean triggerBack = activeTouchTracker.getTriggerBack();
         ProtoLog.d(WM_SHELL_BACK_PREVIEW, "onGestureFinished() mTriggerBack == %s", triggerBack);
 
+        if (triggerBack) {
+            mBackTransitionObserver.update(mBackNavigationInfo != null
+                            ? mBackNavigationInfo.getFocusedTaskId()
+                            : INVALID_TASK_ID);
+        }
         // Reset gesture states.
         mThresholdCrossed = false;
         mPointersPilfered = false;
@@ -1218,6 +1235,7 @@
             }
 
             if (shouldCancelAnimation(info)) {
+                mPrepareOpenTransition = null;
                 return false;
             }
 
@@ -1645,4 +1663,58 @@
     private static boolean canBeTransitionTarget(TransitionInfo.Change change) {
         return findComponentName(change) != null || findTaskId(change) != INVALID_TASK_ID;
     }
+
+    // Record the latest back gesture happen on which task.
+    static class BackTransitionObserver implements Transitions.TransitionObserver {
+        int mFocusedTaskId = INVALID_TASK_ID;
+        IBinder mFocusTaskMonitorToken;
+        private BackTransitionHandler mBackTransitionHandler;
+        void setBackTransitionHandler(BackTransitionHandler handler) {
+            mBackTransitionHandler = handler;
+        }
+
+        void update(int focusedTaskId) {
+            mFocusedTaskId = focusedTaskId;
+        }
+
+        @Override
+        public void onTransitionReady(@NonNull IBinder transition, @NonNull TransitionInfo info,
+                @NonNull SurfaceControl.Transaction startTransaction,
+                @NonNull SurfaceControl.Transaction finishTransaction) {
+            if (mFocusedTaskId == INVALID_TASK_ID) {
+                return;
+            }
+            for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+                final TransitionInfo.Change c = info.getChanges().get(i);
+                if (c.getTaskInfo() != null && c.getTaskInfo().taskId == mFocusedTaskId) {
+                    mFocusTaskMonitorToken = transition;
+                    break;
+                }
+            }
+            // Transition happen but the task isn't involved, reset.
+            if (mFocusTaskMonitorToken == null) {
+                mFocusedTaskId = INVALID_TASK_ID;
+            }
+        }
+
+        @Override
+        public void onTransitionMerged(@NonNull IBinder merged, @NonNull IBinder playing) {
+            if (mFocusTaskMonitorToken == merged) {
+                mFocusTaskMonitorToken = playing;
+            }
+            if (mBackTransitionHandler.mClosePrepareTransition == merged) {
+                mBackTransitionHandler.mClosePrepareTransition = null;
+            }
+        }
+
+        @Override
+        public void onTransitionFinished(@NonNull IBinder transition, boolean aborted) {
+            if (mFocusTaskMonitorToken == transition) {
+                mFocusedTaskId = INVALID_TASK_ID;
+            }
+            if (mBackTransitionHandler.mClosePrepareTransition == transition) {
+                mBackTransitionHandler.mClosePrepareTransition = null;
+            }
+        }
+    }
 }
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 537ef18..810eff8 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
@@ -776,6 +776,10 @@
                 cancelPhysicsAnimation();
                 settlePipBoundsAfterPhysicsAnimation(false /* animatingAfter */);
                 break;
+            case PipTransitionState.CHANGED_PIP_BOUNDS:
+                // Check whether changed bounds imply we need to update stash state too.
+                stashEndActionIfNeeded();
+                break;
         }
     }
 
@@ -829,9 +833,6 @@
         mPipBoundsState.getMotionBoundsState().onPhysicsAnimationEnded();
         mSpringingToTouch = false;
         mDismissalPending = false;
-
-        // Check whether new bounds after fling imply we need to update stash state too.
-        stashEndActionIfNeeded();
     }
 
     private void stashEndActionIfNeeded() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java
index 4d0432e..19d293e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java
@@ -1013,21 +1013,6 @@
             return true;
         }
 
-        private void stashEndAction() {
-            if (mPipBoundsState.getBounds().left < 0
-                    && mPipBoundsState.getStashedState() != STASH_TYPE_LEFT) {
-                mPipUiEventLogger.log(
-                        PipUiEventLogger.PipUiEventEnum.PICTURE_IN_PICTURE_STASH_LEFT);
-                mPipBoundsState.setStashed(STASH_TYPE_LEFT);
-            } else if (mPipBoundsState.getBounds().left >= 0
-                    && mPipBoundsState.getStashedState() != STASH_TYPE_RIGHT) {
-                mPipUiEventLogger.log(
-                        PipUiEventLogger.PipUiEventEnum.PICTURE_IN_PICTURE_STASH_RIGHT);
-                mPipBoundsState.setStashed(STASH_TYPE_RIGHT);
-            }
-            mMenuController.hideMenu();
-        }
-
         private void flingEndAction() {
             if (mShouldHideMenuAfterFling) {
                 // If the menu is not visible, then we can still be showing the activity for the
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHandleManageWindowsMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHandleManageWindowsMenu.kt
index 13a805a..e71b4f3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHandleManageWindowsMenu.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHandleManageWindowsMenu.kt
@@ -95,15 +95,15 @@
     override fun addToContainer(menuView: ManageWindowsView) {
         val menuPosition = calculateMenuPosition()
         menuViewContainer = AdditionalSystemViewContainer(
-            windowManagerWrapper,
-            callerTaskInfo.taskId,
-            menuPosition.x,
-            menuPosition.y,
-            menuView.menuWidth,
-            menuView.menuHeight,
-            WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE or
+            windowManagerWrapper = windowManagerWrapper,
+            taskId = callerTaskInfo.taskId,
+            x = menuPosition.x,
+            y = menuPosition.y,
+            width = menuView.menuWidth,
+            height = menuView.menuHeight,
+            flags = WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE or
                     WindowManager.LayoutParams.FLAG_WATCH_OUTSIDE_TOUCH,
-            menuView.rootView
+            view = menuView.rootView,
         )
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenu.kt
index 05391a8..173bc08 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenu.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenu.kt
@@ -22,14 +22,19 @@
 import android.graphics.Point
 import android.view.SurfaceControl
 import android.view.SurfaceControlViewHost
+import android.view.WindowInsets.Type.systemBars
 import android.view.WindowManager
 import android.view.WindowlessWindowManager
 import android.window.TaskConstants
 import android.window.TaskSnapshot
 import androidx.compose.ui.graphics.toArgb
+import com.android.internal.annotations.VisibleForTesting
+import com.android.window.flags.Flags
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.common.DisplayController
+import com.android.wm.shell.desktopmode.DesktopRepository
 import com.android.wm.shell.shared.desktopmode.ManageWindowsViewContainer
+import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalSystemViewContainer
 import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalViewContainer
 import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalViewHostViewContainer
 import com.android.wm.shell.windowdecor.common.DecorThemeUtil
@@ -41,9 +46,12 @@
  */
 class DesktopHeaderManageWindowsMenu(
     private val callerTaskInfo: RunningTaskInfo,
+    private val x: Int,
+    private val y: Int,
     private val displayController: DisplayController,
     private val rootTdaOrganizer: RootTaskDisplayAreaOrganizer,
     context: Context,
+    private val desktopRepository: DesktopRepository,
     private val surfaceControlBuilderSupplier: Supplier<SurfaceControl.Builder>,
     private val surfaceControlTransactionSupplier: Supplier<SurfaceControl.Transaction>,
     snapshotList: List<Pair<Int, TaskSnapshot>>,
@@ -53,7 +61,8 @@
     context,
     DecorThemeUtil(context).getColorScheme(callerTaskInfo).background.toArgb()
 ) {
-    private var menuViewContainer: AdditionalViewContainer? = null
+    @VisibleForTesting
+    var menuViewContainer: AdditionalViewContainer? = null
 
     init {
         show(snapshotList, onIconClickListener, onOutsideClickListener)
@@ -64,8 +73,37 @@
     }
 
     override fun addToContainer(menuView: ManageWindowsView) {
-        val taskBounds = callerTaskInfo.getConfiguration().windowConfiguration.bounds
-        val menuPosition = Point(taskBounds.left, taskBounds.top)
+        val menuPosition = Point(x, y)
+        val flags = WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE or
+                WindowManager.LayoutParams.FLAG_WATCH_OUTSIDE_TOUCH or
+                WindowManager.LayoutParams.FLAG_SPLIT_TOUCH
+        menuViewContainer = if (Flags.enableFullyImmersiveInDesktop()
+            && desktopRepository.isTaskInFullImmersiveState(callerTaskInfo.taskId)) {
+            // Use system view container so that forcibly shown system bars take effect in
+            // immersive.
+            createAsSystemViewContainer(menuPosition, flags)
+        } else {
+            createAsViewHostContainer(menuPosition, flags)
+        }
+    }
+
+    private fun createAsSystemViewContainer(position: Point, flags: Int): AdditionalViewContainer {
+        return AdditionalSystemViewContainer(
+            windowManagerWrapper = WindowManagerWrapper(
+                context.getSystemService(WindowManager::class.java)
+            ),
+            taskId = callerTaskInfo.taskId,
+            x = position.x,
+            y = position.y,
+            width = menuView.menuWidth,
+            height = menuView.menuHeight,
+            flags = flags,
+            forciblyShownTypes = systemBars(),
+            view = menuView.rootView
+        )
+    }
+
+    private fun createAsViewHostContainer(position: Point, flags: Int): AdditionalViewContainer {
         val builder = surfaceControlBuilderSupplier.get()
         rootTdaOrganizer.attachToDisplayArea(callerTaskInfo.displayId, builder)
         val leash = builder
@@ -73,11 +111,10 @@
             .setContainerLayer()
             .build()
         val lp = WindowManager.LayoutParams(
-            menuView.menuWidth, menuView.menuHeight,
+            menuView.menuWidth,
+            menuView.menuHeight,
             WindowManager.LayoutParams.TYPE_APPLICATION,
-            WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE
-                    or WindowManager.LayoutParams.FLAG_WATCH_OUTSIDE_TOUCH
-                    or WindowManager.LayoutParams.FLAG_SPLIT_TOUCH,
+            flags,
             PixelFormat.TRANSPARENT
         )
         val windowManager = WindowlessWindowManager(
@@ -93,11 +130,12 @@
         menuView.let { viewHost.setView(it.rootView, lp) }
         val t = surfaceControlTransactionSupplier.get()
         t.setLayer(leash, TaskConstants.TASK_CHILD_LAYER_FLOATING_MENU)
-            .setPosition(leash, menuPosition.x.toFloat(), menuPosition.y.toFloat())
+            .setPosition(leash, position.x.toFloat(), position.y.toFloat())
             .show(leash)
         t.apply()
-        menuViewContainer = AdditionalViewHostViewContainer(
-            leash, viewHost,
+        return AdditionalViewHostViewContainer(
+            leash,
+            viewHost,
             surfaceControlTransactionSupplier
         )
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index d3b7ca1..6eb20b9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -454,7 +454,13 @@
         }
 
         if (isHandleMenuActive()) {
-            mHandleMenu.relayout(startT, mResult.mCaptionX);
+            mHandleMenu.relayout(
+                    startT,
+                    mResult.mCaptionX,
+                    // Add top padding to the caption Y so that the menu is shown over what is the
+                    // actual contents of the caption, ignoring padding. This is currently relevant
+                    // to the Header in desktop immersive.
+                    mResult.mCaptionY + mResult.mCaptionTopPadding);
         }
 
         if (isOpenByDefaultDialogActive()) {
@@ -1258,6 +1264,8 @@
                 && Flags.enableDesktopWindowingMultiInstanceFeatures();
         final boolean shouldShowManageWindowsButton = supportsMultiInstance
                 && mMinimumInstancesFound;
+        final boolean inDesktopImmersive = mDesktopRepository
+                .isTaskInFullImmersiveState(mTaskInfo.taskId);
         mHandleMenu = mHandleMenuFactory.create(
                 this,
                 mWindowManagerWrapper,
@@ -1271,7 +1279,11 @@
                 getBrowserLink(),
                 mResult.mCaptionWidth,
                 mResult.mCaptionHeight,
-                mResult.mCaptionX
+                mResult.mCaptionX,
+                // Add top padding to the caption Y so that the menu is shown over what is the
+                // actual contents of the caption, ignoring padding. This is currently relevant
+                // to the Header in desktop immersive.
+                mResult.mCaptionY + mResult.mCaptionTopPadding
         );
         mWindowDecorViewHolder.onHandleMenuOpened();
         mHandleMenu.show(
@@ -1302,7 +1314,8 @@
                 /* onOutsideTouchListener= */ () -> {
                     closeHandleMenu();
                     return Unit.INSTANCE;
-                }
+                },
+                /* forceShowSystemBars= */ inDesktopImmersive
         );
         if (canEnterDesktopMode(mContext) && Flags.enableDesktopWindowingAppHandleEducation()) {
             notifyCaptionStateChanged();
@@ -1316,9 +1329,12 @@
         if (mTaskInfo.isFreeform()) {
             mManageWindowsMenu = new DesktopHeaderManageWindowsMenu(
                     mTaskInfo,
+                    /* x= */ mResult.mCaptionX,
+                    /* y= */ mResult.mCaptionY + mResult.mCaptionTopPadding,
                     mDisplayController,
                     mRootTaskDisplayAreaOrganizer,
                     mContext,
+                    mDesktopRepository,
                     mSurfaceControlBuilderSupplier,
                     mSurfaceControlTransactionSupplier,
                     snapshotList,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt
index 2e32703..93bd929 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt
@@ -32,6 +32,7 @@
 import android.view.MotionEvent.ACTION_OUTSIDE
 import android.view.SurfaceControl
 import android.view.View
+import android.view.WindowInsets.Type.systemBars
 import android.view.WindowManager
 import android.widget.Button
 import android.widget.ImageButton
@@ -73,7 +74,8 @@
     private val openInBrowserIntent: Intent?,
     private val captionWidth: Int,
     private val captionHeight: Int,
-    captionX: Int
+    captionX: Int,
+    captionY: Int
 ) {
     private val context: Context = parentDecor.mDecorWindowContext
     private val taskInfo: RunningTaskInfo = parentDecor.mTaskInfo
@@ -110,7 +112,7 @@
         get() = openInBrowserIntent != null
 
     init {
-        updateHandleMenuPillPositions(captionX)
+        updateHandleMenuPillPositions(captionX, captionY)
     }
 
     fun show(
@@ -123,6 +125,7 @@
         onOpenByDefaultClickListener: () -> Unit,
         onCloseMenuClickListener: () -> Unit,
         onOutsideTouchListener: () -> Unit,
+        forceShowSystemBars: Boolean = false,
     ) {
         val ssg = SurfaceSyncGroup(TAG)
         val t = SurfaceControl.Transaction()
@@ -139,6 +142,7 @@
             onOpenByDefaultClickListener = onOpenByDefaultClickListener,
             onCloseMenuClickListener = onCloseMenuClickListener,
             onOutsideTouchListener = onOutsideTouchListener,
+            forceShowSystemBars = forceShowSystemBars,
         )
         ssg.addTransaction(t)
         ssg.markSyncReady()
@@ -157,7 +161,8 @@
         openInBrowserClickListener: (Intent) -> Unit,
         onOpenByDefaultClickListener: () -> Unit,
         onCloseMenuClickListener: () -> Unit,
-        onOutsideTouchListener: () -> Unit
+        onOutsideTouchListener: () -> Unit,
+        forceShowSystemBars: Boolean = false,
     ) {
         val handleMenuView = HandleMenuView(
             context = context,
@@ -185,7 +190,7 @@
         val x = handleMenuPosition.x.toInt()
         val y = handleMenuPosition.y.toInt()
         handleMenuViewContainer =
-            if (!taskInfo.isFreeform && Flags.enableHandleInputFix()) {
+            if ((!taskInfo.isFreeform && Flags.enableHandleInputFix()) || forceShowSystemBars) {
                 AdditionalSystemViewContainer(
                     windowManagerWrapper = windowManagerWrapper,
                     taskId = taskInfo.taskId,
@@ -196,7 +201,8 @@
                     flags = WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE or
                             WindowManager.LayoutParams.FLAG_WATCH_OUTSIDE_TOUCH or
                             WindowManager.LayoutParams.FLAG_SPLIT_TOUCH,
-                    view = handleMenuView.rootView
+                    view = handleMenuView.rootView,
+                    forciblyShownTypes = if (forceShowSystemBars) { systemBars() } else { 0 }
                 )
             } else {
                 parentDecor.addWindow(
@@ -210,15 +216,15 @@
     /**
      * Updates handle menu's position variables to reflect its next position.
      */
-    private fun updateHandleMenuPillPositions(captionX: Int) {
+    private fun updateHandleMenuPillPositions(captionX: Int, captionY: Int) {
         val menuX: Int
         val menuY: Int
         val taskBounds = taskInfo.getConfiguration().windowConfiguration.bounds
-        updateGlobalMenuPosition(taskBounds, captionX)
+        updateGlobalMenuPosition(taskBounds, captionX, captionY)
         if (layoutResId == R.layout.desktop_mode_app_header) {
             // Align the handle menu to the left side of the caption.
             menuX = marginMenuStart
-            menuY = marginMenuTop
+            menuY = captionY + marginMenuTop
         } else {
             if (Flags.enableHandleInputFix()) {
                 // In a focused decor, we use global coordinates for handle menu. Therefore we
@@ -228,26 +234,26 @@
                 menuY = globalMenuPosition.y
             } else {
                 menuX = (taskBounds.width() / 2) - (menuWidth / 2)
-                menuY = marginMenuTop
+                menuY = captionY + marginMenuTop
             }
         }
         // Handle Menu position setup.
         handleMenuPosition.set(menuX.toFloat(), menuY.toFloat())
     }
 
-    private fun updateGlobalMenuPosition(taskBounds: Rect, captionX: Int) {
+    private fun updateGlobalMenuPosition(taskBounds: Rect, captionX: Int, captionY: Int) {
         val nonFreeformX = captionX + (captionWidth / 2) - (menuWidth / 2)
         when {
             taskInfo.isFreeform -> {
                 globalMenuPosition.set(
                     /* x = */ taskBounds.left + marginMenuStart,
-                    /* y = */ taskBounds.top + marginMenuTop
+                    /* y = */ taskBounds.top + captionY + marginMenuTop
                 )
             }
             taskInfo.isFullscreen -> {
                 globalMenuPosition.set(
                     /* x = */ nonFreeformX,
-                    /* y = */ marginMenuTop
+                    /* y = */ marginMenuTop + captionY
                 )
             }
             taskInfo.isMultiWindow -> {
@@ -261,13 +267,13 @@
                     SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT -> {
                         globalMenuPosition.set(
                             /* x = */ leftOrTopStageBounds.width() + nonFreeformX,
-                            /* y = */ marginMenuTop
+                            /* y = */ captionY + marginMenuTop
                         )
                     }
                     SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT -> {
                         globalMenuPosition.set(
                             /* x = */ nonFreeformX,
-                            /* y = */ marginMenuTop
+                            /* y = */ captionY + marginMenuTop
                         )
                     }
                 }
@@ -280,10 +286,11 @@
      */
     fun relayout(
         t: SurfaceControl.Transaction,
-        captionX: Int
+        captionX: Int,
+        captionY: Int,
     ) {
         handleMenuViewContainer?.let { container ->
-            updateHandleMenuPillPositions(captionX)
+            updateHandleMenuPillPositions(captionX, captionY)
             container.setPosition(t, handleMenuPosition.x, handleMenuPosition.y)
         }
     }
@@ -675,7 +682,8 @@
         openInBrowserIntent: Intent?,
         captionWidth: Int,
         captionHeight: Int,
-        captionX: Int
+        captionX: Int,
+        captionY: Int,
     ): HandleMenu
 }
 
@@ -694,7 +702,8 @@
         openInBrowserIntent: Intent?,
         captionWidth: Int,
         captionHeight: Int,
-        captionX: Int
+        captionX: Int,
+        captionY: Int,
     ): HandleMenu {
         return HandleMenu(
             parentDecor,
@@ -709,7 +718,8 @@
             openInBrowserIntent,
             captionWidth,
             captionHeight,
-            captionX
+            captionX,
+            captionY,
         )
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
index ce5cfd0..6b3b357 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/WindowDecoration.java
@@ -256,6 +256,8 @@
         outResult.mCaptionWidth = params.mCaptionWidthId != Resources.ID_NULL
                 ? loadDimensionPixelSize(resources, params.mCaptionWidthId) : taskBounds.width();
         outResult.mCaptionX = (outResult.mWidth - outResult.mCaptionWidth) / 2;
+        outResult.mCaptionY = 0;
+        outResult.mCaptionTopPadding = params.mCaptionTopPadding;
 
         updateDecorationContainerSurface(startT, outResult);
         updateCaptionContainerSurface(startT, outResult);
@@ -786,6 +788,8 @@
         int mCaptionHeight;
         int mCaptionWidth;
         int mCaptionX;
+        int mCaptionY;
+        int mCaptionTopPadding;
         final Region mCustomizableCaptionRegion = Region.obtain();
         int mWidth;
         int mHeight;
@@ -797,6 +801,8 @@
             mCaptionHeight = 0;
             mCaptionWidth = 0;
             mCaptionX = 0;
+            mCaptionY = 0;
+            mCaptionTopPadding = 0;
             mCustomizableCaptionRegion.setEmpty();
             mRootView = null;
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/additionalviewcontainer/AdditionalSystemViewContainer.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/additionalviewcontainer/AdditionalSystemViewContainer.kt
index 1be26f0..8b6aaaf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/additionalviewcontainer/AdditionalSystemViewContainer.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/additionalviewcontainer/AdditionalSystemViewContainer.kt
@@ -23,6 +23,7 @@
 import android.view.LayoutInflater
 import android.view.SurfaceControl
 import android.view.View
+import android.view.WindowInsets
 import android.view.WindowManager
 import com.android.wm.shell.windowdecor.WindowManagerWrapper
 
@@ -38,6 +39,7 @@
     width: Int,
     height: Int,
     flags: Int,
+    @WindowInsets.Type.InsetsType forciblyShownTypes: Int = 0,
     override val view: View
 ) : AdditionalViewContainer() {
     val lp: WindowManager.LayoutParams = WindowManager.LayoutParams(
@@ -49,6 +51,7 @@
         title = "Additional view container of Task=$taskId"
         gravity = Gravity.LEFT or Gravity.TOP
         setTrustedOverlay()
+        this.forciblyShownTypes = forciblyShownTypes
     }
 
     constructor(
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenuTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenuTest.kt
new file mode 100644
index 0000000..f9f760e
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenuTest.kt
@@ -0,0 +1,111 @@
+/*
+ * 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.windowdecor
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import android.view.SurfaceControl
+import androidx.test.filters.SmallTest
+import com.android.window.flags.Flags
+import com.android.wm.shell.MockToken
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestRunningTaskInfoBuilder
+import com.android.wm.shell.TestShellExecutor
+import com.android.wm.shell.desktopmode.DesktopRepository
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalSystemViewContainer
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+
+/**
+ * Tests for [DesktopHeaderManageWindowsMenu].
+ *
+ * Build/Install/Run:
+ * atest WMShellUnitTests:DesktopHeaderManageWindowsMenuTest
+ */
+@SmallTest
+@TestableLooper.RunWithLooper
+@RunWith(AndroidTestingRunner::class)
+class DesktopHeaderManageWindowsMenuTest : ShellTestCase() {
+
+    @JvmField
+    @Rule
+    val setFlagsRule: SetFlagsRule = SetFlagsRule()
+
+    private lateinit var desktopRepository: DesktopRepository
+    private lateinit var menu: DesktopHeaderManageWindowsMenu
+
+    @Before
+    fun setUp() {
+        desktopRepository = DesktopRepository(
+            context = context,
+            shellInit = ShellInit(TestShellExecutor()),
+            persistentRepository = mock(),
+            mainCoroutineScope = mock()
+        )
+    }
+
+    @After
+    fun tearDown() {
+        menu.close()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP)
+    fun testShow_forImmersiveTask_usesSystemViewContainer() {
+        val task = createFreeformTask()
+        desktopRepository.setTaskInFullImmersiveState(
+            displayId = task.displayId,
+            taskId = task.taskId,
+            immersive = true
+        )
+
+        menu = createMenu(task)
+
+        assertThat(menu.menuViewContainer).isInstanceOf(AdditionalSystemViewContainer::class.java)
+    }
+
+    private fun createMenu(task: RunningTaskInfo) = DesktopHeaderManageWindowsMenu(
+        callerTaskInfo = task,
+        x = 0,
+        y = 0,
+        displayController = mock(),
+        rootTdaOrganizer = mock(),
+        context = context,
+        desktopRepository = desktopRepository,
+        surfaceControlBuilderSupplier = { SurfaceControl.Builder() },
+        surfaceControlTransactionSupplier = { SurfaceControl.Transaction() },
+        snapshotList = emptyList(),
+        onIconClickListener = {},
+        onOutsideClickListener = {},
+    )
+
+    private fun createFreeformTask(): RunningTaskInfo = TestRunningTaskInfoBuilder()
+        .setToken(MockToken().token())
+        .setActivityType(ACTIVITY_TYPE_STANDARD)
+        .setWindowingMode(WINDOWING_MODE_FREEFORM)
+        .build()
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
index 1d11d2e..3208872 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
@@ -259,7 +259,8 @@
         doReturn(defaultDisplay).when(mMockDisplayController).getDisplay(Display.DEFAULT_DISPLAY);
         doReturn(mInsetsState).when(mMockDisplayController).getInsetsState(anyInt());
         when(mMockHandleMenuFactory.create(any(), any(), anyInt(), any(), any(), any(),
-                anyBoolean(), anyBoolean(), anyBoolean(), any(), anyInt(), anyInt(), anyInt()))
+                anyBoolean(), anyBoolean(), anyBoolean(), any(), anyInt(), anyInt(), anyInt(),
+                anyInt()))
                 .thenReturn(mMockHandleMenu);
         when(mMockMultiInstanceHelper.supportsMultiInstanceSplit(any())).thenReturn(false);
         when(mMockAppHeaderViewHolderFactory.create(any(), any(), any(), any(), any(), any(), any(),
@@ -1070,7 +1071,8 @@
                 openInBrowserCaptor.capture(),
                 any(),
                 any(),
-                any()
+                any(),
+                anyBoolean()
         );
         openInBrowserCaptor.getValue().invoke(new Intent(Intent.ACTION_MAIN, TEST_URI1));
 
@@ -1099,7 +1101,8 @@
                 openInBrowserCaptor.capture(),
                 any(),
                 any(),
-                any()
+                any(),
+                anyBoolean()
         );
 
         openInBrowserCaptor.getValue().invoke(new Intent(Intent.ACTION_MAIN, TEST_URI1));
@@ -1151,7 +1154,8 @@
                 any(),
                 any(),
                 closeClickListener.capture(),
-                any()
+                any(),
+                anyBoolean()
         );
 
         closeClickListener.getValue().invoke();
@@ -1161,6 +1165,30 @@
     }
 
     @Test
+    public void createHandleMenu_immersiveWindow_forceShowsSystemBars() {
+        final ActivityManager.RunningTaskInfo taskInfo = createTaskInfo(/* visible= */ true);
+        final DesktopModeWindowDecoration decoration = createWindowDecoration(taskInfo,
+                true /* relayout */);
+        when(mMockDesktopRepository.isTaskInFullImmersiveState(taskInfo.taskId))
+                .thenReturn(true);
+
+        createHandleMenu(decoration);
+
+        verify(mMockHandleMenu).show(
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                /* forceShowSystemBars= */ eq(true)
+        );
+    }
+
+    @Test
     @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_APP_HANDLE_EDUCATION)
     public void notifyCaptionStateChanged_flagDisabled_doNoNotify() {
         when(DesktopModeStatus.canEnterDesktopMode(mContext)).thenReturn(true);
@@ -1301,7 +1329,7 @@
         verify(mMockHandleMenuFactory).create(any(), any(), anyInt(), any(), any(),
                 any(), anyBoolean(), anyBoolean(), anyBoolean(),
                 argThat(intent -> (uri == null && intent == null) || intent.getData().equals(uri)),
-                anyInt(), anyInt(), anyInt());
+                anyInt(), anyInt(), anyInt(), anyInt());
     }
 
     private void createMaximizeMenu(DesktopModeWindowDecoration decoration, MaximizeMenu menu) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt
index 1820133..9544fa8 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt
@@ -33,6 +33,7 @@
 import android.view.SurfaceControl
 import android.view.SurfaceControlViewHost
 import android.view.View
+import android.view.WindowInsets.Type.systemBars
 import android.view.WindowManager
 import androidx.core.graphics.toPointF
 import androidx.test.filters.SmallTest
@@ -186,13 +187,35 @@
         assertEquals(expected.toPointF(), handleMenu.handleMenuPosition)
     }
 
-    private fun createTaskInfo(windowingMode: Int, splitPosition: Int) {
+    @Test
+    fun testCreate_forceShowSystemBars_usesSystemViewContainer() {
+        createTaskInfo(WINDOWING_MODE_FREEFORM)
+
+        handleMenu = createAndShowHandleMenu(forceShowSystemBars = true)
+
+        // Only AdditionalSystemViewContainer supports force showing system bars.
+        assertTrue(handleMenu.handleMenuViewContainer is AdditionalSystemViewContainer)
+    }
+
+    @Test
+    fun testCreate_forceShowSystemBars() {
+        createTaskInfo(WINDOWING_MODE_FREEFORM)
+
+        handleMenu = createAndShowHandleMenu(forceShowSystemBars = true)
+
+        val types = (handleMenu.handleMenuViewContainer as AdditionalSystemViewContainer)
+            .lp.forciblyShownTypes
+        assertTrue((types and systemBars()) != 0)
+    }
+
+    private fun createTaskInfo(windowingMode: Int, splitPosition: Int? = null) {
         val taskDescriptionBuilder = ActivityManager.TaskDescription.Builder()
             .setBackgroundColor(Color.YELLOW)
         val bounds = when (windowingMode) {
             WINDOWING_MODE_FULLSCREEN -> DISPLAY_BOUNDS
             WINDOWING_MODE_FREEFORM -> FREEFORM_BOUNDS
             WINDOWING_MODE_MULTI_WINDOW -> {
+                checkNotNull(splitPosition)
                 if (splitPosition == SPLIT_POSITION_TOP_OR_LEFT) {
                     SPLIT_LEFT_BOUNDS
                 } else {
@@ -208,14 +231,19 @@
             .setBounds(bounds)
             .setVisible(true)
             .build()
-        whenever(splitScreenController.getSplitPosition(any())).thenReturn(splitPosition)
-        whenever(splitScreenController.getStageBounds(any(), any())).thenAnswer {
-            (it.arguments.first() as Rect).set(SPLIT_LEFT_BOUNDS)
-            (it.arguments[1] as Rect).set(SPLIT_RIGHT_BOUNDS)
+        if (windowingMode == WINDOWING_MODE_MULTI_WINDOW) {
+            whenever(splitScreenController.getSplitPosition(any())).thenReturn(splitPosition)
+            whenever(splitScreenController.getStageBounds(any(), any())).thenAnswer {
+                (it.arguments.first() as Rect).set(SPLIT_LEFT_BOUNDS)
+                (it.arguments[1] as Rect).set(SPLIT_RIGHT_BOUNDS)
+            }
         }
     }
 
-    private fun createAndShowHandleMenu(splitPosition: Int): HandleMenu {
+    private fun createAndShowHandleMenu(
+        splitPosition: Int? = null,
+        forceShowSystemBars: Boolean = false,
+    ): HandleMenu {
         val layoutId = if (mockDesktopWindowDecoration.mTaskInfo.isFreeform) {
             R.layout.desktop_mode_app_header
         } else {
@@ -225,6 +253,7 @@
             WINDOWING_MODE_FULLSCREEN -> (DISPLAY_BOUNDS.width() / 2) - (HANDLE_WIDTH / 2)
             WINDOWING_MODE_FREEFORM -> 0
             WINDOWING_MODE_MULTI_WINDOW -> {
+                checkNotNull(splitPosition)
                 if (splitPosition == SPLIT_POSITION_TOP_OR_LEFT) {
                     (SPLIT_LEFT_BOUNDS.width() / 2) - (HANDLE_WIDTH / 2)
                 } else {
@@ -238,9 +267,21 @@
             layoutId, appIcon, appName, splitScreenController, shouldShowWindowingPill = true,
             shouldShowNewWindowButton = true, shouldShowManageWindowsButton = false,
             null /* openInBrowserLink */, captionWidth = HANDLE_WIDTH, captionHeight = 50,
-            captionX = captionX
+            captionX = captionX,
+            captionY = 0,
         )
-        handleMenu.show(mock(), mock(), mock(), mock(), mock(), mock(), mock(), mock(), mock())
+        handleMenu.show(
+            onToDesktopClickListener = mock(),
+            onToFullscreenClickListener = mock(),
+            onToSplitScreenClickListener = mock(),
+            onNewWindowClickListener = mock(),
+            onManageWindowsClickListener = mock(),
+            openInBrowserClickListener = mock(),
+            onOpenByDefaultClickListener = mock(),
+            onCloseMenuClickListener = mock(),
+            onOutsideTouchListener = mock(),
+            forceShowSystemBars = forceShowSystemBars
+        )
         return handleMenu
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
index bb41e9c..fb17ae9 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/WindowDecorationTests.java
@@ -515,6 +515,23 @@
     }
 
     @Test
+    public void testRelayout_withPadding_setsOnResult() {
+        final ActivityManager.RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
+                .setDisplayId(Display.DEFAULT_DISPLAY)
+                .setBounds(TASK_BOUNDS)
+                .setPositionInParent(TASK_POSITION_IN_PARENT.x, TASK_POSITION_IN_PARENT.y)
+                .setVisible(true)
+                .build();
+        final TestWindowDecoration windowDecor = createWindowDecoration(taskInfo);
+        mRelayoutParams.mCaptionTopPadding = 50;
+
+        windowDecor.relayout(taskInfo, false /* applyStartTransactionOnDraw */,
+                true /* hasGlobalFocus */);
+
+        assertEquals(50, mRelayoutResult.mCaptionTopPadding);
+    }
+
+    @Test
     public void testRelayout_fluidResizeEnabled_freeformTask_setTaskSurfaceColor() {
         StaticMockitoSession mockitoSession = mockitoSession().mockStatic(
                 DesktopModeStatus.class).strictness(
diff --git a/libs/hwui/utils/Color.cpp b/libs/hwui/utils/Color.cpp
index 9673c5f..23097f6 100644
--- a/libs/hwui/utils/Color.cpp
+++ b/libs/hwui/utils/Color.cpp
@@ -17,6 +17,7 @@
 #include "Color.h"
 
 #include <Properties.h>
+#include <aidl/android/hardware/graphics/common/Dataspace.h>
 #include <android/hardware_buffer.h>
 #include <android/native_window.h>
 #include <ui/ColorSpace.h>
@@ -25,6 +26,8 @@
 #include <algorithm>
 #include <cmath>
 
+#include "SkColorSpace.h"
+
 namespace android {
 namespace uirenderer {
 
@@ -215,9 +218,13 @@
         return HAL_DATASPACE_ADOBE_RGB;
     }
 
-    if (nearlyEqual(fn, SkNamedTransferFn::kRec2020) &&
-        nearlyEqual(gamut, SkNamedGamut::kRec2020)) {
-        return HAL_DATASPACE_BT2020;
+    if (nearlyEqual(gamut, SkNamedGamut::kRec2020)) {
+        if (nearlyEqual(fn, SkNamedTransferFn::kRec2020)) {
+            return HAL_DATASPACE_BT2020;
+        } else if (nearlyEqual(fn, SkNamedTransferFn::kSRGB)) {
+            return static_cast<android_dataspace>(
+                    ::aidl::android::hardware::graphics::common::Dataspace::DISPLAY_BT2020);
+        }
     }
 
     if (nearlyEqual(fn, k2Dot6) && nearlyEqual(gamut, kDCIP3)) {
diff --git a/media/java/android/media/IAudioService.aidl b/media/java/android/media/IAudioService.aidl
index 02ca307..c22b674 100644
--- a/media/java/android/media/IAudioService.aidl
+++ b/media/java/android/media/IAudioService.aidl
@@ -471,6 +471,7 @@
 
     List<AudioDeviceAttributes> getDevicesForAttributesUnprotected(in AudioAttributes attributes);
 
+    @EnforcePermission(anyOf = {"MODIFY_AUDIO_ROUTING", "QUERY_AUDIO_STATE"})
     void addOnDevicesForAttributesChangedListener(in AudioAttributes attributes,
             in IDevicesForAttributesCallback callback);
 
diff --git a/media/java/android/media/flags/projection.aconfig b/media/java/android/media/flags/projection.aconfig
index 9a9a073..17d1ff6 100644
--- a/media/java/android/media/flags/projection.aconfig
+++ b/media/java/android/media/flags/projection.aconfig
@@ -10,3 +10,11 @@
     bug: "323008518"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "media_projection_connected_display"
+    namespace: "virtual_devices"
+    description: "Enable recording connected display"
+    bug: "362720120"
+    is_exported: true
+}
diff --git a/packages/SettingsLib/Preference/testutils/com/android/settingslib/preference/CatalystScreenTestCase.kt b/packages/SettingsLib/Preference/testutils/com/android/settingslib/preference/CatalystScreenTestCase.kt
index 06214eb..8ef4c58 100644
--- a/packages/SettingsLib/Preference/testutils/com/android/settingslib/preference/CatalystScreenTestCase.kt
+++ b/packages/SettingsLib/Preference/testutils/com/android/settingslib/preference/CatalystScreenTestCase.kt
@@ -99,15 +99,18 @@
         @Suppress("UNCHECKED_CAST")
         val clazz = preferenceScreenCreator.fragmentClass() as Class<PreferenceFragmentCompat>
         val builder = StringBuilder()
-        launchFragmentScenario(clazz).use {
-            it.onFragment { fragment ->
-                taskFinished.set(true)
-                fragment.preferenceScreen.toString(builder)
-            }
+        launchFragment(clazz) { fragment ->
+            taskFinished.set(true)
+            fragment.preferenceScreen.toString(builder)
         }
         return builder.toString()
     }
 
+    protected open fun launchFragment(
+        fragmentClass: Class<PreferenceFragmentCompat>,
+        action: (PreferenceFragmentCompat) -> Unit,
+    ): Unit = launchFragmentScenario(fragmentClass).use { it.onFragment(action) }
+
     protected open fun launchFragmentScenario(fragmentClass: Class<PreferenceFragmentCompat>) =
         FragmentScenario.launch(fragmentClass)
 
diff --git a/packages/SystemUI/AndroidManifest.xml b/packages/SystemUI/AndroidManifest.xml
index 9726ecf..510c9b7 100644
--- a/packages/SystemUI/AndroidManifest.xml
+++ b/packages/SystemUI/AndroidManifest.xml
@@ -1139,11 +1139,5 @@
                 android:name="android.service.dream"
                 android:resource="@xml/home_controls_dream_metadata" />
         </service>
-
-        <activity android:name="com.android.systemui.keyboard.shortcut.ui.view.ShortcutHelperActivity"
-            android:exported="false"
-            android:showForAllUsers="true"
-            android:theme="@style/ShortcutHelperTheme"
-            android:excludeFromRecents="true" />
     </application>
 </manifest>
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index bd015d0..426b24d 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -1550,3 +1550,10 @@
        purpose: PURPOSE_BUGFIX
    }
 }
+
+flag {
+    name: "qs_tile_detailed_view"
+    namespace: "systemui"
+    description: "Enables the tile detailed view UI."
+    bug: "374173773"
+}
diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/RemoteAnimationRunnerCompat.java b/packages/SystemUI/animation/src/com/android/systemui/animation/RemoteAnimationRunnerCompat.java
index 94f8846..0b15d23 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/RemoteAnimationRunnerCompat.java
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/RemoteAnimationRunnerCompat.java
@@ -22,8 +22,14 @@
 import static android.view.WindowManager.TRANSIT_OPEN;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_EXIT_TRANSITIONS;
 import static android.window.TransitionInfo.FLAG_IS_WALLPAPER;
 
+import static com.android.internal.util.Preconditions.checkArgument;
+import static com.android.wm.shell.shared.TransitionUtil.isClosingMode;
+import static com.android.wm.shell.shared.TransitionUtil.isClosingType;
+import static com.android.wm.shell.shared.TransitionUtil.isOpeningMode;
+
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.util.ArrayMap;
@@ -157,6 +163,9 @@
                         t.show(wallpapers[i].leash);
                         t.setAlpha(wallpapers[i].leash, 1.f);
                     }
+                    if (ENABLE_DESKTOP_WINDOWING_EXIT_TRANSITIONS.isTrue()) {
+                        resetLauncherAlphaOnDesktopExit(info, launcherTask, leashMap, t);
+                    }
                 } else {
                     if (launcherTask != null) {
                         counterLauncher.addChild(t, leashMap.get(launcherTask.getLeash()));
@@ -236,4 +245,33 @@
             }
         };
     }
+
+    /**
+     * Reset the alpha of the Launcher leash to give the Launcher time to hide its Views before the
+     * exit-desktop animation starts.
+     *
+     * This method should only be called if the current transition is opening Launcher, otherwise we
+     * might not be exiting Desktop Mode.
+     */
+    private static void resetLauncherAlphaOnDesktopExit(
+            TransitionInfo info,
+            TransitionInfo.Change launcherChange,
+            ArrayMap<SurfaceControl, SurfaceControl> leashMap,
+            SurfaceControl.Transaction startTransaction
+    ) {
+        checkArgument(isOpeningMode(launcherChange.getMode()));
+        if (!isClosingType(info.getType())) {
+            return;
+        }
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            final TransitionInfo.Change change = info.getChanges().get(i);
+            // skip changes that we didn't wrap
+            if (!leashMap.containsKey(change.getLeash())) continue;
+            // Only make the update if we are closing Desktop tasks.
+            if (change.getTaskInfo().isFreeform() && isClosingMode(change.getMode())) {
+                startTransaction.setAlpha(leashMap.get(launcherChange.getLeash()), 0f);
+                return;
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/animation/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShader.kt b/packages/SystemUI/animation/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShader.kt
index 025c8b9..f426aa5 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShader.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShader.kt
@@ -70,6 +70,19 @@
             }
         """
 
+        private const val SIMPLEX_SIMPLE_SHADER =
+            """
+            vec4 main(vec2 p) {
+                vec2 uv = p / in_size.xy;
+                uv.x *= in_aspectRatio;
+
+                // Compute turbulence effect with the uv distorted with simplex noise.
+                vec3 noisePos = vec3(uv + in_noiseMove.xy, in_noiseMove.z) * in_gridNum;
+                float mixFactor = simplex3d(noisePos) * 0.5 + 0.5;
+                return mix(in_color, in_screenColor, mixFactor);
+            }
+        """
+
         private const val FRACTAL_SHADER =
             """
             vec4 main(vec2 p) {
@@ -155,6 +168,8 @@
                 return sparkleLayer;
             }
         """
+        private const val SIMPLEX_NOISE_SIMPLE_SHADER =
+            ShaderUtilLibrary.SHADER_LIB + UNIFORMS + SIMPLEX_SIMPLE_SHADER
         private const val SIMPLEX_NOISE_SHADER =
             ShaderUtilLibrary.SHADER_LIB + UNIFORMS + COMMON_FUNCTIONS + SIMPLEX_SHADER
         private const val FRACTAL_NOISE_SHADER =
@@ -163,17 +178,20 @@
             ShaderUtilLibrary.SHADER_LIB + UNIFORMS + COMMON_FUNCTIONS + SIMPLEX_SPARKLE_SHADER
 
         enum class Type {
-            /** Effect with a simple color noise turbulence. */
+            /** Effect with a color noise turbulence with luma matte. */
             SIMPLEX_NOISE,
+            /** Effect with a noise interpolating between foreground and background colors. */
+            SIMPLEX_NOISE_SIMPLE,
             /** Effect with a simple color noise turbulence, with fractal. */
             SIMPLEX_NOISE_FRACTAL,
             /** Effect with color & sparkle turbulence with screen color layer. */
-            SIMPLEX_NOISE_SPARKLE
+            SIMPLEX_NOISE_SPARKLE,
         }
 
         fun getShader(type: Type): String {
             return when (type) {
                 Type.SIMPLEX_NOISE -> SIMPLEX_NOISE_SHADER
+                Type.SIMPLEX_NOISE_SIMPLE -> SIMPLEX_NOISE_SIMPLE_SHADER
                 Type.SIMPLEX_NOISE_FRACTAL -> FRACTAL_NOISE_SHADER
                 Type.SIMPLEX_NOISE_SPARKLE -> SPARKLE_NOISE_SHADER
             }
@@ -206,15 +224,15 @@
         setFloatUniform("in_pixelDensity", pixelDensity)
     }
 
-    /** Sets the noise color of the effect. Alpha is ignored. */
+    /**
+     * Sets the noise color of the effect. Alpha is ignored for all types except
+     * [Type.SIMPLEX_NOISE_SIMPLE].
+     */
     fun setColor(color: Int) {
         setColorUniform("in_color", color)
     }
 
-    /**
-     * Sets the color that is used for blending on top of the background color/image. Only relevant
-     * to [Type.SIMPLEX_NOISE_SPARKLE].
-     */
+    /** Sets the color that is used for blending on top of the background color/image. */
     fun setScreenColor(color: Int) {
         setColorUniform("in_screenColor", color)
     }
@@ -259,7 +277,7 @@
      */
     fun setLumaMatteFactors(
         lumaMatteBlendFactor: Float = 1f,
-        lumaMatteOverallBrightness: Float = 0f
+        lumaMatteOverallBrightness: Float = 0f,
     ) {
         setFloatUniform("in_lumaMatteBlendFactor", lumaMatteBlendFactor)
         setFloatUniform("in_lumaMatteOverallBrightness", lumaMatteOverallBrightness)
@@ -279,8 +297,10 @@
     /** Current noise movements in x, y, and z axes. */
     var noiseOffsetX: Float = 0f
         private set
+
     var noiseOffsetY: Float = 0f
         private set
+
     var noiseOffsetZ: Float = 0f
         private set
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
index c1c3b1f..d08df26 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
@@ -371,10 +371,7 @@
                 .motionTestValues { animatedAlpha(animatedOffset) exportAs MotionTestValues.alpha }
         }
 
-        UserSwitcher(
-            viewModel = viewModel,
-            modifier = Modifier.weight(1f).swappable().testTag("UserSwitcher"),
-        )
+        UserSwitcher(viewModel = viewModel, modifier = Modifier.weight(1f).swappable())
 
         FoldAware(
             modifier = Modifier.weight(1f).swappable(inversed = true).testTag("FoldAware"),
@@ -738,7 +735,7 @@
     Column(
         horizontalAlignment = Alignment.CenterHorizontally,
         verticalArrangement = Arrangement.Center,
-        modifier = modifier,
+        modifier = modifier.sysuiResTag("UserSwitcher"),
     ) {
         selectedUserImage?.let {
             Image(
@@ -781,7 +778,7 @@
                     Icon(
                         imageVector = Icons.Default.KeyboardArrowDown,
                         contentDescription = null,
-                        modifier = Modifier.size(32.dp),
+                        modifier = Modifier.size(32.dp).sysuiResTag("user_switcher_anchor"),
                     )
                 }
 
@@ -819,11 +816,11 @@
             expanded = isExpanded,
             onDismissRequest = onDismissed,
             offset = DpOffset(x = 0.dp, y = -UserSwitcherDropdownHeight),
-            modifier =
-                Modifier.width(UserSwitcherDropdownWidth).sysuiResTag("user_switcher_dropdown"),
+            modifier = Modifier.width(UserSwitcherDropdownWidth).sysuiResTag("user_list_dropdown"),
         ) {
             items.forEach { userSwitcherDropdownItem ->
                 DropdownMenuItem(
+                    modifier = Modifier.sysuiResTag("user_switcher_item"),
                     leadingIcon = {
                         Icon(
                             icon = userSwitcherDropdownItem.icon,
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 4ab5261..8b0daf6 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
@@ -16,6 +16,7 @@
 
 package com.android.systemui.communal.ui.compose
 
+import android.appwidget.AppWidgetProviderInfo
 import android.content.Context
 import android.content.res.Configuration
 import android.graphics.drawable.Icon
@@ -24,7 +25,6 @@
 import android.view.MotionEvent
 import android.widget.FrameLayout
 import android.widget.RemoteViews
-import androidx.annotation.VisibleForTesting
 import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.core.LinearEasing
 import androidx.compose.animation.core.Spring
@@ -149,9 +149,11 @@
 import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.DpSize
+import androidx.compose.ui.unit.IntRect
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.LayoutDirection
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.round
 import androidx.compose.ui.unit.sp
 import androidx.compose.ui.unit.times
 import androidx.compose.ui.util.fastAll
@@ -183,6 +185,9 @@
 import com.android.systemui.communal.widgets.WidgetConfigurator
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.phone.SystemUIDialogFactory
+import kotlin.math.max
+import kotlin.math.min
+import kotlin.math.roundToInt
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
 
@@ -386,13 +391,19 @@
                             contentOffset = contentOffset,
                             screenWidth = screenWidth,
                             setGridCoordinates = { gridCoordinates = it },
-                            updateDragPositionForRemove = { offset ->
-                                isPointerWithinEnabledRemoveButton(
-                                    removeEnabled = removeButtonEnabled,
-                                    offset =
-                                        gridCoordinates?.let { it.positionInWindow() + offset },
-                                    containerToCheck = removeButtonCoordinates,
-                                )
+                            updateDragPositionForRemove = { boundingBox ->
+                                val gridOffset = gridCoordinates?.positionInWindow()
+                                val removeButtonCenter =
+                                    removeButtonCoordinates?.boundsInWindow()?.center
+                                removeButtonEnabled &&
+                                    gridOffset != null &&
+                                    removeButtonCenter != null &&
+                                    boundingBox
+                                        // The bounding box is relative to the grid, so we need to
+                                        // normalize it by adding the grid offset and the content
+                                        // offset.
+                                        .translate((gridOffset + contentOffset).round())
+                                        .contains(removeButtonCenter.round())
                             },
                             gridState = gridState,
                             contentListState = contentListState,
@@ -644,11 +655,13 @@
 @Composable
 private fun ResizableItemFrameWrapper(
     key: String,
+    currentSpan: GridItemSpan,
     gridState: LazyGridState,
-    minItemSpan: Int,
     gridContentPadding: PaddingValues,
     verticalArrangement: Arrangement.Vertical,
     enabled: Boolean,
+    minHeightPx: Int,
+    maxHeightPx: Int,
     modifier: Modifier = Modifier,
     alpha: () -> Float = { 1f },
     onResize: (info: ResizeInfo) -> Unit = {},
@@ -659,20 +672,48 @@
     } else {
         ResizableItemFrame(
             key = key,
+            currentSpan = currentSpan,
             gridState = gridState,
-            minItemSpan = minItemSpan,
             gridContentPadding = gridContentPadding,
             verticalArrangement = verticalArrangement,
             enabled = enabled,
             alpha = alpha,
             modifier = modifier,
             onResize = onResize,
+            minHeightPx = minHeightPx,
+            maxHeightPx = maxHeightPx,
+            resizeMultiple = CommunalContentSize.HALF.span,
         ) {
             content(Modifier)
         }
     }
 }
 
+@Composable
+fun calculateWidgetSize(item: CommunalContentModel, isResizable: Boolean): WidgetSizeInfo {
+    val density = LocalDensity.current
+
+    return if (isResizable && item is CommunalContentModel.WidgetContent.Widget) {
+        with(density) {
+            val minHeightPx =
+                (min(item.providerInfo.minResizeHeight, item.providerInfo.minHeight)
+                    .coerceAtLeast(CommunalContentSize.HALF.dp().toPx().roundToInt()))
+
+            val maxHeightPx =
+                (if (item.providerInfo.maxResizeHeight > 0) {
+                        max(item.providerInfo.maxResizeHeight, item.providerInfo.minHeight)
+                    } else {
+                        Int.MAX_VALUE
+                    })
+                    .coerceIn(minHeightPx, CommunalContentSize.FULL.dp().toPx().roundToInt())
+
+            WidgetSizeInfo(minHeightPx, maxHeightPx)
+        }
+    } else {
+        WidgetSizeInfo(0, Int.MAX_VALUE)
+    }
+}
+
 @OptIn(ExperimentalFoundationApi::class)
 @Composable
 private fun BoxScope.CommunalHubLazyGrid(
@@ -685,7 +726,7 @@
     gridState: LazyGridState,
     contentListState: ContentListState,
     setGridCoordinates: (coordinates: LayoutCoordinates) -> Unit,
-    updateDragPositionForRemove: (offset: Offset) -> Boolean,
+    updateDragPositionForRemove: (boundingBox: IntRect) -> Boolean,
     widgetConfigurator: WidgetConfigurator?,
     interactionHandler: RemoteViews.InteractionHandler?,
     widgetSection: CommunalAppWidgetSection,
@@ -747,6 +788,12 @@
             val size = SizeF(Dimensions.CardWidth.value, item.size.dp().value)
             val selected = item.key == selectedKey.value
             val dpSize = DpSize(size.width.dp, size.height.dp)
+            val isResizable =
+                if (item is CommunalContentModel.WidgetContent.Widget) {
+                    item.providerInfo.resizeMode and AppWidgetProviderInfo.RESIZE_VERTICAL != 0
+                } else {
+                    false
+                }
 
             if (viewModel.isEditMode && dragDropState != null) {
                 val outlineAlpha by
@@ -755,10 +802,11 @@
                         animationSpec = spring(stiffness = Spring.StiffnessMediumLow),
                         label = "Widget resizing outline alpha",
                     )
+                val widgetSizeInfo = calculateWidgetSize(item, isResizable)
                 ResizableItemFrameWrapper(
                     key = item.key,
+                    currentSpan = GridItemSpan(item.size.span),
                     gridState = gridState,
-                    minItemSpan = CommunalContentSize.HALF.span,
                     gridContentPadding = contentPadding,
                     verticalArrangement = itemArrangement,
                     enabled = selected,
@@ -772,6 +820,8 @@
                             )
                         },
                     onResize = { resizeInfo -> contentListState.resize(index, resizeInfo) },
+                    minHeightPx = widgetSizeInfo.minHeightPx,
+                    maxHeightPx = widgetSizeInfo.maxHeightPx,
                 ) { modifier ->
                     DraggableItem(
                         modifier = modifier,
@@ -781,7 +831,7 @@
                         key = item.key,
                     ) { isDragging ->
                         CommunalContent(
-                            modifier = Modifier.fillMaxSize(),
+                            modifier = Modifier.requiredSize(dpSize),
                             model = item,
                             viewModel = viewModel,
                             size = size,
@@ -1556,23 +1606,6 @@
     }
 }
 
-/**
- * Check whether the pointer position that the item is being dragged at is within the coordinates of
- * the remove button in the toolbar. Returns true if the item is removable.
- */
-@VisibleForTesting
-fun isPointerWithinEnabledRemoveButton(
-    removeEnabled: Boolean,
-    offset: Offset?,
-    containerToCheck: LayoutCoordinates?,
-): Boolean {
-    if (!removeEnabled || offset == null || containerToCheck == null) {
-        return false
-    }
-    val container = containerToCheck.boundsInWindow()
-    return container.contains(offset)
-}
-
 private fun CommunalContentSize.dp(): Dp {
     return when (this) {
         CommunalContentSize.FULL -> Dimensions.CardHeightFull
@@ -1656,6 +1689,8 @@
     }
 }
 
+data class WidgetSizeInfo(val minHeightPx: Int, val maxHeightPx: Int)
+
 private object Colors {
     val DisabledColorFilter by lazy { disabledColorMatrix() }
 
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 0718bc3..5feb63d 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
@@ -58,11 +58,11 @@
 fun rememberGridDragDropState(
     gridState: LazyGridState,
     contentListState: ContentListState,
-    updateDragPositionForRemove: (offset: Offset) -> Boolean,
+    updateDragPositionForRemove: (boundingBox: IntRect) -> Boolean,
 ): GridDragDropState {
     val scope = rememberCoroutineScope()
     val state =
-        remember(gridState, contentListState) {
+        remember(gridState, contentListState, updateDragPositionForRemove) {
             GridDragDropState(
                 state = gridState,
                 contentListState = contentListState,
@@ -92,7 +92,7 @@
     private val state: LazyGridState,
     private val contentListState: ContentListState,
     private val scope: CoroutineScope,
-    private val updateDragPositionForRemove: (offset: Offset) -> Boolean,
+    private val updateDragPositionForRemove: (draggingBoundingBox: IntRect) -> Boolean,
 ) {
     var draggingItemKey by mutableStateOf<Any?>(null)
         private set
@@ -104,7 +104,6 @@
 
     private var draggingItemDraggedDelta by mutableStateOf(Offset.Zero)
     private var draggingItemInitialOffset by mutableStateOf(Offset.Zero)
-    private var dragStartPointerOffset by mutableStateOf(Offset.Zero)
 
     private var previousTargetItemKey: Any? = null
 
@@ -139,7 +138,6 @@
             // before content padding from the initial pointer position
             .firstItemAtOffset(normalizedOffset - contentOffset)
             ?.apply {
-                dragStartPointerOffset = normalizedOffset - this.offset.toOffset()
                 draggingItemKey = key
                 draggingItemInitialOffset = this.offset.toOffset()
                 return true
@@ -155,7 +153,7 @@
                     contentListState.list.indexOfFirst { it.key == draggingItemKey }
                 )
                 isDraggingToRemove = false
-                updateDragPositionForRemove(Offset.Zero)
+                updateDragPositionForRemove(IntRect.Zero)
             }
             // persist list editing changes on dragging ends
             contentListState.onSaveList()
@@ -164,7 +162,6 @@
         previousTargetItemKey = null
         draggingItemDraggedDelta = Offset.Zero
         draggingItemInitialOffset = Offset.Zero
-        dragStartPointerOffset = Offset.Zero
     }
 
     internal fun onDrag(offset: Offset, layoutDirection: LayoutDirection) {
@@ -230,7 +227,7 @@
             if (overscroll != 0f) {
                 scrollChannel.trySend(overscroll)
             }
-            isDraggingToRemove = checkForRemove(startOffset)
+            isDraggingToRemove = checkForRemove(draggingBoundingBox)
             previousTargetItemKey = null
         }
     }
@@ -247,10 +244,12 @@
     }
 
     /** Calls the callback with the updated drag position and returns whether to remove the item. */
-    private fun checkForRemove(startOffset: Offset): Boolean {
-        return if (draggingItemDraggedDelta.y < 0)
-            updateDragPositionForRemove(startOffset + dragStartPointerOffset)
-        else false
+    private fun checkForRemove(draggingItemBoundingBox: IntRect): Boolean {
+        return if (draggingItemDraggedDelta.y < 0) {
+            updateDragPositionForRemove(draggingItemBoundingBox)
+        } else {
+            false
+        }
     }
 }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ResizeableItemFrame.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ResizeableItemFrame.kt
index 97ad4f1..521330f 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ResizeableItemFrame.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/ResizeableItemFrame.kt
@@ -27,11 +27,14 @@
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.lazy.grid.GridItemSpan
 import androidx.compose.foundation.lazy.grid.LazyGridState
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Alignment
@@ -48,6 +51,8 @@
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.util.fastIsFinite
+import androidx.compose.ui.zIndex
+import com.android.compose.modifiers.thenIf
 import com.android.systemui.communal.ui.viewmodel.DragHandle
 import com.android.systemui.communal.ui.viewmodel.ResizeInfo
 import com.android.systemui.communal.ui.viewmodel.ResizeableItemFrameViewModel
@@ -60,9 +65,12 @@
     viewModel: ResizeableItemFrameViewModel,
     key: String,
     gridState: LazyGridState,
-    minItemSpan: Int,
     gridContentPadding: PaddingValues,
     verticalArrangement: Arrangement.Vertical,
+    minHeightPx: Int,
+    maxHeightPx: Int,
+    resizeMultiple: Int,
+    currentSpan: GridItemSpan,
 ) {
     val density = LocalDensity.current
     LaunchedEffect(
@@ -70,9 +78,12 @@
         viewModel,
         key,
         gridState,
-        minItemSpan,
         gridContentPadding,
         verticalArrangement,
+        minHeightPx,
+        maxHeightPx,
+        resizeMultiple,
+        currentSpan,
     ) {
         val verticalItemSpacingPx = with(density) { verticalArrangement.spacing.toPx() }
         val verticalContentPaddingPx =
@@ -92,13 +103,15 @@
             )
             .collectLatest { (maxItemSpan, viewportHeightPx, itemInfo) ->
                 viewModel.setGridLayoutInfo(
-                    verticalItemSpacingPx,
-                    verticalContentPaddingPx,
-                    viewportHeightPx,
-                    maxItemSpan,
-                    minItemSpan,
-                    itemInfo?.row,
-                    itemInfo?.span,
+                    verticalItemSpacingPx = verticalItemSpacingPx,
+                    currentRow = itemInfo?.row,
+                    maxHeightPx = maxHeightPx,
+                    minHeightPx = minHeightPx,
+                    currentSpan = currentSpan.currentLineSpan,
+                    resizeMultiple = resizeMultiple,
+                    totalSpans = maxItemSpan,
+                    viewportHeightPx = viewportHeightPx,
+                    verticalContentPaddingPx = verticalContentPaddingPx,
                 )
             }
     }
@@ -141,10 +154,9 @@
 /**
  * Draws a frame around the content with drag handles on the top and bottom of the content.
  *
- * @param index The index of this item in the [LazyGridState].
+ * @param key The unique key of this element, must be the same key used in the [LazyGridState].
+ * @param currentSpan The current span size of this item in the grid.
  * @param gridState The [LazyGridState] for the grid containing this item.
- * @param minItemSpan The minimum span that an item may occupy. Items are resized in multiples of
- *   this span.
  * @param gridContentPadding The content padding used for the grid, needed for determining offsets.
  * @param verticalArrangement The vertical arrangement of the grid items.
  * @param modifier Optional modifier to apply to the frame.
@@ -153,6 +165,10 @@
  * @param outlineColor Optional color to make the outline around the content.
  * @param cornerRadius Optional radius to give to the outline around the content.
  * @param strokeWidth Optional stroke width to draw the outline with.
+ * @param minHeightPx Optional minimum height in pixels that this widget can be resized to.
+ * @param maxHeightPx Optional maximum height in pixels that this widget can be resized to.
+ * @param resizeMultiple Optional number of spans that we allow resizing by. For example, if set to
+ *   3, then we only allow resizing in multiples of 3 spans.
  * @param alpha Optional function to provide an alpha value for the outline. Can be used to fade the
  *   outline in and out. This is wrapped in a function for performance, as the value is only
  *   accessed during the draw phase.
@@ -162,8 +178,8 @@
 @Composable
 fun ResizableItemFrame(
     key: String,
+    currentSpan: GridItemSpan,
     gridState: LazyGridState,
-    minItemSpan: Int,
     gridContentPadding: PaddingValues,
     verticalArrangement: Arrangement.Vertical,
     modifier: Modifier = Modifier,
@@ -172,6 +188,9 @@
     outlineColor: Color = MaterialTheme.colorScheme.primary,
     cornerRadius: Dp = 37.dp,
     strokeWidth: Dp = 3.dp,
+    minHeightPx: Int = 0,
+    maxHeightPx: Int = Int.MAX_VALUE,
+    resizeMultiple: Int = 1,
     alpha: () -> Float = { 1f },
     onResize: (info: ResizeInfo) -> Unit = {},
     content: @Composable () -> Unit,
@@ -179,15 +198,24 @@
     val brush = SolidColor(outlineColor)
     val onResizeUpdated by rememberUpdatedState(onResize)
     val viewModel =
-        rememberViewModel(traceName = "ResizeableItemFrame.viewModel") {
+        rememberViewModel(key = currentSpan, traceName = "ResizeableItemFrame.viewModel") {
             ResizeableItemFrameViewModel()
         }
 
     val dragHandleHeight = verticalArrangement.spacing - outlinePadding * 2
+    val isDragging by
+        remember(viewModel) {
+            derivedStateOf {
+                val topOffset = viewModel.topDragState.offset.takeIf { it.fastIsFinite() } ?: 0f
+                val bottomOffset =
+                    viewModel.bottomDragState.offset.takeIf { it.fastIsFinite() } ?: 0f
+                topOffset > 0 || bottomOffset > 0
+            }
+        }
 
     // Draw content surrounded by drag handles at top and bottom. Allow drag handles
     // to overlap content.
-    Box(modifier) {
+    Box(modifier.thenIf(isDragging) { Modifier.zIndex(1f) }) {
         content()
 
         if (enabled) {
@@ -230,12 +258,15 @@
             }
 
             UpdateGridLayoutInfo(
-                viewModel,
-                key,
-                gridState,
-                minItemSpan,
-                gridContentPadding,
-                verticalArrangement,
+                viewModel = viewModel,
+                key = key,
+                gridState = gridState,
+                currentSpan = currentSpan,
+                gridContentPadding = gridContentPadding,
+                verticalArrangement = verticalArrangement,
+                minHeightPx = minHeightPx,
+                maxHeightPx = maxHeightPx,
+                resizeMultiple = resizeMultiple,
             )
             LaunchedEffect(viewModel) {
                 viewModel.resizeInfo.collectLatest { info -> onResizeUpdated(info) }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
index a0fed90..339445e 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
@@ -22,7 +22,7 @@
 import androidx.compose.material3.Text
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.input.nestedscroll.NestedScrollConnection
-import androidx.compose.ui.input.nestedscroll.NestedScrollSource
+import androidx.compose.ui.input.nestedscroll.NestedScrollSource.Companion.UserInput
 import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.IntSize
 import androidx.compose.ui.unit.LayoutDirection
@@ -294,15 +294,10 @@
             available: Offset,
             consumedByScroll: Offset = Offset.Zero,
         ) {
-            val consumedByPreScroll =
-                onPreScroll(available = available, source = NestedScrollSource.Drag)
+            val consumedByPreScroll = onPreScroll(available = available, source = UserInput)
             val consumed = consumedByPreScroll + consumedByScroll
 
-            onPostScroll(
-                consumed = consumed,
-                available = available - consumed,
-                source = NestedScrollSource.Drag,
-            )
+            onPostScroll(consumed = consumed, available = available - consumed, source = UserInput)
         }
 
         fun NestedScrollConnection.preFling(
@@ -738,7 +733,7 @@
         val nestedScroll = nestedScrollConnection(nestedScrollBehavior = EdgeWithPreview)
         nestedScroll.onPreScroll(
             available = downOffset(fractionOfScreen = 0.1f),
-            source = NestedScrollSource.Drag,
+            source = UserInput,
         )
         assertIdle(currentScene = SceneA)
     }
@@ -750,7 +745,7 @@
             nestedScroll.onPostScroll(
                 consumed = Offset.Zero,
                 available = Offset.Zero,
-                source = NestedScrollSource.Drag,
+                source = UserInput,
             )
 
         assertIdle(currentScene = SceneA)
@@ -764,7 +759,7 @@
             nestedScroll.onPostScroll(
                 consumed = Offset.Zero,
                 available = downOffset(fractionOfScreen = 0.1f),
-                source = NestedScrollSource.Drag,
+                source = UserInput,
             )
 
         assertTransition(currentScene = SceneA)
@@ -784,16 +779,12 @@
         val consumed =
             nestedScroll.onPreScroll(
                 available = downOffset(fractionOfScreen = 0.1f),
-                source = NestedScrollSource.Drag,
+                source = UserInput,
             )
         assertThat(progress).isEqualTo(0.2f)
 
         // do nothing on postScroll
-        nestedScroll.onPostScroll(
-            consumed = consumed,
-            available = Offset.Zero,
-            source = NestedScrollSource.Drag,
-        )
+        nestedScroll.onPostScroll(consumed = consumed, available = Offset.Zero, source = UserInput)
         assertThat(progress).isEqualTo(0.2f)
 
         nestedScroll.scroll(available = downOffset(fractionOfScreen = 0.1f))
@@ -813,10 +804,7 @@
         nestedScroll.preFling(available = Velocity.Zero)
 
         // a pre scroll event, that could be intercepted by DraggableHandlerImpl
-        nestedScroll.onPreScroll(
-            available = Offset(0f, secondScroll),
-            source = NestedScrollSource.Drag,
-        )
+        nestedScroll.onPreScroll(available = Offset(0f, secondScroll), source = UserInput)
     }
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/compose/CommunalHubUtilsTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/compose/CommunalHubUtilsTest.kt
deleted file mode 100644
index 643063e7..0000000
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/compose/CommunalHubUtilsTest.kt
+++ /dev/null
@@ -1,41 +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.communal.ui.compose
-
-import android.testing.TestableLooper
-import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.layout.LayoutCoordinates
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.google.common.truth.Truth.assertThat
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.mockito.kotlin.mock
-
-@RunWith(AndroidJUnit4::class)
-@TestableLooper.RunWithLooper(setAsMainLooper = true)
-@SmallTest
-class CommunalHubUtilsTest : SysuiTestCase() {
-    @Test
-    fun isPointerWithinEnabledRemoveButton_ensureDisabledStatePriority() {
-        assertThat(
-                isPointerWithinEnabledRemoveButton(false, mock<Offset>(), mock<LayoutCoordinates>())
-            )
-            .isFalse()
-    }
-}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModelTest.kt
index f0d88ab..22b114c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModelTest.kt
@@ -26,7 +26,6 @@
 import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
-import kotlin.time.Duration
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
@@ -53,10 +52,12 @@
             verticalItemSpacingPx = 10f,
             verticalContentPaddingPx = verticalContentPaddingPx,
             viewportHeightPx = viewportHeightPx,
-            maxItemSpan = 1,
-            minItemSpan = 1,
-            currentSpan = 1,
             currentRow = 0,
+            currentSpan = 1,
+            maxHeightPx = Int.MAX_VALUE,
+            minHeightPx = 0,
+            resizeMultiple = 1,
+            totalSpans = 1,
         )
 
     @Before
@@ -79,7 +80,7 @@
 
     @Test
     fun testSingleSpanGrid() =
-        testScope.runTest(timeout = Duration.INFINITE) {
+        testScope.runTest {
             updateGridLayout(singleSpanGrid)
 
             val topState = underTest.topDragState
@@ -98,8 +99,7 @@
     @Test
     fun testTwoSpanGrid_elementInFirstRow_sizeSingleSpan() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2))
-
+            updateGridLayout(singleSpanGrid.copy(currentRow = 0, totalSpans = 2))
             val topState = underTest.topDragState
             assertThat(topState.currentValue).isEqualTo(0)
             assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
@@ -116,8 +116,7 @@
     @Test
     fun testTwoSpanGrid_elementInSecondRow_sizeSingleSpan() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2, currentRow = 1))
-
+            updateGridLayout(singleSpanGrid.copy(currentRow = 1, totalSpans = 2))
             val topState = underTest.topDragState
             assertThat(topState.currentValue).isEqualTo(0)
             assertThat(topState.anchors.toList()).containsExactly(0 to 0f, -1 to -45f)
@@ -134,15 +133,17 @@
     @Test
     fun testTwoSpanGrid_elementInFirstRow_sizeTwoSpan() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2, currentSpan = 2))
+            val adjustedGridLayout = singleSpanGrid.copy(currentSpan = 2, totalSpans = 2)
+
+            updateGridLayout(adjustedGridLayout)
 
             val topState = underTest.topDragState
-            assertThat(topState.currentValue).isEqualTo(0)
             assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
+            assertThat(topState.currentValue).isEqualTo(0)
 
             val bottomState = underTest.bottomDragState
+            assertThat(bottomState.anchors.toList()).containsExactly(-1 to -45f, 0 to 0f)
             assertThat(bottomState.currentValue).isEqualTo(0)
-            assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f, -1 to -45f)
         }
 
     /**
@@ -151,7 +152,10 @@
     @Test
     fun testThreeSpanGrid_elementInMiddleRow_sizeOneSpan() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 3, currentRow = 1))
+            val adjustedGridLayout =
+                singleSpanGrid.copy(currentRow = 1, currentSpan = 1, totalSpans = 3)
+
+            updateGridLayout(adjustedGridLayout)
 
             val topState = underTest.topDragState
             assertThat(topState.currentValue).isEqualTo(0)
@@ -165,7 +169,10 @@
     @Test
     fun testThreeSpanGrid_elementInTopRow_sizeOneSpan() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 3))
+            val adjustedGridLayout =
+                singleSpanGrid.copy(currentRow = 0, currentSpan = 1, totalSpans = 3)
+
+            updateGridLayout(adjustedGridLayout)
 
             val topState = underTest.topDragState
             assertThat(topState.currentValue).isEqualTo(0)
@@ -177,16 +184,17 @@
         }
 
     @Test
-    fun testSixSpanGrid_minSpanThree_itemInThirdRow_sizeThreeSpans() =
+    fun testSixSpanGrid_minSpanThree_itemInFourthRow_sizeThreeSpans() =
         testScope.runTest {
-            updateGridLayout(
+            val adjustedGridLayout =
                 singleSpanGrid.copy(
-                    maxItemSpan = 6,
                     currentRow = 3,
                     currentSpan = 3,
-                    minItemSpan = 3,
+                    resizeMultiple = 3,
+                    totalSpans = 6,
                 )
-            )
+
+            updateGridLayout(adjustedGridLayout)
 
             val topState = underTest.topDragState
             assertThat(topState.currentValue).isEqualTo(0)
@@ -200,7 +208,14 @@
     @Test
     fun testTwoSpanGrid_elementMovesFromFirstRowToSecondRow() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2))
+            val firstRowLayout =
+                singleSpanGrid.copy(
+                    currentRow = 0,
+                    currentSpan = 1,
+                    resizeMultiple = 1,
+                    totalSpans = 2,
+                )
+            updateGridLayout(firstRowLayout)
 
             val topState = underTest.topDragState
             val bottomState = underTest.bottomDragState
@@ -208,7 +223,8 @@
             assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
             assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f, 1 to 45f)
 
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2, currentRow = 1))
+            val secondRowLayout = firstRowLayout.copy(currentRow = 1)
+            updateGridLayout(secondRowLayout)
 
             assertThat(topState.anchors.toList()).containsExactly(0 to 0f, -1 to -45f)
             assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f)
@@ -217,19 +233,21 @@
     @Test
     fun testTwoSpanGrid_expandElementFromBottom() = runTestWithSnapshots {
         val resizeInfo by collectLastValue(underTest.resizeInfo)
-        updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2))
 
-        assertThat(resizeInfo).isNull()
+        val adjustedGridLayout = singleSpanGrid.copy(resizeMultiple = 1, totalSpans = 2)
+
+        updateGridLayout(adjustedGridLayout)
+
         underTest.bottomDragState.anchoredDrag { dragTo(45f) }
+
         assertThat(resizeInfo).isEqualTo(ResizeInfo(1, DragHandle.BOTTOM))
     }
 
     @Test
     fun testThreeSpanGrid_expandMiddleElementUpwards() = runTestWithSnapshots {
         val resizeInfo by collectLastValue(underTest.resizeInfo)
-        updateGridLayout(singleSpanGrid.copy(maxItemSpan = 3, currentRow = 1))
+        updateGridLayout(singleSpanGrid.copy(currentRow = 1, totalSpans = 3))
 
-        assertThat(resizeInfo).isNull()
         underTest.topDragState.anchoredDrag { dragTo(-30f) }
         assertThat(resizeInfo).isEqualTo(ResizeInfo(1, DragHandle.TOP))
     }
@@ -237,7 +255,7 @@
     @Test
     fun testThreeSpanGrid_expandTopElementDownBy2Spans() = runTestWithSnapshots {
         val resizeInfo by collectLastValue(underTest.resizeInfo)
-        updateGridLayout(singleSpanGrid.copy(maxItemSpan = 3))
+        updateGridLayout(singleSpanGrid.copy(totalSpans = 3))
 
         assertThat(resizeInfo).isNull()
         underTest.bottomDragState.anchoredDrag { dragTo(60f) }
@@ -247,7 +265,7 @@
     @Test
     fun testTwoSpanGrid_shrinkElementFromBottom() = runTestWithSnapshots {
         val resizeInfo by collectLastValue(underTest.resizeInfo)
-        updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2, currentSpan = 2))
+        updateGridLayout(singleSpanGrid.copy(totalSpans = 2, currentSpan = 2))
 
         assertThat(resizeInfo).isNull()
         underTest.bottomDragState.anchoredDrag { dragTo(-45f) }
@@ -257,7 +275,7 @@
     @Test
     fun testRowInfoBecomesNull_revertsBackToDefault() =
         testScope.runTest {
-            val gridLayout = singleSpanGrid.copy(maxItemSpan = 3, currentRow = 1)
+            val gridLayout = singleSpanGrid.copy(currentRow = 1, resizeMultiple = 1, totalSpans = 3)
             updateGridLayout(gridLayout)
 
             val topState = underTest.topDragState
@@ -266,44 +284,113 @@
             val bottomState = underTest.bottomDragState
             assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f, 1 to 30f)
 
+            // Set currentRow to null to simulate the row info becoming null
             updateGridLayout(gridLayout.copy(currentRow = null))
             assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
             assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f)
         }
 
-    @Test(expected = IllegalArgumentException::class)
-    fun testIllegalState_maxSpanSmallerThanMinSpan() =
+    @Test
+    fun testEqualMaxAndMinHeight_cannotResize() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2, minItemSpan = 3))
+            val heightPx = 20
+            updateGridLayout(
+                singleSpanGrid.copy(maxHeightPx = heightPx, minHeightPx = heightPx, totalSpans = 2)
+            )
+
+            val topState = underTest.topDragState
+            val bottomState = underTest.bottomDragState
+
+            assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
+            assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f)
+        }
+
+    @Test
+    fun testMinHeightTwoRows_canExpandButNotShrink() =
+        testScope.runTest {
+            val threeRowGrid =
+                singleSpanGrid.copy(
+                    maxHeightPx = 80,
+                    minHeightPx = 50,
+                    totalSpans = 3,
+                    currentSpan = 2,
+                    currentRow = 0,
+                )
+
+            updateGridLayout(threeRowGrid)
+
+            val topState = underTest.topDragState
+            val bottomState = underTest.bottomDragState
+            assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
+            assertThat(bottomState.anchors.toList()).containsAtLeast(0 to 0f, 1 to 30f)
+        }
+
+    @Test
+    fun testMaxHeightTwoRows_canShrinkButNotExpand() =
+        testScope.runTest {
+            val threeRowGrid =
+                singleSpanGrid.copy(
+                    maxHeightPx = 50,
+                    minHeightPx = 20,
+                    totalSpans = 3,
+                    currentSpan = 2,
+                    currentRow = 0,
+                )
+
+            updateGridLayout(threeRowGrid)
+
+            val topState = underTest.topDragState
+            val bottomState = underTest.bottomDragState
+
+            assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f, -1 to -30f)
+
+            assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
+        }
+
+    @Test
+    fun testMinHeightEqualToAvailableSpan_cannotResize() =
+        testScope.runTest {
+            val twoRowGrid =
+                singleSpanGrid.copy(
+                    minHeightPx = (viewportHeightPx - verticalContentPaddingPx.toInt()),
+                    totalSpans = 2,
+                    currentSpan = 2,
+                )
+
+            updateGridLayout(twoRowGrid)
+
+            val topState = underTest.topDragState
+            val bottomState = underTest.bottomDragState
+
+            assertThat(topState.anchors.toList()).containsExactly(0 to 0f)
+            assertThat(bottomState.anchors.toList()).containsExactly(0 to 0f)
         }
 
     @Test(expected = IllegalArgumentException::class)
-    fun testIllegalState_minSpanOfZero() =
+    fun testIllegalState_maxHeightLessThanMinHeight() =
         testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 2, minItemSpan = 0))
+            updateGridLayout(singleSpanGrid.copy(maxHeightPx = 50, minHeightPx = 100))
         }
 
     @Test(expected = IllegalArgumentException::class)
-    fun testIllegalState_maxSpanOfZero() =
-        testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 0, minItemSpan = 0))
-        }
+    fun testIllegalState_currentSpanExceedsTotalSpans() =
+        testScope.runTest { updateGridLayout(singleSpanGrid.copy(currentSpan = 3, totalSpans = 2)) }
 
     @Test(expected = IllegalArgumentException::class)
-    fun testIllegalState_currentRowNotMultipleOfMinSpan() =
-        testScope.runTest {
-            updateGridLayout(singleSpanGrid.copy(maxItemSpan = 6, minItemSpan = 3, currentSpan = 2))
-        }
+    fun testIllegalState_resizeMultipleZeroOrNegative() =
+        testScope.runTest { updateGridLayout(singleSpanGrid.copy(resizeMultiple = 0)) }
 
     private fun TestScope.updateGridLayout(gridLayout: GridLayout) {
         underTest.setGridLayoutInfo(
-            gridLayout.verticalItemSpacingPx,
-            gridLayout.verticalContentPaddingPx,
-            gridLayout.viewportHeightPx,
-            gridLayout.maxItemSpan,
-            gridLayout.minItemSpan,
-            gridLayout.currentRow,
-            gridLayout.currentSpan,
+            verticalItemSpacingPx = gridLayout.verticalItemSpacingPx,
+            currentRow = gridLayout.currentRow,
+            maxHeightPx = gridLayout.maxHeightPx,
+            minHeightPx = gridLayout.minHeightPx,
+            currentSpan = gridLayout.currentSpan,
+            resizeMultiple = gridLayout.resizeMultiple,
+            totalSpans = gridLayout.totalSpans,
+            viewportHeightPx = gridLayout.viewportHeightPx,
+            verticalContentPaddingPx = gridLayout.verticalContentPaddingPx,
         )
         runCurrent()
     }
@@ -332,9 +419,11 @@
         val verticalItemSpacingPx: Float,
         val verticalContentPaddingPx: Float,
         val viewportHeightPx: Int,
-        val maxItemSpan: Int,
-        val minItemSpan: Int,
         val currentRow: Int?,
-        val currentSpan: Int?,
+        val currentSpan: Int,
+        val maxHeightPx: Int,
+        val minHeightPx: Int,
+        val resizeMultiple: Int,
+        val totalSpans: Int,
     )
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
index 8201bbe..e7d2ef1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
@@ -30,8 +30,11 @@
 import com.android.systemui.education.data.repository.contextualEducationRepository
 import com.android.systemui.education.data.repository.fakeEduClock
 import com.android.systemui.education.shared.model.EducationUiType
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
+import com.android.systemui.inputdevice.tutorial.tutorialSchedulerRepository
 import com.android.systemui.keyboard.data.repository.keyboardRepository
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.recents.OverviewProxyService.OverviewProxyListener
 import com.android.systemui.testKosmos
 import com.android.systemui.touchpad.data.repository.touchpadRepository
 import com.android.systemui.user.data.repository.fakeUserRepository
@@ -42,10 +45,13 @@
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import org.junit.After
 import org.junit.Assume.assumeTrue
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.verify
 import platform.test.runner.parameterized.ParameterizedAndroidJunit4
 import platform.test.runner.parameterized.Parameters
 
@@ -56,14 +62,19 @@
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val contextualEduInteractor = kosmos.contextualEducationInteractor
+    private val repository = kosmos.contextualEducationRepository
     private val touchpadRepository = kosmos.touchpadRepository
     private val keyboardRepository = kosmos.keyboardRepository
+    private val tutorialSchedulerRepository = kosmos.tutorialSchedulerRepository
     private val userRepository = kosmos.fakeUserRepository
+    private val overviewProxyService = kosmos.mockOverviewProxyService
 
     private val underTest: KeyboardTouchpadEduInteractor = kosmos.keyboardTouchpadEduInteractor
     private val eduClock = kosmos.fakeEduClock
     private val minDurationForNextEdu =
         KeyboardTouchpadEduInteractor.minIntervalBetweenEdu + 1.seconds
+    private val initialDelayElapsedDuration =
+        KeyboardTouchpadEduInteractor.initialDelayDuration + 1.seconds
 
     @Before
     fun setup() {
@@ -271,6 +282,131 @@
             assertThat(model?.lastShortcutTriggeredTime).isEqualTo(eduClock.instant())
         }
 
+    @Test
+    fun dataUpdatedOnIncrementSignalCountWhenTouchpadConnected() =
+        testScope.runTest {
+            assumeTrue(gestureType != ALL_APPS)
+            setUpForInitialDelayElapse()
+            touchpadRepository.setIsAnyTouchpadConnected(true)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWhenTouchpadDisconnected() =
+        testScope.runTest {
+            setUpForInitialDelayElapse()
+            touchpadRepository.setIsAnyTouchpadConnected(false)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountWhenKeyboardConnected() =
+        testScope.runTest {
+            assumeTrue(gestureType == ALL_APPS)
+            setUpForInitialDelayElapse()
+            keyboardRepository.setIsAnyKeyboardConnected(true)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWhenKeyboardDisconnected() =
+        testScope.runTest {
+            setUpForInitialDelayElapse()
+            keyboardRepository.setIsAnyKeyboardConnected(false)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataAddedOnUpdateShortcutTriggerTime() =
+        testScope.runTest {
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            assertThat(model?.lastShortcutTriggeredTime).isNull()
+
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ true, gestureType)
+
+            assertThat(model?.lastShortcutTriggeredTime).isEqualTo(kosmos.fakeEduClock.instant())
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountAfterInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, eduClock.instant())
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            eduClock.offset(initialDelayElapsedDuration)
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountBeforeInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, eduClock.instant())
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            // No offset to the clock to simulate update before initial delay
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWithoutOobeLaunchTime() =
+        testScope.runTest {
+            // No update to OOBE launch time to simulate no OOBE is launched yet
+            setUpForDeviceConnection()
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    private suspend fun setUpForInitialDelayElapse() {
+        tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, eduClock.instant())
+        tutorialSchedulerRepository.updateLaunchTime(DeviceType.KEYBOARD, eduClock.instant())
+        eduClock.offset(initialDelayElapsedDuration)
+    }
+
+    @After
+    fun clear() {
+        testScope.launch { tutorialSchedulerRepository.clearDataStore() }
+    }
+
     private suspend fun triggerMaxEducationSignals(gestureType: GestureType) {
         // Increment max number of signal to try triggering education
         for (i in 1..KeyboardTouchpadEduInteractor.MAX_SIGNAL_COUNT) {
@@ -288,9 +424,15 @@
         runCurrent()
     }
 
-    private suspend fun setUpForDeviceConnection() {
-        contextualEduInteractor.updateKeyboardFirstConnectionTime()
-        contextualEduInteractor.updateTouchpadFirstConnectionTime()
+    private fun setUpForDeviceConnection() {
+        touchpadRepository.setIsAnyTouchpadConnected(true)
+        keyboardRepository.setIsAnyKeyboardConnected(true)
+    }
+
+    private fun getOverviewProxyListener(): OverviewProxyListener {
+        val listenerCaptor = argumentCaptor<OverviewProxyListener>()
+        verify(overviewProxyService).addCallback(listenerCaptor.capture())
+        return listenerCaptor.firstValue
     }
 
     companion object {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt
deleted file mode 100644
index 98e0947..0000000
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt
+++ /dev/null
@@ -1,172 +0,0 @@
-/*
- * Copyright 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.education.domain.interactor
-
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.contextualeducation.GestureType.ALL_APPS
-import com.android.systemui.contextualeducation.GestureType.BACK
-import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.education.data.repository.contextualEducationRepository
-import com.android.systemui.education.data.repository.fakeEduClock
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
-import com.android.systemui.inputdevice.tutorial.tutorialSchedulerRepository
-import com.android.systemui.keyboard.data.repository.keyboardRepository
-import com.android.systemui.kosmos.testScope
-import com.android.systemui.testKosmos
-import com.android.systemui.touchpad.data.repository.touchpadRepository
-import com.google.common.truth.Truth.assertThat
-import kotlin.time.Duration.Companion.seconds
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.test.runTest
-import org.junit.After
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@SmallTest
-@RunWith(AndroidJUnit4::class)
-class KeyboardTouchpadStatsInteractorTest : SysuiTestCase() {
-    private val kosmos = testKosmos()
-    private val testScope = kosmos.testScope
-    private val underTest = kosmos.keyboardTouchpadEduStatsInteractor
-    private val keyboardRepository = kosmos.keyboardRepository
-    private val touchpadRepository = kosmos.touchpadRepository
-    private val repository = kosmos.contextualEducationRepository
-    private val fakeClock = kosmos.fakeEduClock
-    private val tutorialSchedulerRepository = kosmos.tutorialSchedulerRepository
-    private val initialDelayElapsedDuration =
-        KeyboardTouchpadEduStatsInteractorImpl.initialDelayDuration + 1.seconds
-
-    @Test
-    fun dataUpdatedOnIncrementSignalCountWhenTouchpadConnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            touchpadRepository.setIsAnyTouchpadConnected(true)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountWhenTouchpadDisconnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            touchpadRepository.setIsAnyTouchpadConnected(false)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    @Test
-    fun dataUpdatedOnIncrementSignalCountWhenKeyboardConnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            keyboardRepository.setIsAnyKeyboardConnected(true)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(ALL_APPS)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountWhenKeyboardDisconnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            keyboardRepository.setIsAnyKeyboardConnected(false)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(ALL_APPS)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    @Test
-    fun dataAddedOnUpdateShortcutTriggerTime() =
-        testScope.runTest {
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            assertThat(model?.lastShortcutTriggeredTime).isNull()
-            underTest.updateShortcutTriggerTime(BACK)
-            assertThat(model?.lastShortcutTriggeredTime).isEqualTo(kosmos.fakeEduClock.instant())
-        }
-
-    @Test
-    fun dataUpdatedOnIncrementSignalCountAfterInitialDelay() =
-        testScope.runTest {
-            setUpForDeviceConnection()
-            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, fakeClock.instant())
-
-            fakeClock.offset(initialDelayElapsedDuration)
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountBeforeInitialDelay() =
-        testScope.runTest {
-            setUpForDeviceConnection()
-            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, fakeClock.instant())
-
-            // No offset to the clock to simulate update before initial delay
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountWithoutOobeLaunchTime() =
-        testScope.runTest {
-            // No update to OOBE launch time to simulate no OOBE is launched yet
-            setUpForDeviceConnection()
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    private suspend fun setUpForInitialDelayElapse() {
-        tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, fakeClock.instant())
-        tutorialSchedulerRepository.updateLaunchTime(DeviceType.KEYBOARD, fakeClock.instant())
-        fakeClock.offset(initialDelayElapsedDuration)
-    }
-
-    private fun setUpForDeviceConnection() {
-        touchpadRepository.setIsAnyTouchpadConnected(true)
-        keyboardRepository.setIsAnyKeyboardConnected(true)
-    }
-
-    @After
-    fun clear() {
-        testScope.launch { tutorialSchedulerRepository.clearDataStore() }
-    }
-}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarterTest.kt
similarity index 72%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarterTest.kt
index f8e2f47..d1431ee 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarterTest.kt
@@ -16,25 +16,26 @@
 
 package com.android.systemui.keyboard.shortcut.ui
 
-import android.content.Intent
+import androidx.test.annotation.UiThreadTest
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyboard.shortcut.data.source.FakeKeyboardShortcutGroupsSource
 import com.android.systemui.keyboard.shortcut.data.source.TestShortcuts
-import com.android.systemui.keyboard.shortcut.fakeShortcutHelperStartActivity
-import com.android.systemui.keyboard.shortcut.shortcutHelperActivityStarter
 import com.android.systemui.keyboard.shortcut.shortcutHelperAppCategoriesShortcutsSource
 import com.android.systemui.keyboard.shortcut.shortcutHelperCurrentAppShortcutsSource
 import com.android.systemui.keyboard.shortcut.shortcutHelperInputShortcutsSource
 import com.android.systemui.keyboard.shortcut.shortcutHelperMultiTaskingShortcutsSource
 import com.android.systemui.keyboard.shortcut.shortcutHelperSystemShortcutsSource
 import com.android.systemui.keyboard.shortcut.shortcutHelperTestHelper
-import com.android.systemui.keyboard.shortcut.ui.view.ShortcutHelperActivity
+import com.android.systemui.keyboard.shortcut.shortcutHelperViewModel
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.kosmos.testCase
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.plugins.activityStarter
+import com.android.systemui.statusbar.phone.systemUIDialogFactory
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
@@ -46,7 +47,7 @@
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
-class ShortcutHelperActivityStarterTest : SysuiTestCase() {
+class ShortcutHelperDialogStarterTest : SysuiTestCase() {
 
     private val fakeSystemSource = FakeKeyboardShortcutGroupsSource()
     private val fakeMultiTaskingSource = FakeKeyboardShortcutGroupsSource()
@@ -64,8 +65,14 @@
 
     private val testScope = kosmos.testScope
     private val testHelper = kosmos.shortcutHelperTestHelper
-    private val fakeStartActivity = kosmos.fakeShortcutHelperStartActivity
-    private val starter = kosmos.shortcutHelperActivityStarter
+    private val dialogFactory = kosmos.systemUIDialogFactory
+    private val coroutineScope = kosmos.applicationCoroutineScope
+    private val viewModel = kosmos.shortcutHelperViewModel
+
+    private val starter: ShortcutHelperDialogStarter =
+        with(kosmos) {
+            ShortcutHelperDialogStarter(coroutineScope, viewModel, dialogFactory, activityStarter)
+        }
 
     @Before
     fun setUp() {
@@ -74,21 +81,22 @@
     }
 
     @Test
-    fun start_doesNotStartByDefault() =
+    fun start_doesNotShowDialogByDefault() =
         testScope.runTest {
             starter.start()
 
-            assertThat(fakeStartActivity.startIntents).isEmpty()
+            assertThat(starter.dialog).isNull()
         }
 
     @Test
-    fun start_onToggle_startsActivity() =
+    @UiThreadTest
+    fun start_onToggle_showsDialog() =
         testScope.runTest {
             starter.start()
 
             testHelper.toggle(deviceId = 456)
 
-            verifyShortcutHelperActivityStarted()
+            assertThat(starter.dialog?.isShowing).isTrue()
         }
 
     @Test
@@ -101,34 +109,18 @@
 
             testHelper.toggle(deviceId = 456)
 
-            assertThat(fakeStartActivity.startIntents).isEmpty()
+            assertThat(starter.dialog).isNull()
         }
 
     @Test
-    fun start_onToggle_multipleTimesStartsActivityOnlyWhenNotStarted() =
-        testScope.runTest {
-            starter.start()
-
-            // Starts
-            testHelper.toggle(deviceId = 456)
-            // Stops
-            testHelper.toggle(deviceId = 456)
-            // Starts again
-            testHelper.toggle(deviceId = 456)
-            // Stops
-            testHelper.toggle(deviceId = 456)
-
-            verifyShortcutHelperActivityStarted(numTimes = 2)
-        }
-
-    @Test
+    @UiThreadTest
     fun start_onRequestShowShortcuts_startsActivity() =
         testScope.runTest {
             starter.start()
 
             testHelper.showFromActivity()
 
-            verifyShortcutHelperActivityStarted()
+            assertThat(starter.dialog?.isShowing).isTrue()
         }
 
     @Test
@@ -140,10 +132,11 @@
 
             testHelper.showFromActivity()
 
-            assertThat(fakeStartActivity.startIntents).isEmpty()
+            assertThat(starter.dialog).isNull()
         }
 
     @Test
+    @UiThreadTest
     fun start_onRequestShowShortcuts_multipleTimes_startsActivityOnlyOnce() =
         testScope.runTest {
             starter.start()
@@ -152,40 +145,40 @@
             testHelper.showFromActivity()
             testHelper.showFromActivity()
 
-            verifyShortcutHelperActivityStarted(numTimes = 1)
+            assertThat(starter.dialog?.isShowing).isTrue()
         }
 
     @Test
+    @UiThreadTest
     fun start_onRequestShowShortcuts_multipleTimes_startsActivityOnlyWhenNotStarted() =
         testScope.runTest {
             starter.start()
 
+            assertThat(starter.dialog).isNull()
             // No-op. Already hidden.
             testHelper.hideFromActivity()
+            assertThat(starter.dialog).isNull()
             // No-op. Already hidden.
             testHelper.hideForSystem()
+            assertThat(starter.dialog).isNull()
             // Show 1st time.
             testHelper.toggle(deviceId = 987)
+            assertThat(starter.dialog).isNotNull()
+            assertThat(starter.dialog?.isShowing).isTrue()
             // No-op. Already shown.
             testHelper.showFromActivity()
+            assertThat(starter.dialog?.isShowing).isTrue()
             // Hidden.
             testHelper.hideFromActivity()
+            assertThat(starter.dialog?.isShowing).isFalse()
             // No-op. Already hidden.
             testHelper.hideForSystem()
+            assertThat(starter.dialog?.isShowing).isFalse()
             // Show 2nd time.
             testHelper.toggle(deviceId = 456)
+            assertThat(starter.dialog?.isShowing).isTrue()
             // No-op. Already shown.
             testHelper.showFromActivity()
-
-            verifyShortcutHelperActivityStarted(numTimes = 2)
+            assertThat(starter.dialog?.isShowing).isTrue()
         }
-
-    private fun verifyShortcutHelperActivityStarted(numTimes: Int = 1) {
-        assertThat(fakeStartActivity.startIntents).hasSize(numTimes)
-        fakeStartActivity.startIntents.forEach { intent ->
-            assertThat(intent.flags).isEqualTo(Intent.FLAG_ACTIVITY_NEW_TASK)
-            assertThat(intent.filterEquals(Intent(context, ShortcutHelperActivity::class.java)))
-                .isTrue()
-        }
-    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/AbstractQSFragmentComposeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/AbstractQSFragmentComposeViewModelTest.kt
index 4bbdfa4..d96e664 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/AbstractQSFragmentComposeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/AbstractQSFragmentComposeViewModelTest.kt
@@ -21,10 +21,11 @@
 import androidx.lifecycle.testing.TestLifecycleOwner
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.testKosmos
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.TestResult
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.resetMain
@@ -62,7 +63,7 @@
     ): TestResult {
         return runTest {
             lifecycleOwner.setCurrentState(Lifecycle.State.RESUMED)
-            lifecycleOwner.lifecycleScope.launch { underTest.activate() }
+            underTest.activateIn(kosmos.testScope)
             block().also { lifecycleOwner.setCurrentState(Lifecycle.State.DESTROYED) }
         }
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelForceQSTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelForceQSTest.kt
index da16640..d16da1c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelForceQSTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelForceQSTest.kt
@@ -18,12 +18,13 @@
 
 import android.testing.TestableLooper.RunWithLooper
 import androidx.test.filters.SmallTest
-import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.statusbar.StatusBarState
 import com.android.systemui.statusbar.sysuiStatusBarStateController
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
 import org.junit.Test
 import org.junit.runner.RunWith
 import platform.test.runner.parameterized.ParameterizedAndroidJunit4
@@ -32,6 +33,7 @@
 @SmallTest
 @RunWith(ParameterizedAndroidJunit4::class)
 @RunWithLooper
+@OptIn(ExperimentalCoroutinesApi::class)
 class QSFragmentComposeViewModelForceQSTest(private val testData: TestData) :
     AbstractQSFragmentComposeViewModelTest() {
 
@@ -39,18 +41,18 @@
     fun forceQs_orRealExpansion() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val expansionState by collectLastValue(underTest.expansionState)
-
                 with(testData) {
                     sysuiStatusBarStateController.setState(statusBarState)
-                    underTest.isQSExpanded = expanded
+                    underTest.isQsExpanded = expanded
                     underTest.isStackScrollerOverscrolling = stackScrollerOverScrolling
                     fakeDeviceEntryRepository.setBypassEnabled(bypassEnabled)
                     underTest.isTransitioningToFullShade = transitioningToFullShade
                     underTest.isInSplitShade = inSplitShade
 
-                    underTest.qsExpansionValue = EXPANSION
-                    assertThat(expansionState!!.progress)
+                    underTest.setQsExpansionValue(EXPANSION)
+
+                    runCurrent()
+                    assertThat(underTest.expansionState.progress)
                         .isEqualTo(if (expectedForceQS) 1f else EXPANSION)
                 }
             }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
index c19e4b8..3b00f86 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
@@ -19,6 +19,7 @@
 import android.app.StatusBarManager
 import android.content.testableContext
 import android.testing.TestableLooper.RunWithLooper
+import androidx.compose.runtime.snapshots.Snapshot
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.common.ui.data.repository.fakeConfigurationRepository
@@ -33,6 +34,7 @@
 import com.android.systemui.statusbar.disableflags.data.repository.fakeDisableFlagsRepository
 import com.android.systemui.statusbar.sysuiStatusBarStateController
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.runCurrent
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -40,22 +42,21 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 @RunWithLooper
+@OptIn(ExperimentalCoroutinesApi::class)
 class QSFragmentComposeViewModelTest : AbstractQSFragmentComposeViewModelTest() {
 
     @Test
     fun qsExpansionValueChanges_correctExpansionState() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val expansionState by collectLastValue(underTest.expansionState)
+                underTest.setQsExpansionValue(0f)
+                assertThat(underTest.expansionState.progress).isEqualTo(0f)
 
-                underTest.qsExpansionValue = 0f
-                assertThat(expansionState!!.progress).isEqualTo(0f)
+                underTest.setQsExpansionValue(0.3f)
+                assertThat(underTest.expansionState.progress).isEqualTo(0.3f)
 
-                underTest.qsExpansionValue = 0.3f
-                assertThat(expansionState!!.progress).isEqualTo(0.3f)
-
-                underTest.qsExpansionValue = 1f
-                assertThat(expansionState!!.progress).isEqualTo(1f)
+                underTest.setQsExpansionValue(1f)
+                assertThat(underTest.expansionState.progress).isEqualTo(1f)
             }
         }
 
@@ -63,13 +64,11 @@
     fun qsExpansionValueChanges_clamped() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val expansionState by collectLastValue(underTest.expansionState)
+                underTest.setQsExpansionValue(-1f)
+                assertThat(underTest.expansionState.progress).isEqualTo(0f)
 
-                underTest.qsExpansionValue = -1f
-                assertThat(expansionState!!.progress).isEqualTo(0f)
-
-                underTest.qsExpansionValue = 2f
-                assertThat(expansionState!!.progress).isEqualTo(1f)
+                underTest.setQsExpansionValue(2f)
+                assertThat(underTest.expansionState.progress).isEqualTo(1f)
             }
         }
 
@@ -77,15 +76,13 @@
     fun qqsHeaderHeight_largeScreenHeader_0() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val qqsHeaderHeight by collectLastValue(underTest.qqsHeaderHeight)
-
                 testableContext.orCreateTestableResources.addOverride(
                     R.bool.config_use_large_screen_shade_header,
                     true,
                 )
                 fakeConfigurationRepository.onConfigurationChange()
 
-                assertThat(qqsHeaderHeight).isEqualTo(0)
+                assertThat(underTest.qqsHeaderHeight).isEqualTo(0)
             }
         }
 
@@ -93,15 +90,13 @@
     fun qqsHeaderHeight_noLargeScreenHeader_providedByHelper() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val qqsHeaderHeight by collectLastValue(underTest.qqsHeaderHeight)
-
                 testableContext.orCreateTestableResources.addOverride(
                     R.bool.config_use_large_screen_shade_header,
                     false,
                 )
                 fakeConfigurationRepository.onConfigurationChange()
 
-                assertThat(qqsHeaderHeight)
+                assertThat(underTest.qqsHeaderHeight)
                     .isEqualTo(largeScreenHeaderHelper.getLargeScreenHeaderHeight())
             }
         }
@@ -120,17 +115,17 @@
     fun statusBarState_followsController() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val statusBarState by collectLastValue(underTest.statusBarState)
-                runCurrent()
-
                 sysuiStatusBarStateController.setState(StatusBarState.SHADE)
-                assertThat(statusBarState).isEqualTo(StatusBarState.SHADE)
+                runCurrent()
+                assertThat(underTest.statusBarState).isEqualTo(StatusBarState.SHADE)
 
                 sysuiStatusBarStateController.setState(StatusBarState.KEYGUARD)
-                assertThat(statusBarState).isEqualTo(StatusBarState.KEYGUARD)
+                runCurrent()
+                assertThat(underTest.statusBarState).isEqualTo(StatusBarState.KEYGUARD)
 
                 sysuiStatusBarStateController.setState(StatusBarState.SHADE_LOCKED)
-                assertThat(statusBarState).isEqualTo(StatusBarState.SHADE_LOCKED)
+                runCurrent()
+                assertThat(underTest.statusBarState).isEqualTo(StatusBarState.SHADE_LOCKED)
             }
         }
 
@@ -138,17 +133,18 @@
     fun statusBarState_changesEarlyIfUpcomingStateIsKeyguard() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val statusBarState by collectLastValue(underTest.statusBarState)
-
                 sysuiStatusBarStateController.setState(StatusBarState.SHADE)
                 sysuiStatusBarStateController.setUpcomingState(StatusBarState.SHADE_LOCKED)
-                assertThat(statusBarState).isEqualTo(StatusBarState.SHADE)
+                runCurrent()
+                assertThat(underTest.statusBarState).isEqualTo(StatusBarState.SHADE)
 
                 sysuiStatusBarStateController.setUpcomingState(StatusBarState.KEYGUARD)
-                assertThat(statusBarState).isEqualTo(StatusBarState.KEYGUARD)
+                runCurrent()
+                assertThat(underTest.statusBarState).isEqualTo(StatusBarState.KEYGUARD)
 
                 sysuiStatusBarStateController.setUpcomingState(StatusBarState.SHADE)
-                assertThat(statusBarState).isEqualTo(StatusBarState.KEYGUARD)
+                runCurrent()
+                assertThat(underTest.statusBarState).isEqualTo(StatusBarState.KEYGUARD)
             }
         }
 
@@ -156,16 +152,16 @@
     fun qsEnabled_followsRepository() =
         with(kosmos) {
             testScope.testWithinLifecycle {
-                val qsEnabled by collectLastValue(underTest.qsEnabled)
-
                 fakeDisableFlagsRepository.disableFlags.value =
                     DisableFlagsModel(disable2 = QS_DISABLE_FLAG)
+                runCurrent()
 
-                assertThat(qsEnabled).isFalse()
+                assertThat(underTest.isQsEnabled).isFalse()
 
                 fakeDisableFlagsRepository.disableFlags.value = DisableFlagsModel()
+                runCurrent()
 
-                assertThat(qsEnabled).isTrue()
+                assertThat(underTest.isQsEnabled).isTrue()
             }
         }
 
@@ -175,13 +171,16 @@
             testScope.testWithinLifecycle {
                 val squishiness by collectLastValue(tileSquishinessInteractor.squishiness)
 
-                underTest.squishinessFractionValue = 0.3f
+                underTest.squishinessFraction = 0.3f
+                Snapshot.sendApplyNotifications()
                 assertThat(squishiness).isWithin(epsilon).of(0.3f.constrainSquishiness())
 
-                underTest.squishinessFractionValue = 0f
+                underTest.squishinessFraction = 0f
+                Snapshot.sendApplyNotifications()
                 assertThat(squishiness).isWithin(epsilon).of(0f.constrainSquishiness())
 
-                underTest.squishinessFractionValue = 1f
+                underTest.squishinessFraction = 1f
+                Snapshot.sendApplyNotifications()
                 assertThat(squishiness).isWithin(epsilon).of(1f.constrainSquishiness())
             }
         }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractorTest.kt
new file mode 100644
index 0000000..75d4b91
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractorTest.kt
@@ -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.systemui.qs.panels.domain.interactor
+
+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.kosmos.testScope
+import com.android.systemui.lifecycle.activateIn
+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.qsPreferencesRepository
+import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
+import com.android.systemui.qs.pipeline.shared.TileSpec
+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.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class DynamicIconTilesInteractorTest : SysuiTestCase() {
+    private val kosmos =
+        testKosmos().apply {
+            defaultLargeTilesRepository =
+                object : DefaultLargeTilesRepository {
+                    override val defaultLargeTiles: Set<TileSpec> = setOf(largeTile)
+                }
+            currentTilesInteractor.setTiles(listOf(largeTile, smallTile))
+        }
+    private lateinit var underTest: DynamicIconTilesInteractor
+
+    @Before
+    fun setUp() {
+        with(kosmos) {
+            underTest = dynamicIconTilesInteractorFactory.create()
+            underTest.activateIn(testScope)
+        }
+    }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Test
+    fun removingTile_updatesSharedPreferences() =
+        with(kosmos) {
+            testScope.runTest {
+                val latest by collectLastValue(qsPreferencesRepository.largeTilesSpecs)
+                runCurrent()
+
+                // Remove the large tile from the current tiles
+                currentTilesInteractor.removeTiles(listOf(largeTile))
+                runCurrent()
+
+                // Assert that it resized to small
+                assertThat(latest).doesNotContain(largeTile)
+            }
+        }
+
+    private companion object {
+        private val largeTile = TileSpec.create("large")
+        private val smallTile = TileSpec.create("small")
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractorTest.kt
index 79a303d..ed28dc8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractorTest.kt
@@ -98,23 +98,6 @@
 
     @OptIn(ExperimentalCoroutinesApi::class)
     @Test
-    fun removingTile_updatesSharedPreferences() =
-        with(kosmos) {
-            testScope.runTest {
-                val latest by collectLastValue(qsPreferencesRepository.largeTilesSpecs)
-                runCurrent()
-
-                // Remove the large tile from the current tiles
-                currentTilesInteractor.removeTiles(listOf(largeTile))
-                runCurrent()
-
-                // Assert that it resized to small
-                assertThat(latest).doesNotContain(largeTile)
-            }
-        }
-
-    @OptIn(ExperimentalCoroutinesApi::class)
-    @Test
     fun resizingNonCurrentTile_doesNothing() =
         with(kosmos) {
             testScope.runTest {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractorTest.kt
new file mode 100644
index 0000000..ee7a15e
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractorTest.kt
@@ -0,0 +1,82 @@
+/*
+ * 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.domain.interactor
+
+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.pipeline.data.repository.FakeDefaultTilesRepository
+import com.android.systemui.qs.pipeline.data.repository.fakeDefaultTilesRepository
+import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
+import com.android.systemui.qs.pipeline.shared.TileSpec
+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
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class SizedTilesResetInteractorTest : SysuiTestCase() {
+    private val kosmos =
+        testKosmos().apply {
+            defaultLargeTilesRepository =
+                object : DefaultLargeTilesRepository {
+                    override val defaultLargeTiles: Set<TileSpec> = setOf(largeTile)
+                }
+            fakeDefaultTilesRepository = FakeDefaultTilesRepository(listOf(smallTile, largeTile))
+        }
+    private val underTest = with(kosmos) { sizedTilesResetInteractor }
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    @Test
+    fun changeTiles_resetsCorrectly() {
+        with(kosmos) {
+            testScope.runTest {
+                // Change current tiles and large tiles
+                currentTilesInteractor.setTiles(listOf(largeTile, newTile))
+                iconTilesInteractor.setLargeTiles(setOf(newTile))
+                runCurrent()
+
+                // Assert both current tiles and large tiles changed
+                assertThat(currentTilesInteractor.currentTilesSpecs)
+                    .containsExactly(largeTile, newTile)
+                assertThat(iconTilesInteractor.largeTilesSpecs.value).containsExactly(newTile)
+
+                // Reset to default
+                underTest.reset()
+                runCurrent()
+
+                // Assert both current tiles and large tiles are back to the initial state
+                assertThat(currentTilesInteractor.currentTilesSpecs)
+                    .containsExactly(largeTile, smallTile)
+                assertThat(iconTilesInteractor.largeTilesSpecs.value).containsExactly(largeTile)
+            }
+        }
+    }
+
+    private companion object {
+        private val largeTile = TileSpec.create("large")
+        private val smallTile = TileSpec.create("small")
+        private val newTile = TileSpec.create("newTile")
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/TileSpecSettingsRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/TileSpecSettingsRepositoryTest.kt
index 7ebebd7..23056b2 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/TileSpecSettingsRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/TileSpecSettingsRepositoryTest.kt
@@ -91,7 +91,7 @@
                 context.resources,
                 logger,
                 retailModeRepository,
-                userTileSpecRepositoryFactory
+                userTileSpecRepositoryFactory,
             )
     }
 
@@ -218,6 +218,21 @@
             assertThat(loadTilesForUser(user)).isEqualTo(startingTiles)
         }
 
+    @Test
+    fun resetsDefault() =
+        testScope.runTest {
+            val tiles by collectLastValue(underTest.tilesSpecs(0))
+
+            val startingTiles = listOf(TileSpec.create("e"), TileSpec.create("f"))
+
+            underTest.setTiles(0, startingTiles)
+            runCurrent()
+
+            underTest.resetToDefault(0)
+
+            assertThat(tiles!!).containsExactlyElementsIn(DEFAULT_TILES.toTileSpecs())
+        }
+
     private fun TestScope.storeTilesForUser(specs: String, forUser: Int) {
         secureSettings.putStringForUser(SETTING, specs, forUser)
         runCurrent()
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 1d80826..de3dc57 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
@@ -28,7 +28,6 @@
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.SysuiTestableContext
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
@@ -145,13 +144,13 @@
 
             // Tile starts with the generic Modes icon.
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             // Add an inactive mode -> Still modes icon
             zenModeRepository.addMode(id = "Mode", active = false)
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             // Add an active mode with a default icon: icon should be the mode icon, and the
@@ -159,7 +158,7 @@
             zenModeRepository.addMode(
                 id = "Bedtime with default icon",
                 type = AutomaticZenRule.TYPE_BEDTIME,
-                active = true,
+                active = true
             )
             runCurrent()
             assertThat(tileData?.icon).isEqualTo(BEDTIME_ICON)
@@ -190,7 +189,7 @@
             // Deactivate remaining mode: back to the default modes icon
             zenModeRepository.deactivateMode("Driving with custom icon")
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
         }
 
@@ -205,18 +204,18 @@
                 )
 
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             // Activate a Mode -> Icon doesn't change.
             zenModeRepository.addMode(id = "Mode", active = true)
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             zenModeRepository.deactivateMode(id = "Mode")
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
         }
 
@@ -264,7 +263,7 @@
         val BEDTIME_DRAWABLE = TestStubDrawable("bedtime")
         val CUSTOM_DRAWABLE = TestStubDrawable("custom")
 
-        val MODES_RESOURCE_ICON = Icon.Resource(MODES_DRAWABLE_ID, null)
+        val MODES_ICON = MODES_DRAWABLE.asIcon()
         val BEDTIME_ICON = BEDTIME_DRAWABLE.asIcon()
         val CUSTOM_ICON = CUSTOM_DRAWABLE.asIcon()
     }
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 a58cb9c..c3d45db 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
@@ -22,9 +22,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.common.shared.model.asIcon
-import com.android.systemui.qs.tiles.ModesTile
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
 import com.android.systemui.qs.tiles.viewmodel.QSTileConfigTestBuilder
 import com.android.systemui.qs.tiles.viewmodel.QSTileState
@@ -53,11 +51,6 @@
                 .apply {
                     addOverride(R.drawable.qs_dnd_icon_on, TestStubDrawable())
                     addOverride(R.drawable.qs_dnd_icon_off, TestStubDrawable())
-                    addOverride(
-                        ModesTile.ICON_RES_ID,
-                        TestStubDrawable(ModesTile.ICON_RES_ID.toString()),
-                    )
-                    addOverride(123, TestStubDrawable("123"))
                 }
                 .resources,
             context.theme,
@@ -66,7 +59,12 @@
     @Test
     fun inactiveState() {
         val icon = TestStubDrawable("res123").asIcon()
-        val model = ModesTileModel(isActivated = false, activeModes = emptyList(), icon = icon)
+        val model =
+            ModesTileModel(
+                isActivated = false,
+                activeModes = emptyList(),
+                icon = icon,
+            )
 
         val state = underTest.map(config, model)
 
@@ -78,7 +76,12 @@
     @Test
     fun activeState_oneMode() {
         val icon = TestStubDrawable("res123").asIcon()
-        val model = ModesTileModel(isActivated = true, activeModes = listOf("DND"), icon = icon)
+        val model =
+            ModesTileModel(
+                isActivated = true,
+                activeModes = listOf("DND"),
+                icon = icon,
+            )
 
         val state = underTest.map(config, model)
 
@@ -105,36 +108,19 @@
     }
 
     @Test
-    fun resourceIconModel_whenResIdsIdentical_mapsToLoadedIconWithInputResId() {
-        val icon = Icon.Resource(123, null)
+    fun state_modelHasIconResId_includesIconResId() {
+        val icon = TestStubDrawable("res123").asIcon()
         val model =
             ModesTileModel(
                 isActivated = false,
                 activeModes = emptyList(),
                 icon = icon,
-                iconResId = 123,
+                iconResId = 123
             )
 
         val state = underTest.map(config, model)
 
-        assertThat(state.icon()).isEqualTo(TestStubDrawable("123").asIcon())
-        assertThat(state.iconRes).isEqualTo(123)
-    }
-
-    @Test
-    fun resourceIconModel_whenResIdsNonIdentical_mapsToLoadedIconWithIconResourceId() {
-        val icon = Icon.Resource(123, null)
-        val model =
-            ModesTileModel(
-                isActivated = false,
-                activeModes = emptyList(),
-                icon = icon,
-                iconResId = 321, // Note: NOT 123. This will be ignored.
-            )
-
-        val state = underTest.map(config, model)
-
-        assertThat(state.icon()).isEqualTo(TestStubDrawable("123").asIcon())
+        assertThat(state.icon()).isEqualTo(icon)
         assertThat(state.iconRes).isEqualTo(123)
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
index 254f1e1..4d71dc4 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
@@ -21,8 +21,8 @@
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREE_FORM
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.PIP
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_BOTTOM
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_TOP
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.HORIZONTAL
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.VERTICAL
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.emptyRootSplit
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.freeForm
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.fullScreen
@@ -39,16 +39,14 @@
 
     data class TaskSpec(val taskId: Int, val userId: Int, val name: String)
 
+    val emptyDisplayContent = DisplayContentModel(0, SystemUiState(shadeExpanded = false), listOf())
+
     /** Home screen, with only the launcher visible */
     fun launcherOnly(shadeExpanded: Boolean = false) =
         DisplayContentModel(
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
-            rootTasks =
-                listOf(
-                    launcher(visible = true),
-                    emptyRootSplit,
-                )
+            rootTasks = listOf(launcher(visible = true), emptyRootSplit),
         )
 
     /** A Full screen activity for the personal (primary) user, with launcher behind it */
@@ -57,48 +55,72 @@
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
             rootTasks =
-                listOf(
-                    fullScreen(spec, visible = true),
-                    launcher(visible = false),
-                    emptyRootSplit,
-                )
+                listOf(fullScreen(spec, visible = true), launcher(visible = false), emptyRootSplit),
         )
 
+    enum class Orientation {
+        HORIZONTAL,
+        VERTICAL,
+    }
+
+    internal fun Rect.splitLeft(margin: Int = 0) = Rect(left, top, centerX() - margin, bottom)
+
+    internal fun Rect.splitRight(margin: Int = 0) = Rect(centerX() + margin, top, right, bottom)
+
+    internal fun Rect.splitTop(margin: Int = 0) = Rect(left, top, right, centerY() - margin)
+
+    internal fun Rect.splitBottom(margin: Int = 0) = Rect(left, centerY() + margin, right, bottom)
+
     fun splitScreenApps(
-        top: TaskSpec,
-        bottom: TaskSpec,
+        displayId: Int = 0,
+        parentBounds: Rect = FULL_SCREEN,
+        taskMargin: Int = 0,
+        orientation: Orientation = VERTICAL,
+        first: TaskSpec,
+        second: TaskSpec,
         focusedTaskId: Int,
+        parentTaskId: Int = 2,
         shadeExpanded: Boolean = false,
     ): DisplayContentModel {
-        val topBounds = SPLIT_TOP
-        val bottomBounds = SPLIT_BOTTOM
+
+        val firstBounds =
+            when (orientation) {
+                VERTICAL -> parentBounds.splitTop(taskMargin)
+                HORIZONTAL -> parentBounds.splitLeft(taskMargin)
+            }
+        val secondBounds =
+            when (orientation) {
+                VERTICAL -> parentBounds.splitBottom(taskMargin)
+                HORIZONTAL -> parentBounds.splitRight(taskMargin)
+            }
+
         return DisplayContentModel(
-            displayId = 0,
+            displayId = displayId,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
             rootTasks =
                 listOf(
                     newRootTaskInfo(
-                        taskId = 2,
+                        taskId = parentTaskId,
                         userId = TestUserIds.PERSONAL,
-                        bounds = FULL_SCREEN,
+                        bounds = parentBounds,
                         topActivity =
                             ComponentName.unflattenFromString(
-                                if (top.taskId == focusedTaskId) top.name else bottom.name
+                                if (first.taskId == focusedTaskId) first.name else second.name
                             ),
                     ) {
                         listOf(
                                 newChildTask(
-                                    taskId = top.taskId,
-                                    bounds = topBounds,
-                                    userId = top.userId,
-                                    name = top.name
+                                    taskId = first.taskId,
+                                    bounds = firstBounds,
+                                    userId = first.userId,
+                                    name = first.name,
                                 ),
                                 newChildTask(
-                                    taskId = bottom.taskId,
-                                    bounds = bottomBounds,
-                                    userId = bottom.userId,
-                                    name = bottom.name
-                                )
+                                    taskId = second.taskId,
+                                    bounds = secondBounds,
+                                    userId = second.userId,
+                                    name = second.name,
+                                ),
                             )
                             // Child tasks are ordered bottom-up in RootTaskInfo.
                             // Sort 'focusedTaskId' last.
@@ -106,7 +128,7 @@
                             .sortedBy { it.id == focusedTaskId }
                     },
                     launcher(visible = false),
-                )
+                ),
         )
     }
 
@@ -124,7 +146,7 @@
                     fullScreen?.also { add(fullScreen(it, visible = true)) }
                     add(launcher(visible = (fullScreen == null)))
                     add(emptyRootSplit)
-                }
+                },
         )
     }
 
@@ -142,7 +164,7 @@
         return DisplayContentModel(
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
-            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit
+            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit,
         )
     }
 
@@ -153,11 +175,18 @@
      * somewhat sensible in terms of logical position (Re: PIP, SPLIT, etc).
      */
     object Bounds {
+        // "Phone" size
         val FULL_SCREEN = Rect(0, 0, 1080, 2400)
         val PIP = Rect(440, 1458, 1038, 1794)
         val SPLIT_TOP = Rect(0, 0, 1080, 1187)
         val SPLIT_BOTTOM = Rect(0, 1213, 1080, 2400)
         val FREE_FORM = Rect(119, 332, 1000, 1367)
+
+        // "Tablet" size
+        val FREEFORM_FULL_SCREEN = Rect(0, 0, 2560, 1600)
+        val FREEFORM_MAXIMIZED = Rect(0, 48, 2560, 1480)
+        val FREEFORM_SPLIT_LEFT = Rect(0, 0, 1270, 1600)
+        val FREEFORM_SPLIT_RIGHT = Rect(1290, 0, 2560, 1600)
     }
 
     /** A collection of task names used in test scenarios */
@@ -177,6 +206,8 @@
             "com.google.android.youtube/" +
                 "com.google.android.apps.youtube.app.watchwhile.WatchWhileActivity"
 
+        const val MESSAGES = "com.google.android.apps.messaging/.ui.ConversationListActivity"
+
         /** The NexusLauncher activity */
         const val LAUNCHER =
             "com.google.android.apps.nexuslauncher/" +
@@ -220,7 +251,7 @@
             }
 
         /** NexusLauncher on the default display. Usually below all other visible tasks */
-        fun launcher(visible: Boolean) =
+        fun launcher(visible: Boolean, bounds: Rect = FULL_SCREEN) =
             newRootTaskInfo(
                 taskId = 1,
                 activityType = ActivityType.Home,
@@ -229,43 +260,63 @@
                 topActivity = ComponentName.unflattenFromString(ActivityNames.LAUNCHER),
                 topActivityType = ActivityType.Home,
             ) {
-                listOf(newChildTask(taskId = 1002, name = ActivityNames.LAUNCHER))
+                listOf(newChildTask(taskId = 1002, name = ActivityNames.LAUNCHER, bounds = bounds))
             }
 
         /** A full screen Activity */
-        fun fullScreen(task: TaskSpec, visible: Boolean) =
+        fun fullScreen(task: TaskSpec, visible: Boolean, bounds: Rect = FULL_SCREEN) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
                 visible = visible,
-                bounds = FULL_SCREEN,
+                bounds = bounds,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = task.userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = task.userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
 
         /** An activity in Picture-in-Picture mode */
-        fun pictureInPicture(task: TaskSpec) =
+        fun pictureInPicture(task: TaskSpec, bounds: Rect = PIP) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
-                bounds = PIP,
                 windowingMode = WindowingMode.PictureInPicture,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
 
         /** An activity in FreeForm mode */
-        fun freeForm(task: TaskSpec) =
+        fun freeForm(task: TaskSpec, bounds: Rect = FREE_FORM) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
-                bounds = FREE_FORM,
+                bounds = bounds,
                 windowingMode = WindowingMode.Freeform,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
index 6c35b23..cedf0c8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
@@ -69,7 +69,7 @@
     taskId: Int,
     name: String,
     bounds: Rect? = null,
-    userId: Int? = null
+    userId: Int? = null,
 ): ChildTaskModel {
     return ChildTaskModel(taskId, name, bounds ?: this.bounds, userId ?: this.userId)
 }
@@ -83,7 +83,7 @@
     running: Boolean = true,
     activityType: ActivityType = Standard,
     windowingMode: WindowingMode = FullScreen,
-    bounds: Rect? = null,
+    bounds: Rect = Rect(),
     topActivity: ComponentName? = null,
     topActivityType: ActivityType = Standard,
     numActivities: Int? = null,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
index 6e57761..b7f565d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
@@ -17,8 +17,8 @@
 package com.android.systemui.screenshot.policy
 
 import android.content.ComponentName
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.screenshot.data.model.DisplayContentModel
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
@@ -59,7 +59,7 @@
             policy.check(
                 singleFullScreen(
                     spec = TaskSpec(taskId = 1002, name = YOUTUBE, userId = PRIVATE),
-                    shadeExpanded = true
+                    shadeExpanded = true,
                 )
             )
 
@@ -93,8 +93,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -110,25 +110,20 @@
                         listOf(
                             fullScreen(
                                 TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                                visible = true
+                                visible = true,
                             ),
                             fullScreen(
                                 TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                                visible = false
+                                visible = false,
                             ),
                             launcher(visible = false),
                             emptyRootSplit,
-                        )
+                        ),
                 )
             )
 
         assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    PrivateProfilePolicy.NAME,
-                    PrivateProfilePolicy.NO_VISIBLE_TASKS,
-                )
-            )
+            .isEqualTo(NotMatched(PrivateProfilePolicy.NAME, PrivateProfilePolicy.NO_VISIBLE_TASKS))
     }
 
     @Test
@@ -136,9 +131,9 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                    focusedTaskId = 1003
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1003,
                 )
             )
 
@@ -150,8 +145,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -161,9 +156,9 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                    focusedTaskId = 1002
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
                 )
             )
 
@@ -175,8 +170,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(FILES),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -196,8 +191,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE_PIP),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -220,8 +215,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE_PIP),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt
new file mode 100644
index 0000000..28eb9fc
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt
@@ -0,0 +1,351 @@
+/*
+ * 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.policy
+
+import android.content.ComponentName
+import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.LAUNCHER
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.MESSAGES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.YOUTUBE
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREEFORM_FULL_SCREEN
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.VERTICAL
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.freeFormApps
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.pictureInPictureApp
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitScreenApps
+import com.android.systemui.screenshot.data.repository.profileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
+import com.android.systemui.screenshot.policy.CaptureType.RootTask
+import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
+import com.android.systemui.screenshot.policy.TestUserIds.PRIVATE
+import com.android.systemui.screenshot.policy.TestUserIds.WORK
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class ScreenshotPolicyTest {
+    private val kosmos = Kosmos()
+
+    private val defaultComponent = ComponentName("default", "default")
+    private val defaultOwner = UserHandle.SYSTEM
+
+    @Test
+    fun fullScreen_work() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(TaskSpec(taskId = 1002, name = FILES, userId = WORK)),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_private() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(TaskSpec(taskId = 1002, name = YOUTUBE, userId = PRIVATE)),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_workAndPersonal() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_personalAndPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_workAndPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_twoWorkTasks() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    parentTaskId = 1,
+                    parentBounds = FREEFORM_FULL_SCREEN,
+                    orientation = VERTICAL,
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = WORK),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type =
+                        RootTask(
+                            parentTaskId = 1,
+                            taskBounds = FREEFORM_FULL_SCREEN,
+                            childTaskIds = listOf(1002, 1003),
+                        ),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows_maximized() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows_withPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    TaskSpec(taskId = 1004, name = MESSAGES, userId = PERSONAL),
+                    focusedTaskId = 1004,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floating_workOnly() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(LAUNCHER),
+                    owner = defaultOwner,
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_shadeExpanded() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    shadeExpanded = true,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = defaultComponent,
+                    owner = defaultOwner,
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_with_PictureInPicture() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                pictureInPictureApp(
+                    pip = TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
+                    fullScreen = TaskSpec(taskId = 1003, name = FILES, userId = WORK),
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = IsolatedTask(taskId = 1003, taskBounds = FULL_SCREEN),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
index be9fcc2..30a786c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
@@ -31,13 +31,13 @@
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.YOUTUBE
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREE_FORM
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_TOP
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.freeFormApps
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.pictureInPictureApp
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitScreenApps
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitTop
 import com.android.systemui.screenshot.data.model.SystemUiState
 import com.android.systemui.screenshot.data.repository.profileTypeRepository
 import com.android.systemui.screenshot.policy.CapturePolicy.PolicyResult
@@ -69,6 +69,7 @@
     @JvmField @Rule(order = 2) val mockitoRule: MockitoRule = MockitoJUnit.rule()
 
     @Mock lateinit var mContext: Context
+
     @Mock lateinit var mResources: Resources
 
     private val kosmos = Kosmos()
@@ -94,17 +95,11 @@
                 DisplayContentModel(
                     displayId = 0,
                     systemUiState = SystemUiState(shadeExpanded = false),
-                    rootTasks = listOf(RootTasks.emptyWithNoChildTasks)
+                    rootTasks = listOf(RootTasks.emptyWithNoChildTasks),
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -114,13 +109,7 @@
                 singleFullScreen(TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL))
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -129,17 +118,11 @@
             policy.check(
                 singleFullScreen(
                     TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    shadeExpanded = true
+                    shadeExpanded = true,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    SHADE_EXPANDED,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, SHADE_EXPANDED))
     }
 
     @Test
@@ -156,7 +139,7 @@
                         type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -166,9 +149,11 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
-                    focusedTaskId = 1002
+                    parentBounds = FULL_SCREEN,
+                    taskMargin = 20,
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1002,
                 )
             )
 
@@ -178,10 +163,10 @@
                     policy = WorkProfilePolicy.NAME,
                     reason = WORK_TASK_IS_TOP,
                     CaptureParameters(
-                        type = IsolatedTask(taskId = 1002, taskBounds = SPLIT_TOP),
+                        type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN.splitTop(20)),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -191,19 +176,13 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
-                    focusedTaskId = 1003
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -225,7 +204,7 @@
                         type = IsolatedTask(taskId = 1003, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -238,7 +217,7 @@
                 freeFormApps(
                     TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
                     TaskSpec(taskId = 1003, name = FILES, userId = WORK),
-                    focusedTaskId = 1003
+                    focusedTaskId = 1003,
                 )
             )
 
@@ -251,7 +230,7 @@
                         type = IsolatedTask(taskId = 1003, taskBounds = FREE_FORM),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -264,16 +243,10 @@
                 freeFormApps(
                     TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
                     TaskSpec(taskId = 1003, name = FILES, userId = WORK),
-                    focusedTaskId = 1003
+                    focusedTaskId = 1003,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    DESKTOP_MODE_ENABLED,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, DESKTOP_MODE_ENABLED))
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt
index a629b24..5f3668a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt
@@ -30,7 +30,6 @@
 import com.android.systemui.statusbar.phone.StatusBarContentInsetsChangedListener
 import com.android.systemui.statusbar.phone.StatusBarContentInsetsProvider
 import com.android.systemui.statusbar.window.StatusBarWindowController
-import com.android.systemui.statusbar.window.StatusBarWindowControllerStore
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.capture
@@ -50,11 +49,10 @@
 @RunWith(AndroidJUnit4::class)
 @TestableLooper.RunWithLooper
 class SystemEventChipAnimationControllerTest : SysuiTestCase() {
-    private lateinit var controller: SystemEventChipAnimationController
+    private lateinit var controller: SystemEventChipAnimationControllerImpl
 
     @get:Rule val animatorTestRule = AnimatorTestRule(this)
     @Mock private lateinit var sbWindowController: StatusBarWindowController
-    @Mock private lateinit var sbWindowControllerStore: StatusBarWindowControllerStore
     @Mock private lateinit var insetsProvider: StatusBarContentInsetsProvider
 
     private var testView = TestView(mContext)
@@ -63,7 +61,6 @@
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        whenever(sbWindowControllerStore.defaultDisplay).thenReturn(sbWindowController)
         // StatusBarWindowController is mocked. The addViewToWindow function needs to be mocked to
         // ensure that the chip view is added to a parent view
         whenever(sbWindowController.addViewToWindow(any(), any())).then {
@@ -76,7 +73,7 @@
             )
             statusbarFake.addView(
                 it.arguments[0] as View,
-                it.arguments[1] as FrameLayout.LayoutParams
+                it.arguments[1] as FrameLayout.LayoutParams,
             )
         }
 
@@ -86,16 +83,16 @@
                     /* left= */ insets,
                     /* top= */ insets,
                     /* right= */ insets,
-                    /* bottom= */ 0
+                    /* bottom= */ 0,
                 )
             )
         whenever(insetsProvider.getStatusBarContentAreaForCurrentRotation())
             .thenReturn(portraitArea)
 
         controller =
-            SystemEventChipAnimationController(
+            SystemEventChipAnimationControllerImpl(
                 context = mContext,
-                statusBarWindowControllerStore = sbWindowControllerStore,
+                statusBarWindowController = sbWindowController,
                 contentInsetsProvider = insetsProvider,
             )
     }
@@ -156,10 +153,7 @@
             val lp = it.arguments[1] as FrameLayout.LayoutParams
             assertThat(lp.gravity and Gravity.VERTICAL_GRAVITY_MASK).isEqualTo(Gravity.TOP)
 
-            statusbarFake.addView(
-                it.arguments[0] as View,
-                lp,
-            )
+            statusbarFake.addView(it.arguments[0] as View, lp)
         }
 
         // GIVEN insets provider gives the correct content area
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
index 97fa6eb..75479ad 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationScrimNestedScrollConnectionTest.kt
@@ -17,7 +17,7 @@
 package com.android.systemui.statusbar.notification
 
 import androidx.compose.ui.geometry.Offset
-import androidx.compose.ui.input.nestedscroll.NestedScrollSource
+import androidx.compose.ui.input.nestedscroll.NestedScrollSource.Companion.UserInput
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -58,10 +58,7 @@
         contentHeight = COLLAPSED_CONTENT_HEIGHT
 
         val offsetConsumed =
-            scrollConnection.onPreScroll(
-                available = Offset(x = 0f, y = -1f),
-                source = NestedScrollSource.Drag,
-            )
+            scrollConnection.onPreScroll(available = Offset(x = 0f, y = -1f), source = UserInput)
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
         assertThat(isStarted).isEqualTo(false)
@@ -73,10 +70,7 @@
         scrimOffset = MIN_SCRIM_OFFSET
 
         val offsetConsumed =
-            scrollConnection.onPreScroll(
-                available = Offset(x = 0f, y = -1f),
-                source = NestedScrollSource.Drag,
-            )
+            scrollConnection.onPreScroll(available = Offset(x = 0f, y = -1f), source = UserInput)
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
         assertThat(isStarted).isEqualTo(false)
@@ -88,10 +82,7 @@
 
         val availableOffset = Offset(x = 0f, y = -1f)
         val offsetConsumed =
-            scrollConnection.onPreScroll(
-                available = availableOffset,
-                source = NestedScrollSource.Drag,
-            )
+            scrollConnection.onPreScroll(available = availableOffset, source = UserInput)
 
         assertThat(offsetConsumed).isEqualTo(availableOffset)
         assertThat(isStarted).isEqualTo(true)
@@ -105,10 +96,7 @@
         val availableOffset = Offset(x = 0f, y = -2f)
         val consumableOffset = Offset(x = 0f, y = -1f)
         val offsetConsumed =
-            scrollConnection.onPreScroll(
-                available = availableOffset,
-                source = NestedScrollSource.Drag,
-            )
+            scrollConnection.onPreScroll(available = availableOffset, source = UserInput)
 
         assertThat(offsetConsumed).isEqualTo(consumableOffset)
         assertThat(isStarted).isEqualTo(true)
@@ -120,7 +108,7 @@
             scrollConnection.onPostScroll(
                 consumed = Offset.Zero,
                 available = Offset(x = 0f, y = -1f),
-                source = NestedScrollSource.Drag,
+                source = UserInput,
             )
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
@@ -130,10 +118,7 @@
     @Test
     fun onScrollDown_canStartPreScroll_ignoreScroll() = runTest {
         val offsetConsumed =
-            scrollConnection.onPreScroll(
-                available = Offset(x = 0f, y = 1f),
-                source = NestedScrollSource.Drag,
-            )
+            scrollConnection.onPreScroll(available = Offset(x = 0f, y = 1f), source = UserInput)
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
         assertThat(isStarted).isEqualTo(false)
@@ -148,7 +133,7 @@
             scrollConnection.onPostScroll(
                 consumed = Offset.Zero,
                 available = availableOffset,
-                source = NestedScrollSource.Drag
+                source = UserInput,
             )
 
         assertThat(offsetConsumed).isEqualTo(availableOffset)
@@ -165,7 +150,7 @@
             scrollConnection.onPostScroll(
                 consumed = Offset.Zero,
                 available = availableOffset,
-                source = NestedScrollSource.Drag
+                source = UserInput,
             )
 
         assertThat(offsetConsumed).isEqualTo(consumableOffset)
@@ -180,7 +165,7 @@
             scrollConnection.onPostScroll(
                 consumed = Offset.Zero,
                 available = Offset(x = 0f, y = 1f),
-                source = NestedScrollSource.Drag
+                source = UserInput,
             )
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
@@ -197,7 +182,7 @@
             scrollConnection.onPostScroll(
                 consumed = Offset.Zero,
                 available = Offset(x = 0f, y = 1f),
-                source = NestedScrollSource.Drag
+                source = UserInput,
             )
 
         assertThat(offsetConsumed).isEqualTo(Offset.Zero)
@@ -210,17 +195,11 @@
     fun canContinueScroll_inBetweenMinMaxOffset_true() = runTest {
         scrimOffset = (MIN_SCRIM_OFFSET + MAX_SCRIM_OFFSET) / 2f
         contentHeight = EXPANDED_CONTENT_HEIGHT
-        scrollConnection.onPreScroll(
-            available = Offset(x = 0f, y = -1f),
-            source = NestedScrollSource.Drag
-        )
+        scrollConnection.onPreScroll(available = Offset(x = 0f, y = -1f), source = UserInput)
 
         assertThat(isStarted).isEqualTo(true)
 
-        scrollConnection.onPreScroll(
-            available = Offset(x = 0f, y = 1f),
-            source = NestedScrollSource.Drag
-        )
+        scrollConnection.onPreScroll(available = Offset(x = 0f, y = 1f), source = UserInput)
 
         assertThat(isStarted).isEqualTo(true)
     }
@@ -229,17 +208,11 @@
     fun canContinueScroll_atMaxOffset_false() = runTest {
         scrimOffset = MAX_SCRIM_OFFSET
         contentHeight = EXPANDED_CONTENT_HEIGHT
-        scrollConnection.onPreScroll(
-            available = Offset(x = 0f, y = -1f),
-            source = NestedScrollSource.Drag
-        )
+        scrollConnection.onPreScroll(available = Offset(x = 0f, y = -1f), source = UserInput)
 
         assertThat(isStarted).isEqualTo(true)
 
-        scrollConnection.onPreScroll(
-            available = Offset(x = 0f, y = 1f),
-            source = NestedScrollSource.Drag
-        )
+        scrollConnection.onPreScroll(available = Offset(x = 0f, y = 1f), source = UserInput)
 
         assertThat(isStarted).isEqualTo(false)
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCacheTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCacheTest.kt
new file mode 100644
index 0000000..d2a3a19
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCacheTest.kt
@@ -0,0 +1,232 @@
+/*
+ * 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.notification.collection
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.time.FakeSystemClock
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.assertNotNull
+import kotlin.test.assertNull
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class NotifCollectionCacheTest : SysuiTestCase() {
+    companion object {
+        const val A = "a"
+        const val B = "b"
+        const val C = "c"
+    }
+
+    val systemClock = FakeSystemClock()
+    val underTest =
+        NotifCollectionCache<String>(purgeTimeoutMillis = 200L, systemClock = systemClock)
+
+    @After
+    fun cleanUp() {
+        underTest.clear()
+    }
+
+    @Test
+    fun fetch_isOnlyCalledOncePerEntry() {
+        val fetchList = mutableListOf<String>()
+        val fetch = { key: String ->
+            fetchList.add(key)
+            key
+        }
+
+        // Construct the cache and make sure fetch is called
+        assertThat(underTest.getOrFetch(A, fetch)).isEqualTo(A)
+        assertThat(underTest.getOrFetch(B, fetch)).isEqualTo(B)
+        assertThat(underTest.getOrFetch(C, fetch)).isEqualTo(C)
+        assertThat(fetchList).containsExactly(A, B, C).inOrder()
+
+        // Verify that further calls don't trigger fetch again
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+        assertThat(fetchList).containsExactly(A, B, C).inOrder()
+
+        // Verify that fetch gets called again if the entries are cleared
+        underTest.clear()
+        underTest.getOrFetch(A, fetch)
+        assertThat(fetchList).containsExactly(A, B, C, A).inOrder()
+    }
+
+    @Test
+    fun purge_beforeTimeout_doesNothing() {
+        // Populate cache
+        val fetch = { key: String -> key }
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+
+        // B starts off with ♥ ︎♥︎
+        assertThat(underTest.getLives(B)).isEqualTo(2)
+        // First purge run removes a ︎♥︎
+        underTest.purge(listOf(A, C))
+        assertNotNull(underTest.cache[B])
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+        // Second purge run done too early does nothing to B
+        systemClock.advanceTime(100L)
+        underTest.purge(listOf(A, C))
+        assertNotNull(underTest.cache[B])
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+        // Purge done after timeout (200ms) clears B
+        systemClock.advanceTime(100L)
+        underTest.purge(listOf(A, C))
+        assertNull(underTest.cache[B])
+    }
+
+    @Test
+    fun get_resetsLives() {
+        // Populate cache
+        val fetch = { key: String -> key }
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+
+        // Bring B down to one ︎♥︎
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+
+        // Get should restore B to ♥ ︎♥︎
+        underTest.getOrFetch(B, fetch)
+        assertThat(underTest.getLives(B)).isEqualTo(2)
+
+        // Subsequent purge should remove a life regardless of timing
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+    }
+
+    @Test
+    fun purge_resetsLives() {
+        // Populate cache
+        val fetch = { key: String -> key }
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+
+        // Bring B down to one ︎♥︎
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+
+        // When B is back to wantedKeys, it is restored to to ♥ ︎♥ ︎︎
+        underTest.purge(listOf(B))
+        assertThat(underTest.getLives(B)).isEqualTo(2)
+        assertThat(underTest.getLives(A)).isEqualTo(1)
+        assertThat(underTest.getLives(C)).isEqualTo(1)
+
+        // Subsequent purge should remove a life regardless of timing
+        underTest.purge(listOf(A, C))
+        assertThat(underTest.getLives(B)).isEqualTo(1)
+    }
+
+    @Test
+    fun purge_worksWithMoreLives() {
+        val multiLivesCache =
+            NotifCollectionCache<String>(
+                retainCount = 3,
+                purgeTimeoutMillis = 100L,
+                systemClock = systemClock,
+            )
+
+        // Populate cache
+        val fetch = { key: String -> key }
+        multiLivesCache.getOrFetch(A, fetch)
+        multiLivesCache.getOrFetch(B, fetch)
+        multiLivesCache.getOrFetch(C, fetch)
+
+        // B starts off with ♥ ︎♥︎ ♥ ︎♥︎
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(4)
+        // First purge run removes a ︎♥︎
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(3)
+        // Second purge run done too early does nothing to B
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(3)
+        // Staggered purge runs remove further ︎♥︎
+        systemClock.advanceTime(100L)
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(2)
+        systemClock.advanceTime(100L)
+        multiLivesCache.purge(listOf(A, C))
+        assertNotNull(multiLivesCache.cache[B])
+        assertThat(multiLivesCache.getLives(B)).isEqualTo(1)
+        systemClock.advanceTime(100L)
+        multiLivesCache.purge(listOf(A, C))
+        assertNull(multiLivesCache.cache[B])
+    }
+
+    @Test
+    fun purge_worksWithNoLives() {
+        val noLivesCache =
+            NotifCollectionCache<String>(
+                retainCount = 0,
+                purgeTimeoutMillis = 0L,
+                systemClock = systemClock,
+            )
+
+        val fetch = { key: String -> key }
+        noLivesCache.getOrFetch(A, fetch)
+        noLivesCache.getOrFetch(B, fetch)
+        noLivesCache.getOrFetch(C, fetch)
+
+        // Purge immediately removes entry
+        noLivesCache.purge(listOf(A, C))
+
+        assertNotNull(noLivesCache.cache[A])
+        assertNull(noLivesCache.cache[B])
+        assertNotNull(noLivesCache.cache[C])
+    }
+
+    @Test
+    fun hitsAndMisses_areAccurate() {
+        val fetch = { key: String -> key }
+
+        // Construct the cache
+        assertThat(underTest.getOrFetch(A, fetch)).isEqualTo(A)
+        assertThat(underTest.getOrFetch(B, fetch)).isEqualTo(B)
+        assertThat(underTest.getOrFetch(C, fetch)).isEqualTo(C)
+        assertThat(underTest.hits.get()).isEqualTo(0)
+        assertThat(underTest.misses.get()).isEqualTo(3)
+
+        // Verify that further calls count as hits
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(A, fetch)
+        underTest.getOrFetch(B, fetch)
+        underTest.getOrFetch(C, fetch)
+        assertThat(underTest.hits.get()).isEqualTo(4)
+        assertThat(underTest.misses.get()).isEqualTo(3)
+
+        // Verify that a miss is counted again if the entries are cleared
+        underTest.clear()
+        underTest.getOrFetch(A, fetch)
+        assertThat(underTest.hits.get()).isEqualTo(4)
+        assertThat(underTest.misses.get()).isEqualTo(4)
+    }
+
+    private fun <V> NotifCollectionCache<V>.getLives(key: String) = this.cache[key]?.lives
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt
index 9d990b1..9a6a699 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapperTest.kt
@@ -26,6 +26,7 @@
 import com.android.internal.widget.MessagingGroup
 import com.android.internal.widget.MessagingImageMessage
 import com.android.internal.widget.MessagingLinearLayout
+import com.android.internal.widget.NotificationRowIconView
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper
@@ -90,7 +91,7 @@
 
     private fun fakeConversationLayout(
         mockDrawableGroupMessage: AnimatedImageDrawable,
-        mockDrawableImageMessage: AnimatedImageDrawable
+        mockDrawableImageMessage: AnimatedImageDrawable,
     ): View {
         val mockMessagingImageMessage: MessagingImageMessage =
             mock<MessagingImageMessage>().apply {
@@ -126,6 +127,7 @@
                 whenever(requireViewById<CachingIconView>(R.id.conversation_icon))
                     .thenReturn(mock())
                 whenever(findViewById<CachingIconView>(R.id.icon)).thenReturn(mock())
+                whenever(requireViewById<NotificationRowIconView>(R.id.icon)).thenReturn(mock())
                 whenever(requireViewById<View>(R.id.conversation_icon_badge_bg)).thenReturn(mock())
                 whenever(requireViewById<View>(R.id.expand_button)).thenReturn(mock())
                 whenever(requireViewById<View>(R.id.expand_button_container)).thenReturn(mock())
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt
index 286f017..dbe90a5 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt
@@ -20,6 +20,7 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.surfaceeffects.turbulencenoise.TurbulenceNoiseShader.Companion.Type.SIMPLEX_NOISE
 import com.android.systemui.surfaceeffects.turbulencenoise.TurbulenceNoiseShader.Companion.Type.SIMPLEX_NOISE_FRACTAL
+import com.android.systemui.surfaceeffects.turbulencenoise.TurbulenceNoiseShader.Companion.Type.SIMPLEX_NOISE_SIMPLE
 import com.android.systemui.surfaceeffects.turbulencenoise.TurbulenceNoiseShader.Companion.Type.SIMPLEX_NOISE_SPARKLE
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -28,20 +29,23 @@
 @RunWith(AndroidJUnit4::class)
 class TurbulenceNoiseShaderTest : SysuiTestCase() {
 
-    private lateinit var turbulenceNoiseShader: TurbulenceNoiseShader
-
     @Test
     fun compilesSimplexNoise() {
-        turbulenceNoiseShader = TurbulenceNoiseShader(baseType = SIMPLEX_NOISE)
+        TurbulenceNoiseShader(baseType = SIMPLEX_NOISE)
+    }
+
+    @Test
+    fun compilesSimplexSimpleNoise() {
+        TurbulenceNoiseShader(baseType = SIMPLEX_NOISE_SIMPLE)
     }
 
     @Test
     fun compilesFractalNoise() {
-        turbulenceNoiseShader = TurbulenceNoiseShader(baseType = SIMPLEX_NOISE_FRACTAL)
+        TurbulenceNoiseShader(baseType = SIMPLEX_NOISE_FRACTAL)
     }
 
     @Test
     fun compilesSparkleNoise() {
-        turbulenceNoiseShader = TurbulenceNoiseShader(baseType = SIMPLEX_NOISE_SPARKLE)
+        TurbulenceNoiseShader(baseType = SIMPLEX_NOISE_SPARKLE)
     }
 }
diff --git a/packages/SystemUI/res/anim/shortcut_helper_close_anim.xml b/packages/SystemUI/res/anim/shortcut_helper_close_anim.xml
deleted file mode 100644
index 47fd78a..0000000
--- a/packages/SystemUI/res/anim/shortcut_helper_close_anim.xml
+++ /dev/null
@@ -1,26 +0,0 @@
-<?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.
-  -->
-
-<set xmlns:android="http://schemas.android.com/apk/res/android"
-    android:interpolator="@android:anim/accelerate_interpolator"
-    android:zAdjustment="top">
-
-    <translate
-        android:fromYDelta="0"
-        android:toYDelta="100%"
-        android:duration="@android:integer/config_shortAnimTime" />
-</set>
diff --git a/packages/SystemUI/res/anim/shortcut_helper_launch_anim.xml b/packages/SystemUI/res/anim/shortcut_helper_launch_anim.xml
deleted file mode 100644
index 77edf58..0000000
--- a/packages/SystemUI/res/anim/shortcut_helper_launch_anim.xml
+++ /dev/null
@@ -1,27 +0,0 @@
-<?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.
-  -->
-
-<!-- Animation for when a dock window at the bottom of the screen is entering. -->
-<set xmlns:android="http://schemas.android.com/apk/res/android"
-    android:interpolator="@android:anim/accelerate_decelerate_interpolator"
-    android:zAdjustment="top">
-
-    <translate android:fromYDelta="100%"
-        android:toYDelta="0"
-        android:startOffset="@android:integer/config_shortAnimTime"
-        android:duration="@android:integer/config_mediumAnimTime"/>
-</set>
diff --git a/packages/SystemUI/res/values-night/styles.xml b/packages/SystemUI/res/values-night/styles.xml
index 17ba2e5..d5d52e3 100644
--- a/packages/SystemUI/res/values-night/styles.xml
+++ b/packages/SystemUI/res/values-night/styles.xml
@@ -59,8 +59,4 @@
         -->
         <item name="android:windowOptOutEdgeToEdgeEnforcement">true</item>
     </style>
-
-    <style name="ShortcutHelperTheme" parent="@style/ShortcutHelperThemeCommon">
-        <item name="android:windowLightNavigationBar">false</item>
-    </style>
 </resources>
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 2c5fb56..cdf15ca 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -3895,4 +3895,12 @@
     <string name="qs_edit_mode_category_unknown">
         Unknown
     </string>
+    <!-- Title for the Reset Tiles dialog in QS Edit mode. [CHAR LIMIT=NONE] -->
+    <string name="qs_edit_mode_reset_dialog_title">
+        Reset tiles
+    </string>
+    <!-- Content of the Reset Tiles dialog in QS Edit mode. [CHAR LIMIT=NONE] -->
+    <string name="qs_edit_mode_reset_dialog_content">
+        Reset tiles to their original order and sizes?
+    </string>
 </resources>
diff --git a/packages/SystemUI/res/values/styles.xml b/packages/SystemUI/res/values/styles.xml
index 7cebac2..bda3453 100644
--- a/packages/SystemUI/res/values/styles.xml
+++ b/packages/SystemUI/res/values/styles.xml
@@ -1709,32 +1709,4 @@
     <style name="Theme.SystemUI.Dialog.StickyKeys" parent="@style/Theme.SystemUI.Dialog">
         <item name="android:colorBackground">@color/transparent</item>
     </style>
-
-    <style name="ShortcutHelperBottomSheet" parent="@style/Widget.Material3.BottomSheet">
-        <item name="backgroundTint">?colorSurfaceContainer</item>
-    </style>
-
-    <style name="ShortcutHelperAnimation" parent="@android:style/Animation.Activity">
-        <item name="android:activityOpenEnterAnimation">@anim/shortcut_helper_launch_anim</item>
-        <item name="android:taskOpenEnterAnimation">@anim/shortcut_helper_launch_anim</item>
-        <item name="android:activityOpenExitAnimation">@anim/shortcut_helper_close_anim</item>
-        <item name="android:taskOpenExitAnimation">@anim/shortcut_helper_close_anim</item>
-    </style>
-
-    <style name="ShortcutHelperThemeCommon" parent="@style/Theme.Material3.DynamicColors.DayNight">
-        <item name="android:windowAnimationStyle">@style/ShortcutHelperAnimation</item>
-        <item name="android:windowIsTranslucent">true</item>
-        <item name="android:windowNoTitle">true</item>
-        <item name="android:windowBackground">@android:color/transparent</item>
-        <item name="android:backgroundDimEnabled">true</item>
-        <item name="android:statusBarColor">@android:color/transparent</item>
-        <item name="android:windowContentOverlay">@null</item>
-        <item name="android:navigationBarColor">@android:color/transparent</item>
-        <item name="android:windowLayoutInDisplayCutoutMode">always</item>
-        <item name="enableEdgeToEdge">true</item>
-    </style>
-
-    <style name="ShortcutHelperTheme" parent="@style/ShortcutHelperThemeCommon">
-        <item name="android:windowLightNavigationBar">true</item>
-    </style>
 </resources>
diff --git a/packages/SystemUI/src/com/android/systemui/CameraProtectionLoader.kt b/packages/SystemUI/src/com/android/systemui/CameraProtectionLoader.kt
index 6cee28b..0c29466 100644
--- a/packages/SystemUI/src/com/android/systemui/CameraProtectionLoader.kt
+++ b/packages/SystemUI/src/com/android/systemui/CameraProtectionLoader.kt
@@ -22,15 +22,18 @@
 import android.graphics.RectF
 import android.util.PathParser
 import com.android.systemui.res.R
-import javax.inject.Inject
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import kotlin.math.roundToInt
 
 interface CameraProtectionLoader {
     fun loadCameraProtectionInfoList(): List<CameraProtectionInfo>
 }
 
-class CameraProtectionLoaderImpl @Inject constructor(private val context: Context) :
-    CameraProtectionLoader {
+class CameraProtectionLoaderImpl
+@AssistedInject
+constructor(@Assisted private val context: Context) : CameraProtectionLoader {
 
     override fun loadCameraProtectionInfoList(): List<CameraProtectionInfo> {
         val list = mutableListOf<CameraProtectionInfo>()
@@ -76,7 +79,7 @@
                 computed.left.roundToInt(),
                 computed.top.roundToInt(),
                 computed.right.roundToInt(),
-                computed.bottom.roundToInt()
+                computed.bottom.roundToInt(),
             )
         val displayUniqueId = context.getString(displayUniqueIdRes)
         return CameraProtectionInfo(
@@ -84,7 +87,7 @@
             physicalCameraId,
             protectionPath,
             protectionBounds,
-            displayUniqueId
+            displayUniqueId,
         )
     }
 
@@ -95,4 +98,9 @@
             throw IllegalArgumentException("Invalid protection path", e)
         }
     }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(context: Context): CameraProtectionLoaderImpl
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/CameraProtectionModule.kt b/packages/SystemUI/src/com/android/systemui/CameraProtectionModule.kt
index 58680a8..442a1e4 100644
--- a/packages/SystemUI/src/com/android/systemui/CameraProtectionModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/CameraProtectionModule.kt
@@ -16,11 +16,20 @@
 
 package com.android.systemui
 
-import dagger.Binds
+import android.content.Context
+import com.android.systemui.dagger.SysUISingleton
 import dagger.Module
+import dagger.Provides
 
 @Module
-interface CameraProtectionModule {
+object CameraProtectionModule {
 
-    @Binds fun cameraProtectionLoaderImpl(impl: CameraProtectionLoaderImpl): CameraProtectionLoader
+    @Provides
+    @SysUISingleton
+    fun cameraProtectionLoader(
+        factory: CameraProtectionLoaderImpl.Factory,
+        context: Context,
+    ): CameraProtectionLoader {
+        return factory.create(context)
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/ScreenDecorationsModule.kt b/packages/SystemUI/src/com/android/systemui/ScreenDecorationsModule.kt
index 044312b..6fc50fb 100644
--- a/packages/SystemUI/src/com/android/systemui/ScreenDecorationsModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/ScreenDecorationsModule.kt
@@ -16,9 +16,14 @@
 
 package com.android.systemui
 
+import android.content.Context
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.decor.FaceScanningProviderFactory
+import com.android.systemui.decor.FaceScanningProviderFactoryImpl
 import com.android.systemui.statusbar.policy.ConfigurationController.ConfigurationListener
 import dagger.Binds
 import dagger.Module
+import dagger.Provides
 import dagger.multibindings.ClassKey
 import dagger.multibindings.IntoMap
 import dagger.multibindings.IntoSet
@@ -35,4 +40,15 @@
     @Binds
     @IntoSet
     fun bindScreenDecorationsConfigListener(impl: ScreenDecorations): ConfigurationListener
+
+    companion object {
+        @Provides
+        @SysUISingleton
+        fun faceScanningProviderFactory(
+            creator: FaceScanningProviderFactoryImpl.Creator,
+            context: Context,
+        ): FaceScanningProviderFactory {
+            return creator.create(context)
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/SysUICutoutProvider.kt b/packages/SystemUI/src/com/android/systemui/SysUICutoutProvider.kt
index 7309599..b4cb103 100644
--- a/packages/SystemUI/src/com/android/systemui/SysUICutoutProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/SysUICutoutProvider.kt
@@ -21,21 +21,12 @@
 import android.util.RotationUtils
 import android.view.Display
 import android.view.DisplayCutout
-import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.display.naturalBounds
-import javax.inject.Inject
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 
-@SysUISingleton
-class SysUICutoutProvider
-@Inject
-constructor(
-    private val context: Context,
-    private val cameraProtectionLoader: CameraProtectionLoader,
-) {
-
-    private val cameraProtectionList by lazy {
-        cameraProtectionLoader.loadCameraProtectionInfoList()
-    }
+interface SysUICutoutProvider {
 
     /**
      * Returns the [SysUICutoutInformation] for the current display and the current rotation.
@@ -43,7 +34,21 @@
      * This means that the bounds of the display cutout and the camera protection will be
      * adjusted/rotated for the current rotation.
      */
-    fun cutoutInfoForCurrentDisplayAndRotation(): SysUICutoutInformation? {
+    fun cutoutInfoForCurrentDisplayAndRotation(): SysUICutoutInformation?
+}
+
+class SysUICutoutProviderImpl
+@AssistedInject
+constructor(
+    @Assisted private val context: Context,
+    @Assisted private val cameraProtectionLoader: CameraProtectionLoader,
+) : SysUICutoutProvider {
+
+    private val cameraProtectionList by lazy {
+        cameraProtectionLoader.loadCameraProtectionInfoList()
+    }
+
+    override fun cutoutInfoForCurrentDisplayAndRotation(): SysUICutoutInformation? {
         val display = context.display
         val displayCutout: DisplayCutout = display.cutout ?: return null
         return SysUICutoutInformation(displayCutout, getCameraProtectionForDisplay(display))
@@ -72,8 +77,16 @@
             /* inOutBounds = */ rotatedBoundsOut,
             /* parentWidth = */ displayNaturalBounds.width(),
             /* parentHeight = */ displayNaturalBounds.height(),
-            /* rotation = */ display.rotation
+            /* rotation = */ display.rotation,
         )
         return rotatedBoundsOut
     }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(
+            context: Context,
+            cameraProtectionLoader: CameraProtectionLoader,
+        ): SysUICutoutProviderImpl
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/dagger/AudioSharingModule.kt b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/dagger/AudioSharingModule.kt
index 50970a5..8b5fde3 100644
--- a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/dagger/AudioSharingModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/dagger/AudioSharingModule.kt
@@ -55,11 +55,7 @@
             settingsLibAudioSharingRepository: SettingsLibAudioSharingRepository,
             @Background backgroundDispatcher: CoroutineDispatcher,
         ): AudioSharingRepository =
-            if (
-                Flags.enableLeAudioSharing() &&
-                    Flags.audioSharingQsDialogImprovement() &&
-                    localBluetoothManager != null
-            ) {
+            if (Flags.enableLeAudioSharing() && localBluetoothManager != null) {
                 AudioSharingRepositoryImpl(
                     localBluetoothManager,
                     settingsLibAudioSharingRepository,
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt
index 0bcb58d..5f97391 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt
@@ -175,8 +175,10 @@
                             actions.map { action ->
                                 UserSwitcherDropdownItemViewModel(
                                     icon =
-                                        Icon.Resource(
-                                            action.iconResourceId,
+                                        Icon.Loaded(
+                                            applicationContext.resources.getDrawable(
+                                                action.iconResourceId
+                                            ),
                                             contentDescription = null,
                                         ),
                                     text = Text.Resource(action.textResourceId),
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModel.kt
index a519649..db4bee7 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/ResizeableItemFrameViewModel.kt
@@ -13,7 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
 package com.android.systemui.communal.ui.viewmodel
 
 import androidx.compose.foundation.gestures.AnchoredDraggableState
@@ -21,6 +20,10 @@
 import androidx.compose.runtime.snapshotFlow
 import com.android.app.tracing.coroutines.coroutineScopeTraced as coroutineScope
 import com.android.systemui.lifecycle.ExclusiveActivatable
+import kotlin.math.abs
+import kotlin.math.ceil
+import kotlin.math.floor
+import kotlin.math.sign
 import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -50,14 +53,33 @@
 }
 
 class ResizeableItemFrameViewModel : ExclusiveActivatable() {
-    private data class GridLayoutInfo(
-        val minSpan: Int,
-        val maxSpan: Int,
-        val heightPerSpanPx: Float,
-        val verticalItemSpacingPx: Float,
+    data class GridLayoutInfo(
         val currentRow: Int,
         val currentSpan: Int,
-    )
+        val maxHeightPx: Int,
+        val minHeightPx: Int,
+        val resizeMultiple: Int,
+        val totalSpans: Int,
+        private val heightPerSpanPx: Float,
+        private val verticalItemSpacingPx: Float,
+    ) {
+        fun getPxOffsetForResize(spans: Int): Int =
+            (spans * (heightPerSpanPx + verticalItemSpacingPx)).toInt()
+
+        private fun getSpansForPx(height: Int): Int =
+            ceil((height + verticalItemSpacingPx) / (heightPerSpanPx + verticalItemSpacingPx))
+                .toInt()
+                .coerceIn(resizeMultiple, totalSpans)
+
+        private fun roundDownToMultiple(spans: Int): Int =
+            floor(spans.toDouble() / resizeMultiple).toInt() * resizeMultiple
+
+        val maxSpans: Int
+            get() = roundDownToMultiple(getSpansForPx(maxHeightPx))
+
+        val minSpans: Int
+            get() = roundDownToMultiple(getSpansForPx(minHeightPx))
+    }
 
     /**
      * The layout information necessary in order to calculate the pixel offsets of the drag anchor
@@ -84,37 +106,44 @@
      */
     fun setGridLayoutInfo(
         verticalItemSpacingPx: Float,
-        verticalContentPaddingPx: Float,
-        viewportHeightPx: Int,
-        maxItemSpan: Int,
-        minItemSpan: Int,
         currentRow: Int?,
-        currentSpan: Int?,
+        maxHeightPx: Int,
+        minHeightPx: Int,
+        currentSpan: Int,
+        resizeMultiple: Int,
+        totalSpans: Int,
+        viewportHeightPx: Int,
+        verticalContentPaddingPx: Float,
     ) {
-        if (currentSpan == null || currentRow == null) {
+        if (currentRow == null) {
             gridLayoutInfo.value = null
             return
         }
-        require(maxItemSpan >= minItemSpan) {
-            "Maximum item span of $maxItemSpan cannot be less than the minimum span of $minItemSpan"
+        require(maxHeightPx >= minHeightPx) {
+            "Maximum item span of $maxHeightPx cannot be less than the minimum span of $minHeightPx"
         }
-        require(minItemSpan in 1..maxItemSpan) {
-            "Minimum span must be between 1 and $maxItemSpan, but was $minItemSpan"
+
+        require(currentSpan <= totalSpans) {
+            "Current span ($currentSpan) cannot exceed the total number of spans ($totalSpans)"
         }
-        require(currentSpan % minItemSpan == 0) {
-            "Current span of $currentSpan is not a multiple of the minimum span of $minItemSpan"
+
+        require(resizeMultiple > 0) {
+            "Resize multiple ($resizeMultiple) must be a positive integer"
         }
         val availableHeight = viewportHeightPx - verticalContentPaddingPx
-        val totalSpacing = verticalItemSpacingPx * ((maxItemSpan / minItemSpan) - 1)
-        val heightPerSpanPx = (availableHeight - totalSpacing) / maxItemSpan
+        val heightPerSpanPx =
+            (availableHeight - (totalSpans - 1) * verticalItemSpacingPx) / totalSpans
+
         gridLayoutInfo.value =
             GridLayoutInfo(
-                minSpan = minItemSpan,
-                maxSpan = maxItemSpan,
                 heightPerSpanPx = heightPerSpanPx,
                 verticalItemSpacingPx = verticalItemSpacingPx,
                 currentRow = currentRow,
                 currentSpan = currentSpan,
+                maxHeightPx = maxHeightPx.coerceAtMost(availableHeight.toInt()),
+                minHeightPx = minHeightPx,
+                resizeMultiple = resizeMultiple,
+                totalSpans = totalSpans,
             )
     }
 
@@ -123,50 +152,46 @@
         layoutInfo: GridLayoutInfo?,
     ): DraggableAnchors<Int> {
 
-        if (layoutInfo == null || !isDragAllowed(handle, layoutInfo)) {
+        if (layoutInfo == null || (!isDragAllowed(handle, layoutInfo))) {
             return DraggableAnchors { 0 at 0f }
         }
-
-        val (
-            minItemSpan,
-            maxItemSpan,
-            heightPerSpanPx,
-            verticalSpacingPx,
-            currentRow,
-            currentSpan,
-        ) = layoutInfo
+        val currentRow = layoutInfo.currentRow
+        val currentSpan = layoutInfo.currentSpan
+        val minItemSpan = layoutInfo.minSpans
+        val maxItemSpan = layoutInfo.maxSpans
+        val totalSpans = layoutInfo.totalSpans
 
         // The maximum row this handle can be dragged to.
         val maxRow =
             if (handle == DragHandle.TOP) {
                 (currentRow + currentSpan - minItemSpan).coerceAtLeast(0)
             } else {
-                maxItemSpan
+                (currentRow + maxItemSpan).coerceAtMost(totalSpans)
             }
 
         // The minimum row this handle can be dragged to.
         val minRow =
             if (handle == DragHandle.TOP) {
-                0
+                (currentRow + currentSpan - maxItemSpan).coerceAtLeast(0)
             } else {
-                (currentRow + minItemSpan).coerceAtMost(maxItemSpan)
+                (currentRow + minItemSpan).coerceAtMost(totalSpans)
             }
 
         // The current row position of this handle
         val currentPosition = if (handle == DragHandle.TOP) currentRow else currentRow + currentSpan
 
         return DraggableAnchors {
-            for (targetRow in minRow..maxRow step minItemSpan) {
+            for (targetRow in minRow..maxRow step layoutInfo.resizeMultiple) {
                 val diff = targetRow - currentPosition
-                val spacing = diff / minItemSpan * verticalSpacingPx
-                diff at diff * heightPerSpanPx + spacing
+                val pixelOffset = (layoutInfo.getPxOffsetForResize(abs(diff)) * diff.sign).toFloat()
+                diff at pixelOffset
             }
         }
     }
 
     private fun isDragAllowed(handle: DragHandle, layoutInfo: GridLayoutInfo): Boolean {
-        val minItemSpan = layoutInfo.minSpan
-        val maxItemSpan = layoutInfo.maxSpan
+        val minItemSpan = layoutInfo.minSpans
+        val maxItemSpan = layoutInfo.maxSpans
         val currentRow = layoutInfo.currentRow
         val currentSpan = layoutInfo.currentSpan
         val atMinSize = currentSpan == minItemSpan
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java
index a5b2277..c6be0dd 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java
@@ -34,6 +34,7 @@
 import com.android.systemui.dock.DockManager;
 import com.android.systemui.dock.DockManagerImpl;
 import com.android.systemui.doze.DozeHost;
+import com.android.systemui.education.dagger.ContextualEducationModule;
 import com.android.systemui.inputdevice.tutorial.KeyboardTouchpadTutorialModule;
 import com.android.systemui.keyboard.shortcut.ShortcutHelperModule;
 import com.android.systemui.keyguard.ui.composable.blueprint.DefaultBlueprintModule;
@@ -153,6 +154,7 @@
         VolumeModule.class,
         WallpaperModule.class,
         ShortcutHelperModule.class,
+        ContextualEducationModule.class,
 })
 public abstract class ReferenceSystemUIModule {
 
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index b55108d..450863f 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -63,7 +63,6 @@
 import com.android.systemui.doze.dagger.DozeComponent;
 import com.android.systemui.dreams.dagger.DreamModule;
 import com.android.systemui.dump.DumpManager;
-import com.android.systemui.education.dagger.ContextualEducationModule;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.FlagDependenciesModule;
 import com.android.systemui.flags.FlagsModule;
@@ -272,8 +271,7 @@
         UserModule.class,
         UtilModule.class,
         NoteTaskModule.class,
-        WalletModule.class,
-        ContextualEducationModule.class
+        WalletModule.class
 },
         subcomponents = {
                 ComplicationComponent.class,
diff --git a/packages/SystemUI/src/com/android/systemui/decor/CutoutDecorProviderFactory.kt b/packages/SystemUI/src/com/android/systemui/decor/CutoutDecorProviderFactory.kt
index cbed21c..bfd6b5b 100644
--- a/packages/SystemUI/src/com/android/systemui/decor/CutoutDecorProviderFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/decor/CutoutDecorProviderFactory.kt
@@ -22,18 +22,15 @@
 import android.view.DisplayCutout
 import android.view.DisplayInfo
 
-class CutoutDecorProviderFactory constructor(
-    private val res: Resources,
-    private val display: Display?,
-) : DecorProviderFactory() {
+class CutoutDecorProviderFactory(private val res: Resources, private val display: Display?) :
+    DecorProviderFactory {
 
     val displayInfo = DisplayInfo()
 
     override val hasProviders: Boolean
         get() {
-            display?.getDisplayInfo(displayInfo) ?: run {
-                Log.w(TAG, "display is null, can't update displayInfo")
-            }
+            display?.getDisplayInfo(displayInfo)
+                ?: run { Log.w(TAG, "display is null, can't update displayInfo") }
             return DisplayCutout.getFillBuiltInDisplayCutout(res, displayInfo.uniqueId)
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/decor/DecorProviderFactory.kt b/packages/SystemUI/src/com/android/systemui/decor/DecorProviderFactory.kt
index c60cad8..16e73f5 100644
--- a/packages/SystemUI/src/com/android/systemui/decor/DecorProviderFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/decor/DecorProviderFactory.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.decor
 
-abstract class DecorProviderFactory {
-    abstract val providers: List<DecorProvider>
-    abstract val hasProviders: Boolean
-}
\ No newline at end of file
+interface DecorProviderFactory {
+    val providers: List<DecorProvider>
+    val hasProviders: Boolean
+}
diff --git a/packages/SystemUI/src/com/android/systemui/decor/FaceScanningProviderFactory.kt b/packages/SystemUI/src/com/android/systemui/decor/FaceScanningProviderFactory.kt
index 3bc4f34..88580cf 100644
--- a/packages/SystemUI/src/com/android/systemui/decor/FaceScanningProviderFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/decor/FaceScanningProviderFactory.kt
@@ -33,23 +33,32 @@
 import com.android.systemui.FaceScanningOverlay
 import com.android.systemui.biometrics.AuthController
 import com.android.systemui.biometrics.data.repository.FacePropertyRepository
-import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.log.ScreenDecorationsLogger
 import com.android.systemui.plugins.statusbar.StatusBarStateController
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import java.util.concurrent.Executor
-import javax.inject.Inject
 
-@SysUISingleton
-class FaceScanningProviderFactory @Inject constructor(
-        private val authController: AuthController,
-        private val context: Context,
-        private val statusBarStateController: StatusBarStateController,
-        private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
-        @Main private val mainExecutor: Executor,
-        private val logger: ScreenDecorationsLogger,
-        private val facePropertyRepository: FacePropertyRepository,
-) : DecorProviderFactory() {
+interface FaceScanningProviderFactory : DecorProviderFactory {
+
+    fun canShowFaceScanningAnim(): Boolean
+
+    fun shouldShowFaceScanningAnim(): Boolean
+}
+
+class FaceScanningProviderFactoryImpl
+@AssistedInject
+constructor(
+    private val authController: AuthController,
+    @Assisted private val context: Context,
+    private val statusBarStateController: StatusBarStateController,
+    private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
+    @Main private val mainExecutor: Executor,
+    private val logger: ScreenDecorationsLogger,
+    private val facePropertyRepository: FacePropertyRepository,
+) : FaceScanningProviderFactory {
     private val display = context.display
     private val displayInfo = DisplayInfo()
 
@@ -60,11 +69,12 @@
             }
 
             // update display info
-            display?.getDisplayInfo(displayInfo) ?: run {
-                Log.w(TAG, "display is null, can't update displayInfo")
-            }
+            display?.getDisplayInfo(displayInfo)
+                ?: run { Log.w(TAG, "display is null, can't update displayInfo") }
             return DisplayCutout.getFillBuiltInDisplayCutout(
-                    context.resources, displayInfo.uniqueId)
+                context.resources,
+                displayInfo.uniqueId,
+            )
         }
 
     override val providers: List<DecorProvider>
@@ -81,39 +91,45 @@
                     // Cutout drawing is updated in ScreenDecorations#updateCutout
                     for (bound in bounds) {
                         list.add(
-                                FaceScanningOverlayProviderImpl(
-                                        bound.baseOnRotation0(displayInfo.rotation),
-                                        authController,
-                                        statusBarStateController,
-                                        keyguardUpdateMonitor,
-                                        mainExecutor,
-                                        logger,
-                                        facePropertyRepository,
-                                )
+                            FaceScanningOverlayProviderImpl(
+                                bound.baseOnRotation0(displayInfo.rotation),
+                                authController,
+                                statusBarStateController,
+                                keyguardUpdateMonitor,
+                                mainExecutor,
+                                logger,
+                                facePropertyRepository,
+                            )
                         )
                     }
                 }
             }
         }
 
-    fun canShowFaceScanningAnim(): Boolean {
+    override fun canShowFaceScanningAnim(): Boolean {
         return hasProviders && keyguardUpdateMonitor.isFaceEnabledAndEnrolled
     }
 
-    fun shouldShowFaceScanningAnim(): Boolean {
+    override fun shouldShowFaceScanningAnim(): Boolean {
         return canShowFaceScanningAnim() &&
-                (keyguardUpdateMonitor.isFaceDetectionRunning || authController.isShowing)
+            (keyguardUpdateMonitor.isFaceDetectionRunning || authController.isShowing)
+    }
+
+    // Using the name "Creator" so that it doesn't become "...FactoryFactory".
+    @AssistedFactory
+    interface Creator {
+        fun create(context: Context): FaceScanningProviderFactoryImpl
     }
 }
 
 class FaceScanningOverlayProviderImpl(
-        override val alignedBound: Int,
-        private val authController: AuthController,
-        private val statusBarStateController: StatusBarStateController,
-        private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
-        private val mainExecutor: Executor,
-        private val logger: ScreenDecorationsLogger,
-        private val facePropertyRepository: FacePropertyRepository,
+    override val alignedBound: Int,
+    private val authController: AuthController,
+    private val statusBarStateController: StatusBarStateController,
+    private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
+    private val mainExecutor: Executor,
+    private val logger: ScreenDecorationsLogger,
+    private val facePropertyRepository: FacePropertyRepository,
 ) : BoundDecorProvider() {
     override val viewId: Int = com.android.systemui.res.R.id.face_scanning_anim
 
@@ -122,7 +138,7 @@
         reloadToken: Int,
         @Surface.Rotation rotation: Int,
         tintColor: Int,
-        displayUniqueId: String?
+        displayUniqueId: String?,
     ) {
         (view.layoutParams as FrameLayout.LayoutParams).let {
             updateLayoutParams(it, rotation)
@@ -138,9 +154,10 @@
         context: Context,
         parent: ViewGroup,
         @Surface.Rotation rotation: Int,
-        tintColor: Int
+        tintColor: Int,
     ): View {
-        val view = FaceScanningOverlay(
+        val view =
+            FaceScanningOverlay(
                 context,
                 alignedBound,
                 statusBarStateController,
@@ -148,43 +165,46 @@
                 mainExecutor,
                 logger,
                 authController,
-        )
+            )
         view.id = viewId
         view.setColor(tintColor)
-        FrameLayout.LayoutParams(ViewGroup.LayoutParams.MATCH_PARENT,
-                ViewGroup.LayoutParams.MATCH_PARENT).let {
-            updateLayoutParams(it, rotation)
-            parent.addView(view, it)
-        }
+        FrameLayout.LayoutParams(
+                ViewGroup.LayoutParams.MATCH_PARENT,
+                ViewGroup.LayoutParams.MATCH_PARENT,
+            )
+            .let {
+                updateLayoutParams(it, rotation)
+                parent.addView(view, it)
+            }
         return view
     }
 
     private fun updateLayoutParams(
         layoutParams: FrameLayout.LayoutParams,
-        @Surface.Rotation rotation: Int
+        @Surface.Rotation rotation: Int,
     ) {
         layoutParams.let { lp ->
             lp.width = ViewGroup.LayoutParams.MATCH_PARENT
             lp.height = ViewGroup.LayoutParams.MATCH_PARENT
             logger.faceSensorLocation(facePropertyRepository.sensorLocation.value)
-            facePropertyRepository.sensorLocation.value?.y?.let {
-                faceAuthSensorHeight ->
+            facePropertyRepository.sensorLocation.value?.y?.let { faceAuthSensorHeight ->
                 val faceScanningHeight = (faceAuthSensorHeight * 2)
                 when (rotation) {
-                    Surface.ROTATION_0, Surface.ROTATION_180 ->
-                        lp.height = faceScanningHeight
-                    Surface.ROTATION_90, Surface.ROTATION_270 ->
-                        lp.width = faceScanningHeight
+                    Surface.ROTATION_0,
+                    Surface.ROTATION_180 -> lp.height = faceScanningHeight
+                    Surface.ROTATION_90,
+                    Surface.ROTATION_270 -> lp.width = faceScanningHeight
                 }
             }
 
-            lp.gravity = when (rotation) {
-                Surface.ROTATION_0 -> Gravity.TOP or Gravity.START
-                Surface.ROTATION_90 -> Gravity.LEFT or Gravity.START
-                Surface.ROTATION_180 -> Gravity.BOTTOM or Gravity.END
-                Surface.ROTATION_270 -> Gravity.RIGHT or Gravity.END
-                else -> -1 /* invalid rotation */
-            }
+            lp.gravity =
+                when (rotation) {
+                    Surface.ROTATION_0 -> Gravity.TOP or Gravity.START
+                    Surface.ROTATION_90 -> Gravity.LEFT or Gravity.START
+                    Surface.ROTATION_180 -> Gravity.BOTTOM or Gravity.END
+                    Surface.ROTATION_270 -> Gravity.RIGHT or Gravity.END
+                    else -> -1 /* invalid rotation */
+                }
         }
     }
 }
@@ -209,24 +229,27 @@
 fun Int.baseOnRotation0(@DisplayCutout.BoundsPosition currentRotation: Int): Int {
     return when (currentRotation) {
         Surface.ROTATION_0 -> this
-        Surface.ROTATION_90 -> when (this) {
-            BOUNDS_POSITION_LEFT -> BOUNDS_POSITION_TOP
-            BOUNDS_POSITION_TOP -> BOUNDS_POSITION_RIGHT
-            BOUNDS_POSITION_RIGHT -> BOUNDS_POSITION_BOTTOM
-            else /* BOUNDS_POSITION_BOTTOM */ -> BOUNDS_POSITION_LEFT
-        }
-        Surface.ROTATION_270 -> when (this) {
-            BOUNDS_POSITION_LEFT -> BOUNDS_POSITION_BOTTOM
-            BOUNDS_POSITION_TOP -> BOUNDS_POSITION_LEFT
-            BOUNDS_POSITION_RIGHT -> BOUNDS_POSITION_TOP
-            else /* BOUNDS_POSITION_BOTTOM */ -> BOUNDS_POSITION_RIGHT
-        }
-        else /* Surface.ROTATION_180 */ -> when (this) {
-            BOUNDS_POSITION_LEFT -> BOUNDS_POSITION_RIGHT
-            BOUNDS_POSITION_TOP -> BOUNDS_POSITION_BOTTOM
-            BOUNDS_POSITION_RIGHT -> BOUNDS_POSITION_LEFT
-            else /* BOUNDS_POSITION_BOTTOM */ -> BOUNDS_POSITION_TOP
-        }
+        Surface.ROTATION_90 ->
+            when (this) {
+                BOUNDS_POSITION_LEFT -> BOUNDS_POSITION_TOP
+                BOUNDS_POSITION_TOP -> BOUNDS_POSITION_RIGHT
+                BOUNDS_POSITION_RIGHT -> BOUNDS_POSITION_BOTTOM
+                else /* BOUNDS_POSITION_BOTTOM */ -> BOUNDS_POSITION_LEFT
+            }
+        Surface.ROTATION_270 ->
+            when (this) {
+                BOUNDS_POSITION_LEFT -> BOUNDS_POSITION_BOTTOM
+                BOUNDS_POSITION_TOP -> BOUNDS_POSITION_LEFT
+                BOUNDS_POSITION_RIGHT -> BOUNDS_POSITION_TOP
+                else /* BOUNDS_POSITION_BOTTOM */ -> BOUNDS_POSITION_RIGHT
+            }
+        else /* Surface.ROTATION_180 */ ->
+            when (this) {
+                BOUNDS_POSITION_LEFT -> BOUNDS_POSITION_RIGHT
+                BOUNDS_POSITION_TOP -> BOUNDS_POSITION_BOTTOM
+                BOUNDS_POSITION_RIGHT -> BOUNDS_POSITION_LEFT
+                else /* BOUNDS_POSITION_BOTTOM */ -> BOUNDS_POSITION_TOP
+            }
     }
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/decor/PrivacyDotDecorProviderFactory.kt b/packages/SystemUI/src/com/android/systemui/decor/PrivacyDotDecorProviderFactory.kt
index 14ecc66..9aa7fd1 100644
--- a/packages/SystemUI/src/com/android/systemui/decor/PrivacyDotDecorProviderFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/decor/PrivacyDotDecorProviderFactory.kt
@@ -23,19 +23,18 @@
 import android.view.Surface
 import android.view.View
 import android.view.ViewGroup
-import com.android.systemui.res.R
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.res.R
 import javax.inject.Inject
 
 /**
- * Provides privacy dot views for each orientation. The PrivacyDot orientation and visibility
- * of the privacy dot views are controlled by the PrivacyDotViewController.
+ * Provides privacy dot views for each orientation. The PrivacyDot orientation and visibility of the
+ * privacy dot views are controlled by the PrivacyDotViewController.
  */
 @SysUISingleton
-open class PrivacyDotDecorProviderFactory @Inject constructor(
-    @Main private val res: Resources
-) : DecorProviderFactory() {
+open class PrivacyDotDecorProviderFactory @Inject constructor(@Main private val res: Resources) :
+    DecorProviderFactory {
 
     private val isPrivacyDotEnabled: Boolean
         get() = res.getBoolean(R.bool.config_enablePrivacyDot)
@@ -51,22 +50,26 @@
                         viewId = R.id.privacy_dot_top_left_container,
                         alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
                         alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
-                        layoutId = R.layout.privacy_dot_top_left),
+                        layoutId = R.layout.privacy_dot_top_left,
+                    ),
                     PrivacyDotCornerDecorProviderImpl(
                         viewId = R.id.privacy_dot_top_right_container,
                         alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
                         alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
-                        layoutId = R.layout.privacy_dot_top_right),
+                        layoutId = R.layout.privacy_dot_top_right,
+                    ),
                     PrivacyDotCornerDecorProviderImpl(
                         viewId = R.id.privacy_dot_bottom_left_container,
                         alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
                         alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
-                        layoutId = R.layout.privacy_dot_bottom_left),
+                        layoutId = R.layout.privacy_dot_bottom_left,
+                    ),
                     PrivacyDotCornerDecorProviderImpl(
                         viewId = R.id.privacy_dot_bottom_right_container,
                         alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
                         alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
-                        layoutId = R.layout.privacy_dot_bottom_right)
+                        layoutId = R.layout.privacy_dot_bottom_right,
+                    ),
                 )
             } else {
                 emptyList()
@@ -78,7 +81,7 @@
     override val viewId: Int,
     @DisplayCutout.BoundsPosition override val alignedBound1: Int,
     @DisplayCutout.BoundsPosition override val alignedBound2: Int,
-    private val layoutId: Int
+    private val layoutId: Int,
 ) : CornerDecorProvider() {
 
     override fun onReloadResAndMeasure(
@@ -86,7 +89,7 @@
         reloadToken: Int,
         rotation: Int,
         tintColor: Int,
-        displayUniqueId: String?
+        displayUniqueId: String?,
     ) {
         // Do nothing here because it is handled inside PrivacyDotViewController
     }
@@ -95,7 +98,7 @@
         context: Context,
         parent: ViewGroup,
         @Surface.Rotation rotation: Int,
-        tintColor: Int
+        tintColor: Int,
     ): View {
         LayoutInflater.from(context).inflate(layoutId, parent, true)
         return parent.getChildAt(parent.childCount - 1 /* latest new added child */)
diff --git a/packages/SystemUI/src/com/android/systemui/decor/RoundedCornerDecorProviderFactory.kt b/packages/SystemUI/src/com/android/systemui/decor/RoundedCornerDecorProviderFactory.kt
index 2f2c952f..39fd551 100644
--- a/packages/SystemUI/src/com/android/systemui/decor/RoundedCornerDecorProviderFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/decor/RoundedCornerDecorProviderFactory.kt
@@ -21,65 +21,74 @@
 
 class RoundedCornerDecorProviderFactory(
     private val roundedCornerResDelegate: RoundedCornerResDelegate
-) : DecorProviderFactory() {
+) : DecorProviderFactory {
 
     override val hasProviders: Boolean
-        get() = roundedCornerResDelegate.run {
-            hasTop || hasBottom
-        }
+        get() = roundedCornerResDelegate.run { hasTop || hasBottom }
 
     override val providers: List<DecorProvider>
-    get() {
-        val hasTop = roundedCornerResDelegate.hasTop
-        val hasBottom = roundedCornerResDelegate.hasBottom
-        return when {
-            hasTop && hasBottom -> listOf(
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_top_left,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
-                    roundedCornerResDelegate = roundedCornerResDelegate),
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_top_right,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
-                    roundedCornerResDelegate = roundedCornerResDelegate),
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_bottom_left,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
-                    roundedCornerResDelegate = roundedCornerResDelegate),
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_bottom_right,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
-                    roundedCornerResDelegate = roundedCornerResDelegate)
-            )
-            hasTop -> listOf(
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_top_left,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
-                    roundedCornerResDelegate = roundedCornerResDelegate),
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_top_right,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
-                    roundedCornerResDelegate = roundedCornerResDelegate)
-            )
-            hasBottom -> listOf(
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_bottom_left,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
-                    roundedCornerResDelegate = roundedCornerResDelegate),
-                RoundedCornerDecorProviderImpl(
-                    viewId = R.id.rounded_corner_bottom_right,
-                    alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
-                    alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
-                    roundedCornerResDelegate = roundedCornerResDelegate)
-            )
-            else -> emptyList()
+        get() {
+            val hasTop = roundedCornerResDelegate.hasTop
+            val hasBottom = roundedCornerResDelegate.hasBottom
+            return when {
+                hasTop && hasBottom ->
+                    listOf(
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_top_left,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_top_right,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_bottom_left,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_bottom_right,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                    )
+                hasTop ->
+                    listOf(
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_top_left,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_top_right,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_TOP,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                    )
+                hasBottom ->
+                    listOf(
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_bottom_left,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_LEFT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                        RoundedCornerDecorProviderImpl(
+                            viewId = R.id.rounded_corner_bottom_right,
+                            alignedBound1 = DisplayCutout.BOUNDS_POSITION_BOTTOM,
+                            alignedBound2 = DisplayCutout.BOUNDS_POSITION_RIGHT,
+                            roundedCornerResDelegate = roundedCornerResDelegate,
+                        ),
+                    )
+                else -> emptyList()
+            }
         }
-    }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt b/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
index 7fa7da1..abe0289 100644
--- a/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
@@ -18,15 +18,12 @@
 
 import com.android.systemui.CoreStartable
 import com.android.systemui.Flags
-import com.android.systemui.contextualeducation.GestureType
 import com.android.systemui.coroutines.newTracingContext
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.education.data.repository.ContextualEducationRepository
 import com.android.systemui.education.data.repository.UserContextualEducationRepository
 import com.android.systemui.education.domain.interactor.ContextualEducationInteractor
 import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduInteractor
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractor
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractorImpl
 import com.android.systemui.education.ui.view.ContextualEduUiCoordinator
 import dagger.Binds
 import dagger.Lazy
@@ -83,18 +80,6 @@
         }
 
         @Provides
-        fun provideKeyboardTouchpadEduStatsInteractor(
-            implLazy: Lazy<KeyboardTouchpadEduStatsInteractorImpl>
-        ): KeyboardTouchpadEduStatsInteractor {
-            return if (Flags.keyboardTouchpadContextualEducation()) {
-                implLazy.get()
-            } else {
-                // No-op implementation when the flag is disabled.
-                return NoOpKeyboardTouchpadEduStatsInteractor
-            }
-        }
-
-        @Provides
         @IntoMap
         @ClassKey(KeyboardTouchpadEduInteractor::class)
         fun provideKeyboardTouchpadEduInteractor(
@@ -124,12 +109,6 @@
     }
 }
 
-private object NoOpKeyboardTouchpadEduStatsInteractor : KeyboardTouchpadEduStatsInteractor {
-    override fun incrementSignalCount(gestureType: GestureType) {}
-
-    override fun updateShortcutTriggerTime(gestureType: GestureType) {}
-}
-
 private object NoOpCoreStartable : CoreStartable {
     override fun start() {}
 }
diff --git a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
index faee326..c17f3fb 100644
--- a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
@@ -18,6 +18,9 @@
 
 import android.os.SystemProperties
 import com.android.systemui.CoreStartable
+import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
+import com.android.systemui.contextualeducation.GestureType
+import com.android.systemui.contextualeducation.GestureType.ALL_APPS
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.education.dagger.ContextualEducationModule.EduClock
@@ -25,6 +28,13 @@
 import com.android.systemui.education.shared.model.EducationInfo
 import com.android.systemui.education.shared.model.EducationUiType
 import com.android.systemui.inputdevice.data.repository.UserInputDeviceRepository
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.KEYBOARD
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.TOUCHPAD
+import com.android.systemui.inputdevice.tutorial.data.repository.TutorialSchedulerRepository
+import com.android.systemui.recents.OverviewProxyService
+import com.android.systemui.recents.OverviewProxyService.OverviewProxyListener
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import java.time.Clock
 import javax.inject.Inject
 import kotlin.time.Duration
@@ -33,9 +43,11 @@
 import kotlin.time.toDuration
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.launch
@@ -48,6 +60,8 @@
     @Background private val backgroundScope: CoroutineScope,
     private val contextualEducationInteractor: ContextualEducationInteractor,
     private val userInputDeviceRepository: UserInputDeviceRepository,
+    private val tutorialRepository: TutorialSchedulerRepository,
+    private val overviewProxyService: OverviewProxyService,
     @EduClock private val clock: Clock,
 ) : CoreStartable {
 
@@ -59,14 +73,16 @@
             getDurationForConfig("persist.contextual_edu.usage_session_sec", 3.days)
         val minIntervalBetweenEdu =
             getDurationForConfig("persist.contextual_edu.edu_interval_sec", 7.days)
+        val initialDelayDuration =
+            getDurationForConfig("persist.contextual_edu.initial_delay_sec", 7.days)
 
         private fun getDurationForConfig(
             systemPropertyKey: String,
-            defaultDuration: Duration
+            defaultDuration: Duration,
         ): Duration =
             SystemProperties.getLong(
                     systemPropertyKey,
-                    /* defaultValue= */ defaultDuration.inWholeSeconds
+                    /* defaultValue= */ defaultDuration.inWholeSeconds,
                 )
                 .toDuration(DurationUnit.SECONDS)
     }
@@ -74,6 +90,24 @@
     private val _educationTriggered = MutableStateFlow<EducationInfo?>(null)
     val educationTriggered = _educationTriggered.asStateFlow()
 
+    private val statsUpdateRequests: Flow<StatsUpdateRequest> = conflatedCallbackFlow {
+        val listener: OverviewProxyListener =
+            object : OverviewProxyListener {
+                override fun updateContextualEduStats(
+                    isTrackpadGesture: Boolean,
+                    gestureType: GestureType,
+                ) {
+                    trySendWithFailureLogging(
+                        StatsUpdateRequest(isTrackpadGesture, gestureType),
+                        TAG,
+                    )
+                }
+            }
+
+        overviewProxyService.addCallback(listener)
+        awaitClose { overviewProxyService.removeCallback(listener) }
+    }
+
     @OptIn(ExperimentalCoroutinesApi::class)
     override fun start() {
         backgroundScope.launch {
@@ -133,6 +167,16 @@
                 contextualEducationInteractor.updateShortcutTriggerTime(it)
             }
         }
+
+        backgroundScope.launch {
+            statsUpdateRequests.collect {
+                if (it.isTrackpadGesture) {
+                    contextualEducationInteractor.updateShortcutTriggerTime(it.gestureType)
+                } else {
+                    incrementSignalCount(it.gestureType)
+                }
+            }
+        }
     }
 
     private fun isEducationNeeded(model: GestureEduModel): Boolean {
@@ -160,4 +204,41 @@
 
     private fun getEduType(model: GestureEduModel) =
         if (model.educationShownCount > 0) EducationUiType.Notification else EducationUiType.Toast
+
+    private suspend fun incrementSignalCount(gestureType: GestureType) {
+        val targetDevice = getTargetDevice(gestureType)
+        if (isTargetDeviceConnected(targetDevice) && hasInitialDelayElapsed(targetDevice)) {
+            contextualEducationInteractor.incrementSignalCount(gestureType)
+        }
+    }
+
+    private suspend fun isTargetDeviceConnected(deviceType: DeviceType): Boolean {
+        return when (deviceType) {
+            KEYBOARD -> userInputDeviceRepository.isAnyKeyboardConnectedForUser.first().isConnected
+            TOUCHPAD -> userInputDeviceRepository.isAnyTouchpadConnectedForUser.first().isConnected
+        }
+    }
+
+    /**
+     * Keyboard shortcut education would be provided for All Apps. Touchpad gesture education would
+     * be provided for the rest of the gesture types (i.e. Home, Overview, Back). This method maps
+     * gesture to its target education device.
+     */
+    private fun getTargetDevice(gestureType: GestureType) =
+        when (gestureType) {
+            ALL_APPS -> KEYBOARD
+            else -> TOUCHPAD
+        }
+
+    private suspend fun hasInitialDelayElapsed(deviceType: DeviceType): Boolean {
+        val oobeLaunchTime = tutorialRepository.launchTime(deviceType) ?: return false
+        return clock
+            .instant()
+            .isAfter(oobeLaunchTime.plusSeconds(initialDelayDuration.inWholeSeconds))
+    }
+
+    private data class StatsUpdateRequest(
+        val isTrackpadGesture: Boolean,
+        val gestureType: GestureType,
+    )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt
deleted file mode 100644
index 43e39cf..0000000
--- a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * Copyright 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.education.domain.interactor
-
-import android.os.SystemProperties
-import com.android.systemui.contextualeducation.GestureType
-import com.android.systemui.contextualeducation.GestureType.ALL_APPS
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.education.dagger.ContextualEducationModule.EduClock
-import com.android.systemui.inputdevice.data.repository.UserInputDeviceRepository
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.KEYBOARD
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.TOUCHPAD
-import com.android.systemui.inputdevice.tutorial.data.repository.TutorialSchedulerRepository
-import java.time.Clock
-import javax.inject.Inject
-import kotlin.time.Duration
-import kotlin.time.Duration.Companion.days
-import kotlin.time.DurationUnit
-import kotlin.time.toDuration
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.launch
-
-/**
- * Encapsulates the update functions of KeyboardTouchpadEduStatsInteractor. This encapsulation is
- * for having a different implementation of interactor when the feature flag is off.
- */
-interface KeyboardTouchpadEduStatsInteractor {
-    fun incrementSignalCount(gestureType: GestureType)
-
-    fun updateShortcutTriggerTime(gestureType: GestureType)
-}
-
-/** Allow update to education data related to keyboard/touchpad. */
-@SysUISingleton
-class KeyboardTouchpadEduStatsInteractorImpl
-@Inject
-constructor(
-    @Background private val backgroundScope: CoroutineScope,
-    private val contextualEducationInteractor: ContextualEducationInteractor,
-    private val inputDeviceRepository: UserInputDeviceRepository,
-    private val tutorialRepository: TutorialSchedulerRepository,
-    @EduClock private val clock: Clock,
-) : KeyboardTouchpadEduStatsInteractor {
-
-    companion object {
-        val initialDelayDuration: Duration
-            get() =
-                SystemProperties.getLong(
-                        "persist.contextual_edu.initial_delay_sec",
-                        /* defaultValue= */ 7.days.inWholeSeconds,
-                    )
-                    .toDuration(DurationUnit.SECONDS)
-    }
-
-    override fun incrementSignalCount(gestureType: GestureType) {
-        backgroundScope.launch {
-            val targetDevice = getTargetDevice(gestureType)
-            if (isTargetDeviceConnected(targetDevice) && hasInitialDelayElapsed(targetDevice)) {
-                contextualEducationInteractor.incrementSignalCount(gestureType)
-            }
-        }
-    }
-
-    override fun updateShortcutTriggerTime(gestureType: GestureType) {
-        backgroundScope.launch {
-            contextualEducationInteractor.updateShortcutTriggerTime(gestureType)
-        }
-    }
-
-    private suspend fun isTargetDeviceConnected(deviceType: DeviceType): Boolean {
-        return when (deviceType) {
-            KEYBOARD -> inputDeviceRepository.isAnyKeyboardConnectedForUser.first().isConnected
-            TOUCHPAD -> inputDeviceRepository.isAnyTouchpadConnectedForUser.first().isConnected
-        }
-    }
-
-    /**
-     * Keyboard shortcut education would be provided for All Apps. Touchpad gesture education would
-     * be provided for the rest of the gesture types (i.e. Home, Overview, Back). This method maps
-     * gesture to its target education device.
-     */
-    private fun getTargetDevice(gestureType: GestureType) =
-        when (gestureType) {
-            ALL_APPS -> KEYBOARD
-            else -> TOUCHPAD
-        }
-
-    private suspend fun hasInitialDelayElapsed(deviceType: DeviceType): Boolean {
-        val oobeLaunchTime = tutorialRepository.launchTime(deviceType) ?: return false
-        return clock
-            .instant()
-            .isAfter(oobeLaunchTime.plusSeconds(initialDelayDuration.inWholeSeconds))
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ShortcutHelperModule.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ShortcutHelperModule.kt
index 906f600..7b3380a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ShortcutHelperModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ShortcutHelperModule.kt
@@ -16,7 +16,6 @@
 
 package com.android.systemui.keyboard.shortcut
 
-import android.app.Activity
 import com.android.systemui.CoreStartable
 import com.android.systemui.Flags.keyboardShortcutHelperRewrite
 import com.android.systemui.keyboard.shortcut.data.repository.ShortcutHelperStateRepository
@@ -31,8 +30,7 @@
 import com.android.systemui.keyboard.shortcut.qualifiers.InputShortcuts
 import com.android.systemui.keyboard.shortcut.qualifiers.MultitaskingShortcuts
 import com.android.systemui.keyboard.shortcut.qualifiers.SystemShortcuts
-import com.android.systemui.keyboard.shortcut.ui.ShortcutHelperActivityStarter
-import com.android.systemui.keyboard.shortcut.ui.view.ShortcutHelperActivity
+import com.android.systemui.keyboard.shortcut.ui.ShortcutHelperDialogStarter
 import dagger.Binds
 import dagger.Lazy
 import dagger.Module
@@ -44,11 +42,6 @@
 interface ShortcutHelperModule {
 
     @Binds
-    @IntoMap
-    @ClassKey(ShortcutHelperActivity::class)
-    fun activity(impl: ShortcutHelperActivity): Activity
-
-    @Binds
     @SystemShortcuts
     fun systemShortcutsSource(impl: SystemShortcutsSource): KeyboardShortcutGroupsSource
 
@@ -73,8 +66,8 @@
     companion object {
         @Provides
         @IntoMap
-        @ClassKey(ShortcutHelperActivityStarter::class)
-        fun starter(implLazy: Lazy<ShortcutHelperActivityStarter>): CoreStartable {
+        @ClassKey(ShortcutHelperDialogStarter::class)
+        fun starter(implLazy: Lazy<ShortcutHelperDialogStarter>): CoreStartable {
             return if (keyboardShortcutHelperRewrite()) {
                 implLazy.get()
             } else {
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperActivityStarter.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperActivityStarter.kt
deleted file mode 100644
index fbf52e7..0000000
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperActivityStarter.kt
+++ /dev/null
@@ -1,66 +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.keyboard.shortcut.ui
-
-import android.content.Context
-import android.content.Intent
-import com.android.systemui.CoreStartable
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.keyboard.shortcut.ui.view.ShortcutHelperActivity
-import com.android.systemui.keyboard.shortcut.ui.viewmodel.ShortcutHelperViewModel
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.launch
-
-@SysUISingleton
-class ShortcutHelperActivityStarter(
-    private val context: Context,
-    @Application private val applicationScope: CoroutineScope,
-    private val viewModel: ShortcutHelperViewModel,
-    private val startActivity: (Intent) -> Unit,
-) : CoreStartable {
-
-    @Inject
-    constructor(
-        context: Context,
-        @Application applicationScope: CoroutineScope,
-        viewModel: ShortcutHelperViewModel,
-    ) : this(
-        context,
-        applicationScope,
-        viewModel,
-        startActivity = { intent -> context.startActivity(intent) }
-    )
-
-    override fun start() {
-        applicationScope.launch {
-            viewModel.shouldShow.collect { shouldShow ->
-                if (shouldShow) {
-                    startShortcutHelperActivity()
-                }
-            }
-        }
-    }
-
-    private fun startShortcutHelperActivity() {
-        startActivity(
-            Intent(context, ShortcutHelperActivity::class.java)
-                .setFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
-        )
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarter.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarter.kt
new file mode 100644
index 0000000..d33ab2a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/ShortcutHelperDialogStarter.kt
@@ -0,0 +1,102 @@
+/*
+ * 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.keyboard.shortcut.ui
+
+import android.app.Dialog
+import android.content.ActivityNotFoundException
+import android.content.Intent
+import android.content.Intent.FLAG_ACTIVITY_NEW_TASK
+import android.os.UserHandle
+import android.provider.Settings
+import androidx.annotation.VisibleForTesting
+import androidx.compose.foundation.layout.width
+import androidx.compose.runtime.getValue
+import androidx.compose.ui.Modifier
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.keyboard.shortcut.ui.composable.ShortcutHelper
+import com.android.systemui.keyboard.shortcut.ui.composable.ShortcutHelperBottomSheet
+import com.android.systemui.keyboard.shortcut.ui.composable.getWidth
+import com.android.systemui.keyboard.shortcut.ui.viewmodel.ShortcutHelperViewModel
+import com.android.systemui.plugins.ActivityStarter
+import com.android.systemui.statusbar.phone.SystemUIDialogFactory
+import com.android.systemui.statusbar.phone.createBottomSheet
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.map
+
+@SysUISingleton
+class ShortcutHelperDialogStarter
+@Inject
+constructor(
+    @Application private val applicationScope: CoroutineScope,
+    private val viewModel: ShortcutHelperViewModel,
+    private val dialogFactory: SystemUIDialogFactory,
+    private val activityStarter: ActivityStarter,
+) : CoreStartable {
+
+    @VisibleForTesting var dialog: Dialog? = null
+
+    override fun start() {
+        viewModel.shouldShow
+            .map { shouldShow ->
+                if (shouldShow) {
+                    dialog = createShortcutHelperDialog().also { it.show() }
+                } else {
+                    dialog?.dismiss()
+                }
+            }
+            .launchIn(applicationScope)
+    }
+
+    private fun createShortcutHelperDialog(): Dialog {
+        return dialogFactory.createBottomSheet(
+            content = { dialog ->
+                val shortcutsUiState by viewModel.shortcutsUiState.collectAsStateWithLifecycle()
+                ShortcutHelper(
+                    modifier = Modifier.width(getWidth()),
+                    shortcutsUiState = shortcutsUiState,
+                    onKeyboardSettingsClicked = { onKeyboardSettingsClicked(dialog) },
+                    onSearchQueryChanged = { viewModel.onSearchQueryChanged(it) },
+                )
+                dialog.setOnDismissListener { viewModel.onViewClosed() }
+            },
+            maxWidth = ShortcutHelperBottomSheet.LargeScreenWidthLandscape
+        )
+    }
+
+    private fun onKeyboardSettingsClicked(dialog: Dialog) {
+        try {
+            activityStarter.startActivity(
+                Intent(Settings.ACTION_HARD_KEYBOARD_SETTINGS).addFlags(FLAG_ACTIVITY_NEW_TASK),
+                /* dismissShade= */ true,
+                /* animationController = */ null,
+                /* showOverLockscreenWhenLocked = */ false,
+                UserHandle.CURRENT,
+            )
+        } catch (e: ActivityNotFoundException) {
+            // From the Settings docs: In some cases, a matching Activity may not exist, so ensure
+            // you safeguard against this.
+            e.printStackTrace()
+            return
+        }
+        dialog.dismiss()
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelperUtils.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelperUtils.kt
index 1f0d696..e295564 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelperUtils.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelperUtils.kt
@@ -16,9 +16,13 @@
 
 package com.android.systemui.keyboard.shortcut.ui.composable
 
+import android.content.res.Configuration
 import androidx.compose.material3.windowsizeclass.WindowHeightSizeClass
 import androidx.compose.material3.windowsizeclass.WindowWidthSizeClass
 import androidx.compose.runtime.Composable
+import androidx.compose.ui.platform.LocalConfiguration
+import androidx.compose.ui.unit.Dp
+import androidx.compose.ui.unit.dp
 import com.android.compose.windowsizeclass.LocalWindowSizeClass
 
 /**
@@ -29,3 +33,21 @@
 fun hasCompactWindowSize() =
     LocalWindowSizeClass.current.widthSizeClass == WindowWidthSizeClass.Compact ||
         LocalWindowSizeClass.current.heightSizeClass == WindowHeightSizeClass.Compact
+
+@Composable
+fun getWidth(): Dp {
+    return if (hasCompactWindowSize()) {
+        ShortcutHelperBottomSheet.DefaultWidth
+    } else
+        when (LocalConfiguration.current.orientation) {
+            Configuration.ORIENTATION_LANDSCAPE ->
+                ShortcutHelperBottomSheet.LargeScreenWidthLandscape
+            else -> ShortcutHelperBottomSheet.LargeScreenWidthPortrait
+        }
+}
+
+object ShortcutHelperBottomSheet {
+    val DefaultWidth = 412.dp
+    val LargeScreenWidthPortrait = 704.dp
+    val LargeScreenWidthLandscape = 864.dp
+}
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
deleted file mode 100644
index 52263ce..0000000
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
+++ /dev/null
@@ -1,184 +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.keyboard.shortcut.ui.view
-
-import android.content.ActivityNotFoundException
-import android.content.Intent
-import android.content.Intent.FLAG_ACTIVITY_NEW_TASK
-import android.content.res.Configuration
-import android.os.Bundle
-import android.provider.Settings
-import androidx.activity.ComponentActivity
-import androidx.activity.compose.setContent
-import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.padding
-import androidx.compose.foundation.layout.size
-import androidx.compose.foundation.layout.width
-import androidx.compose.material3.ExperimentalMaterial3Api
-import androidx.compose.material3.MaterialTheme
-import androidx.compose.material3.ModalBottomSheet
-import androidx.compose.material3.Surface
-import androidx.compose.material3.rememberModalBottomSheetState
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.CompositionLocalProvider
-import androidx.compose.runtime.getValue
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.input.key.Key
-import androidx.compose.ui.input.key.key
-import androidx.compose.ui.input.key.onKeyEvent
-import androidx.compose.ui.platform.LocalConfiguration
-import androidx.compose.ui.platform.LocalContext
-import androidx.compose.ui.res.stringResource
-import androidx.compose.ui.semantics.contentDescription
-import androidx.compose.ui.semantics.semantics
-import androidx.compose.ui.unit.Dp
-import androidx.compose.ui.unit.dp
-import androidx.lifecycle.compose.collectAsStateWithLifecycle
-import androidx.lifecycle.flowWithLifecycle
-import androidx.lifecycle.lifecycleScope
-import com.android.compose.theme.PlatformTheme
-import com.android.systemui.keyboard.shortcut.ui.composable.ShortcutHelper
-import com.android.systemui.keyboard.shortcut.ui.composable.hasCompactWindowSize
-import com.android.systemui.keyboard.shortcut.ui.viewmodel.ShortcutHelperViewModel
-import com.android.systemui.res.R
-import com.android.systemui.settings.UserTracker
-import javax.inject.Inject
-import kotlinx.coroutines.launch
-
-/**
- * Activity that hosts the new version of the keyboard shortcut helper. It will be used both for
- * small and large screen devices.
- */
-class ShortcutHelperActivity
-@Inject
-constructor(private val userTracker: UserTracker, private val viewModel: ShortcutHelperViewModel) :
-    ComponentActivity() {
-
-    override fun onCreate(savedInstanceState: Bundle?) {
-        setupEdgeToEdge()
-        super.onCreate(savedInstanceState)
-        setContent { Content() }
-        observeFinishRequired()
-        viewModel.onViewOpened()
-    }
-
-    @Composable
-    private fun Content() {
-        CompositionLocalProvider(LocalContext provides userTracker.userContext) {
-            PlatformTheme { BottomSheet { finish() } }
-        }
-    }
-
-    @OptIn(ExperimentalMaterial3Api::class)
-    @Composable
-    private fun BottomSheet(onDismiss: () -> Unit) {
-        ModalBottomSheet(
-            onDismissRequest = { onDismiss() },
-            modifier =
-                Modifier.width(getWidth()).padding(top = getTopPadding()).onKeyEvent {
-                    if (it.key == Key.Escape) {
-                        onDismiss()
-                        true
-                    } else false
-                },
-            sheetState = rememberModalBottomSheetState(skipPartiallyExpanded = true),
-            dragHandle = { DragHandle() },
-        ) {
-            val shortcutsUiState by viewModel.shortcutsUiState.collectAsStateWithLifecycle()
-            ShortcutHelper(
-                shortcutsUiState = shortcutsUiState,
-                onKeyboardSettingsClicked = ::onKeyboardSettingsClicked,
-                onSearchQueryChanged = { viewModel.onSearchQueryChanged(it) },
-            )
-        }
-    }
-
-    @Composable
-    fun DragHandle() {
-        val dragHandleContentDescription =
-            stringResource(id = R.string.shortcut_helper_content_description_drag_handle)
-        Surface(
-            modifier =
-                Modifier.padding(top = 16.dp, bottom = 6.dp).semantics {
-                    contentDescription = dragHandleContentDescription
-                },
-            color = MaterialTheme.colorScheme.outlineVariant,
-            shape = MaterialTheme.shapes.extraLarge,
-        ) {
-            Box(Modifier.size(width = 32.dp, height = 4.dp))
-        }
-    }
-
-    private fun onKeyboardSettingsClicked() {
-        try {
-            startActivityAsUser(
-                Intent(Settings.ACTION_HARD_KEYBOARD_SETTINGS).addFlags(FLAG_ACTIVITY_NEW_TASK),
-                userTracker.userHandle,
-            )
-        } catch (e: ActivityNotFoundException) {
-            // From the Settings docs: In some cases, a matching Activity may not exist, so ensure
-            // you safeguard against this.
-            e.printStackTrace()
-        }
-    }
-
-    override fun onDestroy() {
-        super.onDestroy()
-        if (isFinishing) {
-            viewModel.onViewClosed()
-        }
-    }
-
-    private fun observeFinishRequired() {
-        lifecycleScope.launch {
-            viewModel.shouldShow.flowWithLifecycle(lifecycle).collect { shouldShow ->
-                if (!shouldShow) {
-                    finish()
-                }
-            }
-        }
-    }
-
-    private fun setupEdgeToEdge() {
-        // Draw behind system bars
-        window.setDecorFitsSystemWindows(false)
-    }
-
-    @Composable
-    private fun getTopPadding(): Dp {
-        return if (hasCompactWindowSize()) DefaultTopPadding else LargeScreenTopPadding
-    }
-
-    @Composable
-    private fun getWidth(): Dp {
-        return if (hasCompactWindowSize()) {
-            DefaultWidth
-        } else
-            when (LocalConfiguration.current.orientation) {
-                Configuration.ORIENTATION_LANDSCAPE -> LargeScreenWidthLandscape
-                else -> LargeScreenWidthPortrait
-            }
-    }
-
-    companion object {
-        private val DefaultTopPadding = 64.dp
-        private val LargeScreenTopPadding = 72.dp
-        private val DefaultWidth = 412.dp
-        private val LargeScreenWidthPortrait = 704.dp
-        private val LargeScreenWidthLandscape = 864.dp
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java
index 2961d05..742f435 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputAdapter.java
@@ -170,12 +170,11 @@
                 // Set different layout for each device
                 if (device.isMutingExpectedDevice()
                         && !mController.isCurrentConnectedDeviceRemote()) {
-                    updateTitleIcon(R.drawable.media_output_icon_volume,
-                            mController.getColorItemContent());
+                    updateUnmutedVolumeIcon(device);
                     mCurrentActivePosition = position;
                     updateFullItemClickListener(v -> onItemClick(v, device));
                     setSingleLineLayout(getItemTitle(device));
-                    initFakeActiveDevice();
+                    initFakeActiveDevice(device);
                 } else if (device.hasSubtext()) {
                     boolean isActiveWithOngoingSession =
                             (device.hasOngoingSession() && (currentlyConnected || isDeviceIncluded(
@@ -184,8 +183,7 @@
                             && isActiveWithOngoingSession;
                     if (isActiveWithOngoingSession) {
                         mCurrentActivePosition = position;
-                        updateTitleIcon(R.drawable.media_output_icon_volume,
-                                mController.getColorItemContent());
+                        updateUnmutedVolumeIcon(device);
                         mSubTitleText.setText(device.getSubtextString());
                         updateTwoLineLayoutContentAlpha(DEVICE_CONNECTED_ALPHA);
                         updateEndClickAreaAsSessionEditing(device,
@@ -199,9 +197,7 @@
                     } else {
                         if (currentlyConnected) {
                             mCurrentActivePosition = position;
-                            updateTitleIcon(R.drawable.media_output_icon_volume,
-                                    mController.getColorItemContent());
-                            initSeekbar(device, isCurrentSeekbarInvisible);
+                            updateUnmutedVolumeIcon(device);
                         } else {
                             setUpDeviceIcon(device);
                         }
@@ -243,8 +239,7 @@
                     // selected device in group
                     boolean isDeviceDeselectable = isDeviceIncluded(
                             mController.getDeselectableMediaDevice(), device);
-                    updateTitleIcon(R.drawable.media_output_icon_volume,
-                            mController.getColorItemContent());
+                    updateUnmutedVolumeIcon(device);
                     updateGroupableCheckBox(true, isDeviceDeselectable, device);
                     updateEndClickArea(device, isDeviceDeselectable);
                     disableFocusPropertyForView(mContainerLayout);
@@ -264,8 +259,7 @@
                         setSingleLineLayout(getItemTitle(device));
                     } else if (device.hasOngoingSession()) {
                         mCurrentActivePosition = position;
-                        updateTitleIcon(R.drawable.media_output_icon_volume,
-                                mController.getColorItemContent());
+                        updateUnmutedVolumeIcon(device);
                         updateEndClickAreaAsSessionEditing(device, device.isHostForOngoingSession()
                                 ? R.drawable.media_output_status_edit_session
                                 : R.drawable.ic_sound_bars_anim);
@@ -278,8 +272,7 @@
                             && !mController.getSelectableMediaDevice().isEmpty()) {
                         //If device is connected and there's other selectable devices, layout as
                         // one of selected devices.
-                        updateTitleIcon(R.drawable.media_output_icon_volume,
-                                mController.getColorItemContent());
+                        updateUnmutedVolumeIcon(device);
                         boolean isDeviceDeselectable = isDeviceIncluded(
                                 mController.getDeselectableMediaDevice(), device);
                         updateGroupableCheckBox(true, isDeviceDeselectable, device);
@@ -291,8 +284,7 @@
                                 true /* showEndTouchArea */);
                         initSeekbar(device, isCurrentSeekbarInvisible);
                     } else {
-                        updateTitleIcon(R.drawable.media_output_icon_volume,
-                                mController.getColorItemContent());
+                        updateUnmutedVolumeIcon(device);
                         disableFocusPropertyForView(mContainerLayout);
                         setUpContentDescriptionForView(mSeekBar, device);
                         mCurrentActivePosition = position;
diff --git a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java
index 63a7e01..574ccee 100644
--- a/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/media/dialog/MediaOutputBaseAdapter.java
@@ -47,6 +47,7 @@
 import androidx.annotation.VisibleForTesting;
 import androidx.recyclerview.widget.RecyclerView;
 
+import com.android.settingslib.media.InputMediaDevice;
 import com.android.settingslib.media.MediaDevice;
 import com.android.settingslib.utils.ThreadUtils;
 import com.android.systemui.res.R;
@@ -321,18 +322,20 @@
                     // Check if response volume match with the latest request, to ignore obsolete
                     // response
                     if (isCurrentSeekbarInvisible && !mIsInitVolumeFirstTime) {
-                        updateTitleIcon(currentVolume == 0 ? R.drawable.media_output_icon_volume_off
-                                        : R.drawable.media_output_icon_volume,
-                                mController.getColorItemContent());
+                        if (currentVolume == 0) {
+                            updateMutedVolumeIcon(device);
+                        } else {
+                            updateUnmutedVolumeIcon(device);
+                        }
                     } else {
                         if (!mVolumeAnimator.isStarted()) {
                             int percentage =
                                     (int) ((double) currentVolume * VOLUME_PERCENTAGE_SCALE_SIZE
                                             / (double) mSeekBar.getMax());
                             if (percentage == 0) {
-                                updateMutedVolumeIcon();
+                                updateMutedVolumeIcon(device);
                             } else {
-                                updateUnmutedVolumeIcon();
+                                updateUnmutedVolumeIcon(device);
                             }
                             mSeekBar.setVolume(currentVolume);
                             mLatestUpdateVolume = -1;
@@ -340,7 +343,7 @@
                     }
                 } else if (currentVolume == 0) {
                     mSeekBar.resetVolume();
-                    updateMutedVolumeIcon();
+                    updateMutedVolumeIcon(device);
                 }
                 if (currentVolume == mLatestUpdateVolume) {
                     mLatestUpdateVolume = -1;
@@ -365,7 +368,7 @@
                             R.string.media_output_dialog_volume_percentage, percentage));
                     mVolumeValueText.setVisibility(View.VISIBLE);
                     if (mStartFromMute) {
-                        updateUnmutedVolumeIcon();
+                        updateUnmutedVolumeIcon(device);
                         mStartFromMute = false;
                     }
                     if (progressToVolume != deviceVolume) {
@@ -390,9 +393,9 @@
                             seekBar.getProgress());
                     if (currentVolume == 0) {
                         seekBar.setProgress(0);
-                        updateMutedVolumeIcon();
+                        updateMutedVolumeIcon(device);
                     } else {
-                        updateUnmutedVolumeIcon();
+                        updateUnmutedVolumeIcon(device);
                     }
                     mTitleIcon.setVisibility(View.VISIBLE);
                     mVolumeValueText.setVisibility(View.GONE);
@@ -402,36 +405,48 @@
             });
         }
 
-        void updateMutedVolumeIcon() {
+        void updateMutedVolumeIcon(MediaDevice device) {
             mIconAreaLayout.setBackground(
                     mContext.getDrawable(R.drawable.media_output_item_background_active));
-            updateTitleIcon(R.drawable.media_output_icon_volume_off,
-                    mController.getColorItemContent());
+            updateTitleIcon(device, true /* isMutedVolumeIcon */);
         }
 
-        void updateUnmutedVolumeIcon() {
+        void updateUnmutedVolumeIcon(MediaDevice device) {
             mIconAreaLayout.setBackground(
                     mContext.getDrawable(R.drawable.media_output_title_icon_area)
             );
-            updateTitleIcon(R.drawable.media_output_icon_volume,
-                    mController.getColorItemContent());
+            updateTitleIcon(device, false /* isMutedVolumeIcon */);
         }
 
-        void updateTitleIcon(@DrawableRes int id, int color) {
+        void updateTitleIcon(MediaDevice device, boolean isMutedVolumeIcon) {
+            boolean isInputMediaDevice = device instanceof InputMediaDevice;
+            int id = getDrawableId(isInputMediaDevice, isMutedVolumeIcon);
             mTitleIcon.setImageDrawable(mContext.getDrawable(id));
-            mTitleIcon.setImageTintList(ColorStateList.valueOf(color));
+            mTitleIcon.setImageTintList(ColorStateList.valueOf(mController.getColorItemContent()));
             mIconAreaLayout.setBackgroundTintList(
                     ColorStateList.valueOf(mController.getColorSeekbarProgress()));
         }
 
+        @VisibleForTesting
+        int getDrawableId(boolean isInputDevice, boolean isMutedVolumeIcon) {
+            // Returns the microphone icon when the flag is enabled and the device is an input
+            // device.
+            if (com.android.media.flags.Flags.enableAudioInputDeviceRoutingAndVolumeControl()
+                    && isInputDevice) {
+                return isMutedVolumeIcon ? R.drawable.ic_mic_off : R.drawable.ic_mic_26dp;
+            }
+            return isMutedVolumeIcon
+                    ? R.drawable.media_output_icon_volume_off
+                    : R.drawable.media_output_icon_volume;
+        }
+
         void updateIconAreaClickListener(View.OnClickListener listener) {
             mIconAreaLayout.setOnClickListener(listener);
         }
 
-        void initFakeActiveDevice() {
+        void initFakeActiveDevice(MediaDevice device) {
             disableSeekBar();
-            updateTitleIcon(R.drawable.media_output_icon_volume,
-                    mController.getColorItemContent());
+            updateTitleIcon(device, false /* isMutedIcon */);
             final Drawable backgroundDrawable = mContext.getDrawable(
                                     R.drawable.media_output_item_background_active)
                             .mutate();
@@ -518,13 +533,13 @@
                     mController.logInteractionUnmuteDevice(device);
                     mSeekBar.setVolume(UNMUTE_DEFAULT_VOLUME);
                     mController.adjustVolume(device, UNMUTE_DEFAULT_VOLUME);
-                    updateUnmutedVolumeIcon();
+                    updateUnmutedVolumeIcon(device);
                     mIconAreaLayout.setOnTouchListener(((iconV, event) -> false));
                 } else {
                     mController.logInteractionMuteDevice(device);
                     mSeekBar.resetVolume();
                     mController.adjustVolume(device, 0);
-                    updateMutedVolumeIcon();
+                    updateMutedVolumeIcon(device);
                     mIconAreaLayout.setOnTouchListener(((iconV, event) -> {
                         mSeekBar.dispatchTouchEvent(event);
                         return false;
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt b/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt
index 62694ce..ef7e7eb 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt
@@ -22,18 +22,13 @@
 
 /**
  * Creates a [QSTile.Icon] from an [Icon].
- * * [Icon.Loaded] && [resId] null -> [QSTileImpl.DrawableIcon]
- * * [Icon.Loaded] && [resId] available -> [QSTileImpl.DrawableIconWithRes]
+ * * [Icon.Loaded] -> [QSTileImpl.DrawableIcon]
  * * [Icon.Resource] -> [QSTileImpl.ResourceIcon]
  */
-fun Icon.asQSTileIcon(resId: Int?): QSTile.Icon {
+fun Icon.asQSTileIcon(): QSTile.Icon {
     return when (this) {
         is Icon.Loaded -> {
-            if (resId != null) {
-                QSTileImpl.DrawableIconWithRes(this.drawable, resId)
-            } else {
-                QSTileImpl.DrawableIcon(this.drawable)
-            }
+            QSTileImpl.DrawableIcon(this.drawable)
         }
         is Icon.Resource -> {
             QSTileImpl.ResourceIcon.get(this.res)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
index 49b44cb..6db91ac 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
@@ -45,6 +45,7 @@
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.navigationBars
+import androidx.compose.foundation.layout.offset
 import androidx.compose.foundation.layout.windowInsetsPadding
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.runtime.Composable
@@ -55,9 +56,12 @@
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.graphicsLayer
 import androidx.compose.ui.layout.approachLayout
 import androidx.compose.ui.layout.onPlaced
+import androidx.compose.ui.layout.onSizeChanged
 import androidx.compose.ui.layout.positionInRoot
 import androidx.compose.ui.platform.ComposeView
 import androidx.compose.ui.res.dimensionResource
@@ -65,7 +69,9 @@
 import androidx.compose.ui.semantics.CustomAccessibilityAction
 import androidx.compose.ui.semantics.customActions
 import androidx.compose.ui.semantics.semantics
+import androidx.compose.ui.unit.IntOffset
 import androidx.compose.ui.unit.round
+import androidx.compose.ui.util.fastRoundToInt
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import androidx.lifecycle.lifecycleScope
@@ -143,9 +149,6 @@
 
     private lateinit var viewModel: QSFragmentComposeViewModel
 
-    // Starting with a non-zero value makes it so that it has a non-zero height on first expansion
-    // This is important for `QuickSettingsControllerImpl.mMinExpansionHeight` to detect a "change".
-    private val qqsHeight = MutableStateFlow(1)
     private val qsHeight = MutableStateFlow(0)
     private val qqsVisible = MutableStateFlow(false)
     private val qqsPositionOnRoot = Rect()
@@ -218,7 +221,7 @@
                 { notificationScrimClippingParams.params.top },
                 // Only allow scrolling when we are fully expanded. That way, we don't intercept
                 // swipes in lockscreen (when somehow QS is receiving touches).
-                { scrollState.canScrollForward && viewModel.expansionState.value.progress >= 1f },
+                { scrollState.canScrollForward && viewModel.isQsFullyExpanded },
             )
         frame.addView(
             composeView,
@@ -231,16 +234,20 @@
     @Composable
     private fun Content() {
         PlatformTheme {
-            val visible by viewModel.qsVisible.collectAsStateWithLifecycle()
-
             AnimatedVisibility(
-                visible = visible,
+                visible = viewModel.isQsVisible,
                 modifier =
-                    Modifier.windowInsetsPadding(WindowInsets.navigationBars).thenIf(
-                        notificationScrimClippingParams.isEnabled
-                    ) {
-                        Modifier.notificationScrimClip { notificationScrimClippingParams.params }
-                    },
+                    Modifier.graphicsLayer { alpha = viewModel.viewAlpha }
+                        .windowInsetsPadding(WindowInsets.navigationBars)
+                        // Clipping before translation to match QSContainerImpl.onDraw
+                        .offset {
+                            IntOffset(x = 0, y = viewModel.viewTranslationY.fastRoundToInt())
+                        }
+                        .thenIf(notificationScrimClippingParams.isEnabled) {
+                            Modifier.notificationScrimClip {
+                                notificationScrimClippingParams.params
+                            }
+                        },
             ) {
                 val isEditing by
                     viewModel.containerViewModel.editModeViewModel.isEditing
@@ -254,7 +261,7 @@
                     label = "EditModeAnimatedContent",
                 ) { editing ->
                     if (editing) {
-                        val qqsPadding by viewModel.qqsHeaderHeight.collectAsStateWithLifecycle()
+                        val qqsPadding = viewModel.qqsHeaderHeight
                         EditMode(
                             viewModel = viewModel.containerViewModel.editModeViewModel,
                             modifier =
@@ -283,7 +290,7 @@
     private fun CollapsableQuickSettingsSTL() {
         val sceneState = remember {
             MutableSceneTransitionLayoutState(
-                viewModel.expansionState.value.toIdleSceneKey(),
+                viewModel.expansionState.toIdleSceneKey(),
                 transitions =
                     transitions {
                         from(QuickQuickSettings, QuickSettings) {
@@ -294,7 +301,10 @@
         }
 
         LaunchedEffect(Unit) {
-            synchronizeQsState(sceneState, viewModel.expansionState.map { it.progress })
+            synchronizeQsState(
+                sceneState,
+                snapshotFlow { viewModel.expansionState }.map { it.progress },
+            )
         }
 
         SceneTransitionLayout(state = sceneState, modifier = Modifier.fillMaxSize()) {
@@ -315,7 +325,7 @@
 
     override fun getQsMinExpansionHeight(): Int {
         // TODO (b/353253277) implement split screen
-        return qqsHeight.value
+        return viewModel.qqsHeight
     }
 
     override fun getDesiredHeight(): Int {
@@ -329,7 +339,7 @@
     }
 
     override fun setHeightOverride(desiredHeight: Int) {
-        viewModel.heightOverrideValue = desiredHeight
+        viewModel.heightOverride = desiredHeight
     }
 
     override fun setHeaderClickable(qsExpansionEnabled: Boolean) {
@@ -349,7 +359,7 @@
     }
 
     override fun setExpanded(qsExpanded: Boolean) {
-        viewModel.isQSExpanded = qsExpanded
+        viewModel.isQsExpanded = qsExpanded
     }
 
     override fun setListening(listening: Boolean) {
@@ -357,7 +367,7 @@
     }
 
     override fun setQsVisible(qsVisible: Boolean) {
-        viewModel.isQSVisible = qsVisible
+        viewModel.isQsVisible = qsVisible
     }
 
     override fun isShowingDetail(): Boolean {
@@ -378,11 +388,10 @@
         headerTranslation: Float,
         squishinessFraction: Float,
     ) {
-        viewModel.qsExpansionValue = qsExpansionFraction
-        viewModel.panelExpansionFractionValue = panelExpansionFraction
-        viewModel.squishinessFractionValue = squishinessFraction
-
-        // TODO(b/353254353) Handle header translation
+        viewModel.setQsExpansionValue(qsExpansionFraction)
+        viewModel.panelExpansionFraction = panelExpansionFraction
+        viewModel.squishinessFraction = squishinessFraction
+        viewModel.proposedTranslation = headerTranslation
     }
 
     override fun setHeaderListening(listening: Boolean) {
@@ -402,7 +411,7 @@
     }
 
     override fun getHeightDiff(): Int {
-        return 0 // For now TODO(b/353254353)
+        return viewModel.heightDiff
     }
 
     override fun getHeader(): View? {
@@ -415,8 +424,8 @@
         // TODO (b/353253280)
     }
 
-    override fun setInSplitShade(shouldTranslate: Boolean) {
-        // TODO (b/356435605)
+    override fun setInSplitShade(isInSplitShade: Boolean) {
+        viewModel.isInSplitShade = isInSplitShade
     }
 
     override fun setTransitionToFullShadeProgress(
@@ -425,9 +434,9 @@
         qsSquishinessFraction: Float,
     ) {
         viewModel.isTransitioningToFullShade = isTransitioningToFullShade
-        viewModel.lockscreenToShadeProgressValue = qsTransitionFraction
+        viewModel.lockscreenToShadeProgress = qsTransitionFraction
         if (isTransitioningToFullShade) {
-            viewModel.squishinessFractionValue = qsSquishinessFraction
+            viewModel.squishinessFraction = qsSquishinessFraction
         }
     }
 
@@ -452,7 +461,7 @@
     }
 
     override fun isFullyCollapsed(): Boolean {
-        return viewModel.qsExpansionValue <= 0f
+        return viewModel.isQsFullyCollapsed
     }
 
     override fun setCollapsedMediaVisibilityChangedListener(listener: Consumer<Boolean>?) {
@@ -464,11 +473,11 @@
     }
 
     override fun setOverScrollAmount(overScrollAmount: Int) {
-        super.setOverScrollAmount(overScrollAmount)
+        viewModel.overScrollAmount = overScrollAmount
     }
 
     override fun setIsNotificationPanelFullWidth(isFullWidth: Boolean) {
-        viewModel.isSmallScreenValue = isFullWidth
+        viewModel.isSmallScreen = isFullWidth
     }
 
     override fun getHeaderTop(): Int {
@@ -522,8 +531,8 @@
 
     @Composable
     private fun SceneScope.QuickQuickSettingsElement() {
-        val qqsPadding by viewModel.qqsHeaderHeight.collectAsStateWithLifecycle()
-        val bottomPadding = dimensionResource(id = R.dimen.qqs_layout_padding_bottom)
+        val qqsPadding = viewModel.qqsHeaderHeight
+        val bottomPadding = viewModel.qqsBottomPadding
         DisposableEffect(Unit) {
             qqsVisible.value = true
 
@@ -553,14 +562,13 @@
                         .approachLayout(isMeasurementApproachInProgress = { squishiness < 1f }) {
                             measurable,
                             constraints ->
-                            qqsHeight.value = lookaheadSize.height
+                            viewModel.qqsHeight = lookaheadSize.height
                             val placeable = measurable.measure(constraints)
                             layout(placeable.width, placeable.height) { placeable.place(0, 0) }
                         }
-                        .padding(top = { qqsPadding }, bottom = { bottomPadding.roundToPx() })
+                        .padding(top = { qqsPadding }, bottom = { bottomPadding })
             ) {
-                val qsEnabled by viewModel.qsEnabled.collectAsStateWithLifecycle()
-                if (qsEnabled) {
+                if (viewModel.isQsEnabled) {
                     QuickQuickSettings(
                         viewModel = viewModel.containerViewModel.quickQuickSettingsViewModel,
                         modifier =
@@ -583,7 +591,7 @@
 
     @Composable
     private fun SceneScope.QuickSettingsElement() {
-        val qqsPadding by viewModel.qqsHeaderHeight.collectAsStateWithLifecycle()
+        val qqsPadding = viewModel.qqsHeaderHeight
         val qsExtraPadding = dimensionResource(R.dimen.qs_panel_padding_top)
         Column(
             modifier =
@@ -591,8 +599,7 @@
                     stringResource(id = R.string.accessibility_quick_settings_collapse)
                 )
         ) {
-            val qsEnabled by viewModel.qsEnabled.collectAsStateWithLifecycle()
-            if (qsEnabled) {
+            if (viewModel.isQsEnabled) {
                 Box(
                     modifier =
                         Modifier.element(ElementKeys.QuickSettingsContent).fillMaxSize().weight(1f)
@@ -602,7 +609,17 @@
                         onDispose { lifecycleScope.launch { scrollState.scrollTo(0) } }
                     }
 
-                    Column(modifier = Modifier.verticalScroll(scrollState)) {
+                    Column(
+                        modifier =
+                            Modifier.offset {
+                                    IntOffset(
+                                        x = 0,
+                                        y = viewModel.qsScrollTranslationY.fastRoundToInt(),
+                                    )
+                                }
+                                .onSizeChanged { viewModel.qsScrollHeight = it.height }
+                                .verticalScroll(scrollState)
+                    ) {
                         Spacer(
                             modifier = Modifier.height { qqsPadding + qsExtraPadding.roundToPx() }
                         )
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
index 7a8b2c2..e21485b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
@@ -20,19 +20,31 @@
 import android.graphics.Rect
 import androidx.annotation.FloatRange
 import androidx.annotation.VisibleForTesting
+import androidx.compose.runtime.derivedStateOf
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshotFlow
 import androidx.lifecycle.LifecycleCoroutineScope
+import com.android.keyguard.BouncerPanelExpansionCalculator
 import com.android.systemui.Dumpable
+import com.android.systemui.animation.ShadeInterpolation
 import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
+import com.android.systemui.keyguard.shared.model.Edge
+import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.lifecycle.ExclusiveActivatable
+import com.android.systemui.lifecycle.Hydrator
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.qs.FooterActionsController
-import com.android.systemui.qs.composefragment.viewmodel.QSFragmentComposeViewModel.QSExpansionState
 import com.android.systemui.qs.footer.ui.viewmodel.FooterActionsViewModel
 import com.android.systemui.qs.panels.domain.interactor.TileSquishinessInteractor
 import com.android.systemui.qs.panels.ui.viewmodel.PaginatedGridViewModel
 import com.android.systemui.qs.ui.viewmodel.QuickSettingsContainerViewModel
+import com.android.systemui.res.R
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.LargeScreenHeaderHelper
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator
 import com.android.systemui.statusbar.StatusBarState
@@ -47,246 +59,128 @@
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import java.io.PrintWriter
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.channels.awaitClose
-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.coroutineScope
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
-import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 
+@OptIn(ExperimentalCoroutinesApi::class)
 class QSFragmentComposeViewModel
 @AssistedInject
 constructor(
     val containerViewModel: QuickSettingsContainerViewModel,
     @Main private val resources: Resources,
-    private val footerActionsViewModelFactory: FooterActionsViewModel.Factory,
+    footerActionsViewModelFactory: FooterActionsViewModel.Factory,
     private val footerActionsController: FooterActionsController,
     private val sysuiStatusBarStateController: SysuiStatusBarStateController,
-    private val deviceEntryInteractor: DeviceEntryInteractor,
-    private val disableFlagsRepository: DisableFlagsRepository,
+    deviceEntryInteractor: DeviceEntryInteractor,
+    disableFlagsRepository: DisableFlagsRepository,
+    keyguardTransitionInteractor: KeyguardTransitionInteractor,
     private val largeScreenShadeInterpolator: LargeScreenShadeInterpolator,
-    private val configurationInteractor: ConfigurationInteractor,
+    configurationInteractor: ConfigurationInteractor,
     private val largeScreenHeaderHelper: LargeScreenHeaderHelper,
     private val squishinessInteractor: TileSquishinessInteractor,
     private val paginatedGridViewModel: PaginatedGridViewModel,
     @Assisted private val lifecycleScope: LifecycleCoroutineScope,
 ) : Dumpable, ExclusiveActivatable() {
+
+    private val hydrator = Hydrator("QSFragmentComposeViewModel.hydrator")
+
     val footerActionsViewModel =
         footerActionsViewModelFactory.create(lifecycleScope).also {
             lifecycleScope.launch { footerActionsController.init() }
         }
 
-    private val _qsBounds = MutableStateFlow(Rect())
+    var isQsExpanded by mutableStateOf(false)
 
-    private val _qsExpanded = MutableStateFlow(false)
-    var isQSExpanded: Boolean
-        get() = _qsExpanded.value
-        set(value) {
-            _qsExpanded.value = value
-        }
-
-    private val _qsVisible = MutableStateFlow(false)
-    val qsVisible = _qsVisible.asStateFlow()
-    var isQSVisible: Boolean
-        get() = qsVisible.value
-        set(value) {
-            _qsVisible.value = value
-        }
+    var isQsVisible by mutableStateOf(false)
 
     // This can only be negative if undefined (in which case it will be -1f), else it will be
     // in [0, 1]. In some cases, it could be set back to -1f internally to indicate that it's
     // different to every value in [0, 1].
-    @FloatRange(from = -1.0, to = 1.0) private val _qsExpansion = MutableStateFlow(-1f)
-    var qsExpansionValue: Float
-        get() = _qsExpansion.value
-        set(value) {
-            if (value < 0f) {
-                _qsExpansion.value = -1f
-            }
-            _qsExpansion.value = value.coerceIn(0f, 1f)
+    private var qsExpansion by mutableStateOf(-1f)
+
+    fun setQsExpansionValue(value: Float) {
+        if (value < 0f) {
+            qsExpansion = -1f
+        } else {
+            qsExpansion = value.coerceIn(0f, 1f)
         }
+    }
 
-    private val _panelFraction = MutableStateFlow(0f)
-    var panelExpansionFractionValue: Float
-        get() = _panelFraction.value
-        set(value) {
-            _panelFraction.value = value
-        }
+    val isQsFullyCollapsed by derivedStateOf { qsExpansion <= 0f }
 
-    private val _squishinessFraction = MutableStateFlow(1f)
-    var squishinessFractionValue: Float
-        get() = _squishinessFraction.value
-        set(value) {
-            _squishinessFraction.value = value
-        }
+    var panelExpansionFraction by mutableStateOf(0f)
 
-    val qqsHeaderHeight =
-        configurationInteractor.onAnyConfigurationChange
-            .map {
-                if (LargeScreenUtils.shouldUseLargeScreenShadeHeader(resources)) {
-                    0
-                } else {
-                    largeScreenHeaderHelper.getLargeScreenHeaderHeight()
-                }
-            }
-            .stateIn(lifecycleScope, SharingStarted.WhileSubscribed(), 0)
+    var squishinessFraction by mutableStateOf(1f)
 
-    private val _headerAnimating = MutableStateFlow(false)
+    val qqsHeaderHeight by
+        hydrator.hydratedStateOf(
+            traceName = "qqsHeaderHeight",
+            initialValue = 0,
+            source =
+                configurationInteractor.onAnyConfigurationChange.map {
+                    if (LargeScreenUtils.shouldUseLargeScreenShadeHeader(resources)) {
+                        0
+                    } else {
+                        largeScreenHeaderHelper.getLargeScreenHeaderHeight()
+                    }
+                },
+        )
 
-    private val _stackScrollerOverscrolling = MutableStateFlow(false)
-    var isStackScrollerOverscrolling: Boolean
-        get() = _stackScrollerOverscrolling.value
-        set(value) {
-            _stackScrollerOverscrolling.value = value
-        }
+    val qqsBottomPadding by
+        hydrator.hydratedStateOf(
+            traceName = "qqsBottomPadding",
+            initialValue = resources.getDimensionPixelSize(R.dimen.qqs_layout_padding_bottom),
+            source = configurationInteractor.dimensionPixelSize(R.dimen.qqs_layout_padding_bottom),
+        )
+
+    // Starting with a non-zero value makes it so that it has a non-zero height on first expansion
+    // This is important for `QuickSettingsControllerImpl.mMinExpansionHeight` to detect a "change".
+    var qqsHeight by mutableStateOf(1)
+
+    var qsScrollHeight by mutableStateOf(0)
+
+    val heightDiff: Int
+        get() = qsScrollHeight - qqsHeight + qqsBottomPadding
+
+    var isStackScrollerOverscrolling by mutableStateOf(false)
+
+    var proposedTranslation by mutableStateOf(0f)
 
     /**
      * Whether QS is enabled by policy. This is normally true, except when it's disabled by some
      * policy. See [DisableFlagsRepository].
      */
-    val qsEnabled =
-        disableFlagsRepository.disableFlags
-            .map { it.isQuickSettingsEnabled() }
-            .stateIn(
-                lifecycleScope,
-                SharingStarted.WhileSubscribed(),
-                disableFlagsRepository.disableFlags.value.isQuickSettingsEnabled(),
-            )
+    val isQsEnabled by
+        hydrator.hydratedStateOf(
+            traceName = "isQsEnabled",
+            initialValue = disableFlagsRepository.disableFlags.value.isQuickSettingsEnabled(),
+            source = disableFlagsRepository.disableFlags.map { it.isQuickSettingsEnabled() },
+        )
 
-    private val _keyguardAndExpanded = MutableStateFlow(false)
+    var isInSplitShade by mutableStateOf(false)
 
-    /**
-     * Tracks the current [StatusBarState]. It will switch early if the upcoming state is
-     * [StatusBarState.KEYGUARD]
-     */
-    @get:VisibleForTesting
-    val statusBarState =
-        conflatedCallbackFlow {
-                val callback =
-                    object : StatusBarStateController.StateListener {
-                        override fun onStateChanged(newState: Int) {
-                            trySend(newState)
-                        }
+    var isTransitioningToFullShade by mutableStateOf(false)
 
-                        override fun onUpcomingStateChanged(upcomingState: Int) {
-                            if (upcomingState == StatusBarState.KEYGUARD) {
-                                trySend(upcomingState)
-                            }
-                        }
-                    }
-                sysuiStatusBarStateController.addCallback(callback)
+    var lockscreenToShadeProgress by mutableStateOf(0f)
 
-                awaitClose { sysuiStatusBarStateController.removeCallback(callback) }
-            }
-            .onStart { emit(sysuiStatusBarStateController.state) }
-            .stateIn(
-                lifecycleScope,
-                SharingStarted.WhileSubscribed(),
-                sysuiStatusBarStateController.state,
-            )
+    var isSmallScreen by mutableStateOf(false)
 
-    private val isKeyguardState =
-        statusBarState
-            .map { it == StatusBarState.KEYGUARD }
-            .stateIn(
-                lifecycleScope,
-                SharingStarted.WhileSubscribed(),
-                statusBarState.value == StatusBarState.KEYGUARD,
-            )
+    var heightOverride by mutableStateOf(-1)
 
-    private val _viewHeight = MutableStateFlow(0)
-
-    private val _headerTranslation = MutableStateFlow(0f)
-
-    private val _inSplitShade = MutableStateFlow(false)
-    var isInSplitShade: Boolean
-        get() = _inSplitShade.value
-        set(value) {
-            _inSplitShade.value = value
+    val expansionState by derivedStateOf {
+        if (forceQs) {
+            QSExpansionState(1f)
+        } else {
+            QSExpansionState(qsExpansion.coerceIn(0f, 1f))
         }
+    }
 
-    private val _transitioningToFullShade = MutableStateFlow(false)
-    var isTransitioningToFullShade: Boolean
-        get() = _transitioningToFullShade.value
-        set(value) {
-            _transitioningToFullShade.value = value
-        }
-
-    private val isBypassEnabled = deviceEntryInteractor.isBypassEnabled
-
-    private val showCollapsedOnKeyguard =
-        combine(
-                isBypassEnabled,
-                _transitioningToFullShade,
-                _inSplitShade,
-                ::calculateShowCollapsedOnKeyguard,
-            )
-            .stateIn(
-                lifecycleScope,
-                SharingStarted.WhileSubscribed(),
-                calculateShowCollapsedOnKeyguard(
-                    isBypassEnabled.value,
-                    isTransitioningToFullShade,
-                    isInSplitShade,
-                ),
-            )
-
-    private val _lockscreenToShadeProgress = MutableStateFlow(0.0f)
-    var lockscreenToShadeProgressValue: Float
-        get() = _lockscreenToShadeProgress.value
-        set(value) {
-            _lockscreenToShadeProgress.value = value
-        }
-
-    private val _overscrolling = MutableStateFlow(false)
-
-    private val _isSmallScreen = MutableStateFlow(false)
-    var isSmallScreenValue: Boolean
-        get() = _isSmallScreen.value
-        set(value) {
-            _isSmallScreen.value = value
-        }
-
-    private val _shouldUpdateMediaSquishiness = MutableStateFlow(false)
-
-    private val _heightOverride = MutableStateFlow(-1)
-    val heightOverride = _heightOverride.asStateFlow()
-    var heightOverrideValue: Int
-        get() = heightOverride.value
-        set(value) {
-            _heightOverride.value = value
-        }
-
-    private val forceQS =
-        combine(
-                _qsExpanded,
-                _stackScrollerOverscrolling,
-                isKeyguardState,
-                showCollapsedOnKeyguard,
-                ::calculateForceQs,
-            )
-            .stateIn(
-                lifecycleScope,
-                SharingStarted.WhileSubscribed(),
-                calculateForceQs(
-                    isQSExpanded,
-                    isStackScrollerOverscrolling,
-                    isKeyguardState.value,
-                    showCollapsedOnKeyguard.value,
-                ),
-            )
-
-    val expansionState: StateFlow<QSExpansionState> =
-        combine(_qsExpansion, forceQS, ::calculateExpansionState)
-            .stateIn(
-                lifecycleScope,
-                SharingStarted.WhileSubscribed(),
-                calculateExpansionState(_qsExpansion.value, forceQS.value),
-            )
+    val isQsFullyExpanded by derivedStateOf { expansionState.progress >= 1f }
 
     /**
      * Accessibility action for collapsing/expanding QS. The provided runnable is responsible for
@@ -297,40 +191,181 @@
     val inFirstPage: Boolean
         get() = paginatedGridViewModel.inFirstPage
 
-    override suspend fun onActivated(): Nothing {
-        hydrateSquishinessInteractor()
+    var overScrollAmount by mutableStateOf(0)
+
+    val viewTranslationY by derivedStateOf {
+        if (isOverscrolling) {
+            overScrollAmount.toFloat()
+        } else {
+            if (onKeyguardAndExpanded) {
+                translationScaleY * qqsHeight
+            } else {
+                headerTranslation
+            }
+        }
     }
 
-    private suspend fun hydrateSquishinessInteractor(): Nothing {
-        _squishinessFraction.collect {
-            squishinessInteractor.setSquishinessValue(it.constrainSquishiness())
+    val qsScrollTranslationY by derivedStateOf {
+        val panelTranslationY = translationScaleY * heightDiff
+        if (onKeyguardAndExpanded) panelTranslationY else 0f
+    }
+
+    val viewAlpha by derivedStateOf {
+        when {
+            isInBouncerTransit ->
+                BouncerPanelExpansionCalculator.aboutToShowBouncerProgress(alphaProgress)
+            isKeyguardState -> alphaProgress
+            isSmallScreen -> ShadeInterpolation.getContentAlpha(alphaProgress)
+            else -> largeScreenShadeInterpolator.getQsAlpha(alphaProgress)
         }
     }
 
+    private var qsBounds by mutableStateOf(Rect())
+
+    private val constrainedSquishinessFraction: Float
+        get() = squishinessFraction.constrainSquishiness()
+
+    private var _headerAnimating by mutableStateOf(false)
+
+    /**
+     * Tracks the current [StatusBarState]. It will switch early if the upcoming state is
+     * [StatusBarState.KEYGUARD]
+     */
+    @get:VisibleForTesting
+    val statusBarState by
+        hydrator.hydratedStateOf(
+            traceName = "statusBarState",
+            initialValue = sysuiStatusBarStateController.state,
+            source =
+                conflatedCallbackFlow {
+                        val callback =
+                            object : StatusBarStateController.StateListener {
+                                override fun onStateChanged(newState: Int) {
+                                    trySend(newState)
+                                }
+
+                                override fun onUpcomingStateChanged(upcomingState: Int) {
+                                    if (upcomingState == StatusBarState.KEYGUARD) {
+                                        trySend(upcomingState)
+                                    }
+                                }
+                            }
+                        sysuiStatusBarStateController.addCallback(callback)
+
+                        awaitClose { sysuiStatusBarStateController.removeCallback(callback) }
+                    }
+                    .onStart { emit(sysuiStatusBarStateController.state) },
+        )
+
+    private val isKeyguardState: Boolean
+        get() = statusBarState == StatusBarState.KEYGUARD
+
+    private var viewHeight by mutableStateOf(0)
+
+    private val isBypassEnabled by
+        hydrator.hydratedStateOf(
+            traceName = "isBypassEnabled",
+            source = deviceEntryInteractor.isBypassEnabled,
+        )
+
+    private val showCollapsedOnKeyguard by derivedStateOf {
+        isBypassEnabled || (isTransitioningToFullShade && !isInSplitShade)
+    }
+
+    private val onKeyguardAndExpanded: Boolean
+        get() = isKeyguardState && !showCollapsedOnKeyguard
+
+    private val isOverscrolling: Boolean
+        get() = overScrollAmount != 0
+
+    private var shouldUpdateMediaSquishiness by mutableStateOf(false)
+
+    private val forceQs by derivedStateOf {
+        (isQsExpanded || isStackScrollerOverscrolling) &&
+            (isKeyguardState && !showCollapsedOnKeyguard)
+    }
+
+    private val translationScaleY: Float
+        get() = ((qsExpansion - 1) * (if (isInSplitShade) 1f else SHORT_PARALLAX_AMOUNT))
+
+    private val headerTranslation by derivedStateOf {
+        if (isTransitioningToFullShade) 0f else proposedTranslation
+    }
+
+    private val alphaProgress by derivedStateOf {
+        when {
+            isSmallScreen -> 1f
+            isInSplitShade ->
+                if (isTransitioningToFullShade || isKeyguardState) {
+                    lockscreenToShadeProgress
+                } else {
+                    panelExpansionFraction
+                }
+            isTransitioningToFullShade -> lockscreenToShadeProgress
+            else -> panelExpansionFraction
+        }
+    }
+
+    private val isInBouncerTransit by
+        hydrator.hydratedStateOf(
+            traceName = "isInBouncerTransit",
+            initialValue = false,
+            source =
+                keyguardTransitionInteractor.isInTransition(
+                    Edge.create(to = Scenes.Bouncer),
+                    Edge.create(to = KeyguardState.PRIMARY_BOUNCER),
+                ),
+        )
+
+    override suspend fun onActivated(): Nothing {
+        coroutineScope {
+            launch { hydrateSquishinessInteractor() }
+            launch { hydrator.activate() }
+            awaitCancellation()
+        }
+    }
+
+    private suspend fun hydrateSquishinessInteractor() {
+        snapshotFlow { constrainedSquishinessFraction }
+            .collect { squishinessInteractor.setSquishinessValue(it) }
+    }
+
     override fun dump(pw: PrintWriter, args: Array<out String>) {
         pw.asIndenting().run {
             printSection("Quick Settings state") {
-                println("isQSExpanded", isQSExpanded)
-                println("isQSVisible", isQSVisible)
-                println("isQSEnabled", qsEnabled.value)
+                println("isQSExpanded", isQsExpanded)
+                println("isQSVisible", isQsVisible)
+                println("isQSEnabled", isQsEnabled)
                 println("isCustomizing", containerViewModel.editModeViewModel.isEditing.value)
             }
             printSection("Expansion state") {
-                println("qsExpansion", qsExpansionValue)
-                println("panelExpansionFraction", panelExpansionFractionValue)
-                println("squishinessFraction", squishinessFractionValue)
-                println("expansionState", expansionState.value)
-                println("forceQS", forceQS.value)
+                println("qsExpansion", qsExpansion)
+                println("panelExpansionFraction", panelExpansionFraction)
+                println("squishinessFraction", squishinessFraction)
+                println("proposedTranslation", proposedTranslation)
+                println("expansionState", expansionState)
+                println("forceQS", forceQs)
+                printSection("Derived values") {
+                    println("headerTranslation", headerTranslation)
+                    println("translationScaleY", translationScaleY)
+                    println("viewTranslationY", viewTranslationY)
+                    println("qsScrollTranslationY", qsScrollTranslationY)
+                    println("viewAlpha", viewAlpha)
+                }
             }
             printSection("Shade state") {
                 println("stackOverscrolling", isStackScrollerOverscrolling)
-                println("statusBarState", StatusBarState.toString(statusBarState.value))
-                println("isKeyguardState", isKeyguardState.value)
-                println("isSmallScreen", isSmallScreenValue)
-                println("heightOverride", "${heightOverrideValue}px")
-                println("qqsHeaderHeight", "${qqsHeaderHeight.value}px")
+                println("overscrollAmount", overScrollAmount)
+                println("statusBarState", StatusBarState.toString(statusBarState))
+                println("isKeyguardState", isKeyguardState)
+                println("isSmallScreen", isSmallScreen)
+                println("heightOverride", "${heightOverride}px")
+                println("qqsHeaderHeight", "${qqsHeaderHeight}px")
+                println("qqsBottomPadding", "${qqsBottomPadding}px")
                 println("isSplitShade", isInSplitShade)
-                println("showCollapsedOnKeyguard", showCollapsedOnKeyguard.value)
+                println("showCollapsedOnKeyguard", showCollapsedOnKeyguard)
+                println("qqsHeight", "${qqsHeight}px")
+                println("qsScrollHeight", "${qsScrollHeight}px")
             }
         }
     }
@@ -340,7 +375,7 @@
         fun create(lifecycleScope: LifecycleCoroutineScope): QSFragmentComposeViewModel
     }
 
-    // In the future, this will have other relevant elements like squishiness.
+    // In the future, this may have other relevant elements.
     data class QSExpansionState(@FloatRange(0.0, 1.0) val progress: Float)
 }
 
@@ -348,30 +383,4 @@
     return (0.1f + this * 0.9f).coerceIn(0f, 1f)
 }
 
-// Helper methods for combining flows.
-
-private fun calculateExpansionState(expansion: Float, forceQs: Boolean): QSExpansionState {
-    return if (forceQs) {
-        QSExpansionState(1f)
-    } else {
-        QSExpansionState(expansion.coerceIn(0f, 1f))
-    }
-}
-
-private fun calculateForceQs(
-    isQSExpanded: Boolean,
-    isStackOverScrolling: Boolean,
-    isKeyguardShowing: Boolean,
-    shouldShowCollapsedOnKeyguard: Boolean,
-): Boolean {
-    return (isQSExpanded || isStackOverScrolling) &&
-        (isKeyguardShowing && !shouldShowCollapsedOnKeyguard)
-}
-
-private fun calculateShowCollapsedOnKeyguard(
-    isBypassEnabled: Boolean,
-    isTransitioningToFullShade: Boolean,
-    isInSplitShade: Boolean,
-): Boolean {
-    return isBypassEnabled || (isTransitioningToFullShade && !isInSplitShade)
-}
+private val SHORT_PARALLAX_AMOUNT = 0.1f
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 31e867e..ef30cbf 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
@@ -23,6 +23,8 @@
 import com.android.systemui.qs.panels.data.repository.DefaultLargeTilesRepositoryImpl
 import com.android.systemui.qs.panels.data.repository.GridLayoutTypeRepository
 import com.android.systemui.qs.panels.data.repository.GridLayoutTypeRepositoryImpl
+import com.android.systemui.qs.panels.domain.interactor.EditTilesResetInteractor
+import com.android.systemui.qs.panels.domain.interactor.SizedTilesResetInteractor
 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.PaginatedGridLayoutType
@@ -53,6 +55,9 @@
     @Binds
     fun bindGridLayoutTypeRepository(impl: GridLayoutTypeRepositoryImpl): GridLayoutTypeRepository
 
+    @Binds
+    fun bindEditTilesResetInteractor(impl: SizedTilesResetInteractor): EditTilesResetInteractor
+
     @Binds fun bindIconTilesViewModel(impl: IconTilesViewModelImpl): IconTilesViewModel
 
     @Binds fun bindQSColumnsViewModel(impl: QSColumnsSizeViewModelImpl): QSColumnsViewModel
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractor.kt
new file mode 100644
index 0000000..ee38dfb
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractor.kt
@@ -0,0 +1,47 @@
+/*
+ * 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.domain.interactor
+
+import com.android.systemui.lifecycle.ExclusiveActivatable
+import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+/** Interactor to resize QS tiles down to icons when removed from the current tiles. */
+class DynamicIconTilesInteractor
+@AssistedInject
+constructor(
+    private val iconTilesInteractor: IconTilesInteractor,
+    private val currentTilesInteractor: CurrentTilesInteractor,
+) : ExclusiveActivatable() {
+
+    override suspend fun onActivated(): Nothing {
+        currentTilesInteractor.currentTiles.collect { currentTiles ->
+            // Only current tiles can be resized, so observe the current tiles and find the
+            // intersection between them and the large tiles.
+            val newLargeTiles =
+                iconTilesInteractor.largeTilesSpecs.value intersect
+                    currentTiles.map { it.spec }.toSet()
+            iconTilesInteractor.setLargeTiles(newLargeTiles)
+        }
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): DynamicIconTilesInteractor
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/EditTilesResetInteractor.kt
similarity index 69%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt
rename to packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/EditTilesResetInteractor.kt
index 3190171..b523897 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/EditTilesResetInteractor.kt
@@ -14,15 +14,9 @@
  * limitations under the License.
  */
 
-package com.android.systemui.keyboard.shortcut
+package com.android.systemui.qs.panels.domain.interactor
 
-import android.content.Intent
-
-class FakeShortcutHelperStartActivity : (Intent) -> Unit {
-
-    val startIntents = mutableListOf<Intent>()
-
-    override fun invoke(intent: Intent) {
-        startIntents += intent
-    }
+/** Interactor for resetting QS tiles to the default state. */
+interface EditTilesResetInteractor {
+    fun reset()
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractor.kt
index fc59a50..ec61a0d 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/IconTilesInteractor.kt
@@ -27,7 +27,6 @@
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
 
@@ -36,34 +35,27 @@
 class IconTilesInteractor
 @Inject
 constructor(
-    repo: DefaultLargeTilesRepository,
+    private val repo: DefaultLargeTilesRepository,
     private val currentTilesInteractor: CurrentTilesInteractor,
     private val preferencesInteractor: QSPreferencesInteractor,
     @PanelsLog private val logBuffer: LogBuffer,
     @Application private val applicationScope: CoroutineScope,
 ) {
-
     val largeTilesSpecs =
-        combine(preferencesInteractor.largeTilesSpecs, currentTilesInteractor.currentTiles) {
-                largeTiles,
-                currentTiles ->
-                if (currentTiles.isEmpty()) {
-                    largeTiles
-                } else {
-                    // Only current tiles can be resized, so observe the current tiles and find the
-                    // intersection between them and the large tiles.
-                    val newLargeTiles = largeTiles intersect currentTiles.map { it.spec }.toSet()
-                    if (newLargeTiles != largeTiles) {
-                        preferencesInteractor.setLargeTilesSpecs(newLargeTiles)
-                    }
-                    newLargeTiles
-                }
-            }
+        preferencesInteractor.largeTilesSpecs
             .onEach { logChange(it) }
             .stateIn(applicationScope, SharingStarted.Eagerly, repo.defaultLargeTiles)
 
     fun isIconTile(spec: TileSpec): Boolean = !largeTilesSpecs.value.contains(spec)
 
+    fun setLargeTiles(specs: Set<TileSpec>) {
+        preferencesInteractor.setLargeTilesSpecs(specs)
+    }
+
+    fun resetToDefault() {
+        preferencesInteractor.setLargeTilesSpecs(repo.defaultLargeTiles)
+    }
+
     fun resize(spec: TileSpec, toIcon: Boolean) {
         if (!isCurrent(spec)) {
             return
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractor.kt
new file mode 100644
index 0000000..a402587
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractor.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.qs.panels.domain.interactor
+
+import com.android.internal.logging.UiEventLogger
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.qs.QSEditEvent
+import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor
+import javax.inject.Inject
+
+/**
+ * This implementation of [EditTilesResetInteractor] resets both the current tiles and the sizes to
+ * the default state.
+ */
+@SysUISingleton
+class SizedTilesResetInteractor
+@Inject
+constructor(
+    private val currentTilesInteractor: CurrentTilesInteractor,
+    private val iconTilesInteractor: IconTilesInteractor,
+    private val uiEventLogger: UiEventLogger,
+) : EditTilesResetInteractor {
+    override fun reset() {
+        uiEventLogger.log(QSEditEvent.QS_EDIT_RESET)
+        currentTilesInteractor.resetTiles()
+        iconTilesInteractor.resetToDefault()
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/EditMode.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/EditMode.kt
index 1674865..e990d9d 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/EditMode.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/EditMode.kt
@@ -26,10 +26,7 @@
 import com.android.systemui.qs.panels.ui.viewmodel.EditModeViewModel
 
 @Composable
-fun EditMode(
-    viewModel: EditModeViewModel,
-    modifier: Modifier = Modifier,
-) {
+fun EditMode(viewModel: EditModeViewModel, modifier: Modifier = Modifier) {
     val gridLayout by viewModel.gridLayout.collectAsStateWithLifecycle()
     val tiles by viewModel.tiles.collectAsStateWithLifecycle(emptyList())
 
@@ -44,6 +41,7 @@
             viewModel::addTile,
             viewModel::removeTile,
             viewModel::setTiles,
+            viewModel::stopEditing,
         )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
index 0c02b40..0d37581 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
@@ -41,6 +41,7 @@
         onAddTile: (TileSpec, Int) -> Unit,
         onRemoveTile: (TileSpec) -> Unit,
         onSetTiles: (List<TileSpec>) -> Unit,
+        onStopEditing: () -> Unit,
     )
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/EditTile.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/EditTile.kt
index 5c2a2bd..d2ec958 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/EditTile.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/EditTile.kt
@@ -53,11 +53,19 @@
 import androidx.compose.foundation.shape.RoundedCornerShape
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.automirrored.filled.ArrowBack
+import androidx.compose.material.icons.filled.ArrowBack
 import androidx.compose.material.icons.filled.Clear
+import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.Icon
+import androidx.compose.material3.IconButton
 import androidx.compose.material3.LocalContentColor
 import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Scaffold
 import androidx.compose.material3.Text
+import androidx.compose.material3.TextButton
+import androidx.compose.material3.TopAppBar
+import androidx.compose.material3.TopAppBarDefaults
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.LaunchedEffect
@@ -133,6 +141,31 @@
 
 object TileType
 
+@OptIn(ExperimentalMaterial3Api::class)
+@Composable
+private fun EditModeTopBar(onStopEditing: () -> Unit, onReset: (() -> Unit)?) {
+    TopAppBar(
+        colors = TopAppBarDefaults.topAppBarColors(containerColor = Color.Black),
+        title = { Text(text = stringResource(id = R.string.qs_edit)) },
+        navigationIcon = {
+            IconButton(onClick = onStopEditing) {
+                Icon(
+                    Icons.AutoMirrored.Filled.ArrowBack,
+                    contentDescription =
+                        stringResource(id = com.android.internal.R.string.action_bar_up_description),
+                )
+            }
+        },
+        actions = {
+            if (onReset != null) {
+                TextButton(onClick = onReset) {
+                    Text(stringResource(id = com.android.internal.R.string.reset))
+                }
+            }
+        },
+    )
+}
+
 @Composable
 fun DefaultEditTileGrid(
     listState: EditTileListState,
@@ -142,6 +175,8 @@
     onRemoveTile: (TileSpec) -> Unit,
     onSetTiles: (List<TileSpec>) -> Unit,
     onResize: (TileSpec, toIcon: Boolean) -> Unit,
+    onStopEditing: () -> Unit,
+    onReset: (() -> Unit)?,
 ) {
     val currentListState by rememberUpdatedState(listState)
     val selectionState =
@@ -152,53 +187,71 @@
                 currentListState.isIcon(spec)?.let { onResize(spec, it) }
             },
         )
+    val reset: (() -> Unit)? =
+        if (onReset != null) {
+            {
+                selectionState.unSelect()
+                onReset()
+            }
+        } else {
+            null
+        }
 
-    CompositionLocalProvider(LocalOverscrollConfiguration provides null) {
-        Column(
-            verticalArrangement =
-                spacedBy(dimensionResource(id = R.dimen.qs_label_container_margin)),
-            modifier = modifier.fillMaxSize().verticalScroll(rememberScrollState()),
-        ) {
-            AnimatedContent(
-                targetState = listState.dragInProgress,
-                modifier = Modifier.wrapContentSize(),
-                label = "",
-            ) { dragIsInProgress ->
-                EditGridHeader(Modifier.dragAndDropRemoveZone(listState, onRemoveTile)) {
-                    if (dragIsInProgress) {
-                        RemoveTileTarget()
-                    } else {
-                        Text(text = "Hold and drag to rearrange tiles.")
+    Scaffold(
+        containerColor = Color.Transparent,
+        topBar = { EditModeTopBar(onStopEditing = onStopEditing, onReset = reset) },
+    ) { innerPadding ->
+        CompositionLocalProvider(LocalOverscrollConfiguration provides null) {
+            Column(
+                verticalArrangement =
+                    spacedBy(dimensionResource(id = R.dimen.qs_label_container_margin)),
+                modifier =
+                    modifier
+                        .fillMaxSize()
+                        .verticalScroll(rememberScrollState())
+                        .padding(innerPadding),
+            ) {
+                AnimatedContent(
+                    targetState = listState.dragInProgress,
+                    modifier = Modifier.wrapContentSize(),
+                    label = "",
+                ) { dragIsInProgress ->
+                    EditGridHeader(Modifier.dragAndDropRemoveZone(listState, onRemoveTile)) {
+                        if (dragIsInProgress) {
+                            RemoveTileTarget()
+                        } else {
+                            Text(text = "Hold and drag to rearrange tiles.")
+                        }
                     }
                 }
-            }
 
-            CurrentTilesGrid(listState, selectionState, columns, onResize, onSetTiles)
+                CurrentTilesGrid(listState, selectionState, columns, onResize, onSetTiles)
 
-            // Hide available tiles when dragging
-            AnimatedVisibility(
-                visible = !listState.dragInProgress,
-                enter = fadeIn(),
-                exit = fadeOut(),
-            ) {
-                Column(
-                    verticalArrangement =
-                        spacedBy(dimensionResource(id = R.dimen.qs_label_container_margin)),
-                    modifier = modifier.fillMaxSize(),
+                // Hide available tiles when dragging
+                AnimatedVisibility(
+                    visible = !listState.dragInProgress,
+                    enter = fadeIn(),
+                    exit = fadeOut(),
                 ) {
-                    EditGridHeader { Text(text = "Hold and drag to add tiles.") }
+                    Column(
+                        verticalArrangement =
+                            spacedBy(dimensionResource(id = R.dimen.qs_label_container_margin)),
+                        modifier = modifier.fillMaxSize(),
+                    ) {
+                        EditGridHeader { Text(text = "Hold and drag to add tiles.") }
 
-                    AvailableTileGrid(otherTiles, selectionState, columns, listState)
+                        AvailableTileGrid(otherTiles, selectionState, columns, listState)
+                    }
                 }
-            }
 
-            // Drop zone to remove tiles dragged out of the tile grid
-            Spacer(
-                modifier =
-                    Modifier.fillMaxWidth()
-                        .weight(1f)
-                        .dragAndDropRemoveZone(listState, onRemoveTile)
-            )
+                // Drop zone to remove tiles dragged out of the tile grid
+                Spacer(
+                    modifier =
+                        Modifier.fillMaxWidth()
+                            .weight(1f)
+                            .dragAndDropRemoveZone(listState, onRemoveTile)
+                )
+            }
         }
     }
 }
@@ -269,7 +322,7 @@
                 .border(
                     width = 1.dp,
                     color = MaterialTheme.colorScheme.onBackground.copy(alpha = .5f),
-                    shape = RoundedCornerShape(48.dp),
+                    shape = RoundedCornerShape((TileHeight / 2) + CurrentTilesGridPadding),
                 )
                 .dragAndDropTileList(gridState, { gridContentOffset }, listState) { spec ->
                     onSetTiles(currentListState.tileSpecs())
@@ -313,10 +366,7 @@
                 text = category.label.load() ?: "",
                 fontSize = 20.sp,
                 color = labelColors.label,
-                modifier =
-                    Modifier.fillMaxWidth()
-                        .background(Color.Black)
-                        .padding(start = 16.dp, bottom = 8.dp, top = 8.dp),
+                modifier = Modifier.fillMaxWidth().padding(start = 16.dp, bottom = 8.dp, top = 8.dp),
             )
             tiles.chunked(columns).forEach { row ->
                 Row(
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
index e5c2135..366bc9a 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
@@ -28,6 +28,7 @@
 import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.grid.ui.compose.VerticalSpannedGrid
+import com.android.systemui.lifecycle.rememberViewModel
 import com.android.systemui.qs.panels.shared.model.SizedTileImpl
 import com.android.systemui.qs.panels.ui.compose.PaginatableGridLayout
 import com.android.systemui.qs.panels.ui.compose.bounceableInfo
@@ -35,8 +36,7 @@
 import com.android.systemui.qs.panels.ui.viewmodel.BounceableTileViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.EditTileViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.IconTilesViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.QSColumnsViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.TileSquishinessViewModel
+import com.android.systemui.qs.panels.ui.viewmodel.InfiniteGridViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
 import com.android.systemui.qs.pipeline.shared.TileSpec
 import com.android.systemui.qs.shared.ui.ElementKeys.toElementKey
@@ -48,8 +48,7 @@
 @Inject
 constructor(
     private val iconTilesViewModel: IconTilesViewModel,
-    private val gridSizeViewModel: QSColumnsViewModel,
-    private val squishinessViewModel: TileSquishinessViewModel,
+    private val viewModelFactory: InfiniteGridViewModel.Factory,
 ) : PaginatableGridLayout {
 
     @Composable
@@ -63,11 +62,19 @@
             tiles.forEach { it.startListening(token) }
             onDispose { tiles.forEach { it.stopListening(token) } }
         }
-        val columns by gridSizeViewModel.columns.collectAsStateWithLifecycle()
+        val viewModel =
+            rememberViewModel(traceName = "InfiniteGridLayout.TileGrid") {
+                viewModelFactory.create()
+            }
+        val iconTilesViewModel =
+            rememberViewModel(traceName = "InfiniteGridLayout.TileGrid") {
+                viewModel.dynamicIconTilesViewModelFactory.create()
+            }
+        val columns by viewModel.gridSizeViewModel.columns.collectAsStateWithLifecycle()
         val sizedTiles = tiles.map { SizedTileImpl(it, it.spec.width()) }
         val bounceables =
             remember(sizedTiles) { List(sizedTiles.size) { BounceableTileViewModel() } }
-        val squishiness by squishinessViewModel.squishiness.collectAsStateWithLifecycle()
+        val squishiness by viewModel.squishinessViewModel.squishiness.collectAsStateWithLifecycle()
         val scope = rememberCoroutineScope()
         var cellIndex = 0
 
@@ -98,8 +105,17 @@
         onAddTile: (TileSpec, Int) -> Unit,
         onRemoveTile: (TileSpec) -> Unit,
         onSetTiles: (List<TileSpec>) -> Unit,
+        onStopEditing: () -> Unit,
     ) {
-        val columns by gridSizeViewModel.columns.collectAsStateWithLifecycle()
+        val viewModel =
+            rememberViewModel(traceName = "InfiniteGridLayout.EditTileGrid") {
+                viewModelFactory.create()
+            }
+        val iconTilesViewModel =
+            rememberViewModel(traceName = "InfiniteGridLayout.EditTileGrid") {
+                viewModel.dynamicIconTilesViewModelFactory.create()
+            }
+        val columns by viewModel.gridSizeViewModel.columns.collectAsStateWithLifecycle()
         val largeTiles by iconTilesViewModel.largeTiles.collectAsStateWithLifecycle()
 
         // Non-current tiles should always be displayed as icon tiles.
@@ -123,6 +139,8 @@
             onRemoveTile = onRemoveTile,
             onSetTiles = onSetTiles,
             onResize = iconTilesViewModel::resize,
+            onStopEditing = onStopEditing,
+            onReset = viewModel::showResetDialog,
         )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/dialog/QSResetDialogDelegate.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/dialog/QSResetDialogDelegate.kt
new file mode 100644
index 0000000..03fc425
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/dialog/QSResetDialogDelegate.kt
@@ -0,0 +1,105 @@
+/*
+ * 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.dialog
+
+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.systemui.dagger.SysUISingleton
+import com.android.systemui.dialog.ui.composable.AlertDialogContent
+import com.android.systemui.qs.panels.domain.interactor.EditTilesResetInteractor
+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.util.Assert
+import javax.inject.Inject
+
+@SysUISingleton
+class QSResetDialogDelegate
+@Inject
+constructor(
+    private val sysuiDialogFactory: SystemUIDialogFactory,
+    private val resetInteractor: EditTilesResetInteractor,
+) : SystemUIDialog.Delegate {
+    private var currentDialog: ComponentSystemUIDialog? = null
+
+    override fun createDialog(): SystemUIDialog {
+        Assert.isMainThread()
+        if (currentDialog != null) {
+            Log.d(TAG, "Dialog is already open, dismissing it and creating a new one.")
+            currentDialog?.dismiss()
+        }
+
+        currentDialog =
+            sysuiDialogFactory
+                .create { ResetConfirmationDialog(it) }
+                .also {
+                    it.lifecycle.addObserver(
+                        object : DefaultLifecycleObserver {
+                            override fun onStop(owner: LifecycleOwner) {
+                                Assert.isMainThread()
+                                currentDialog = null
+                            }
+                        }
+                    )
+                }
+        return currentDialog!!
+    }
+
+    @Composable
+    private fun ResetConfirmationDialog(dialog: SystemUIDialog) {
+        AlertDialogContent(
+            title = { Text(text = stringResource(id = R.string.qs_edit_mode_reset_dialog_title)) },
+            content = {
+                Text(text = stringResource(id = R.string.qs_edit_mode_reset_dialog_content))
+            },
+            positiveButton = {
+                PlatformButton(
+                    onClick = {
+                        dialog.dismiss()
+                        resetInteractor.reset()
+                    }
+                ) {
+                    Text(stringResource(id = android.R.string.ok))
+                }
+            },
+            neutralButton = {
+                PlatformOutlinedButton(onClick = { dialog.dismiss() }) {
+                    Text(stringResource(id = android.R.string.cancel))
+                }
+            },
+        )
+    }
+
+    fun showDialog() {
+        if (currentDialog == null) {
+            createDialog()
+        }
+        currentDialog?.show()
+    }
+
+    companion object {
+        private const val TAG = "ResetDialogDelegate"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/DynamicIconTilesViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/DynamicIconTilesViewModel.kt
new file mode 100644
index 0000000..9feaab8
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/DynamicIconTilesViewModel.kt
@@ -0,0 +1,41 @@
+/*
+ * 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.viewmodel
+
+import com.android.systemui.lifecycle.ExclusiveActivatable
+import com.android.systemui.qs.panels.domain.interactor.DynamicIconTilesInteractor
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+/** View model to resize QS tiles down to icons when removed from the current tiles. */
+class DynamicIconTilesViewModel
+@AssistedInject
+constructor(
+    interactorFactory: DynamicIconTilesInteractor.Factory,
+    iconTilesViewModel: IconTilesViewModel,
+) : IconTilesViewModel by iconTilesViewModel, ExclusiveActivatable() {
+    private val interactor = interactorFactory.create()
+
+    override suspend fun onActivated(): Nothing {
+        interactor.activate()
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): DynamicIconTilesViewModel
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModel.kt
index 4a8aa83e..7fe856b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModel.kt
@@ -73,11 +73,7 @@
     val gridLayout: StateFlow<GridLayout> =
         gridLayoutTypeInteractor.layout
             .map { gridLayoutMap[it] ?: defaultGridLayout }
-            .stateIn(
-                applicationScope,
-                SharingStarted.WhileSubscribed(),
-                defaultGridLayout,
-            )
+            .stateIn(applicationScope, SharingStarted.WhileSubscribed(), defaultGridLayout)
 
     /**
      * Flow of view models for each tile that should be visible in edit mode (or empty flow when not
@@ -196,9 +192,4 @@
     fun setTiles(tileSpecs: List<TileSpec>) {
         currentTilesInteractor.setTiles(tileSpecs)
     }
-
-    /** Immediately resets the current tiles to the default list. */
-    fun resetCurrentTilesToDefault() {
-        throw NotImplementedError("This is not supported yet")
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/InfiniteGridViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/InfiniteGridViewModel.kt
new file mode 100644
index 0000000..0d12067
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/InfiniteGridViewModel.kt
@@ -0,0 +1,40 @@
+/*
+ * 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.viewmodel
+
+import com.android.systemui.qs.panels.ui.dialog.QSResetDialogDelegate
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+class InfiniteGridViewModel
+@AssistedInject
+constructor(
+    val dynamicIconTilesViewModelFactory: DynamicIconTilesViewModel.Factory,
+    val gridSizeViewModel: QSColumnsViewModel,
+    val squishinessViewModel: TileSquishinessViewModel,
+    private val resetDialogDelegate: QSResetDialogDelegate,
+) {
+
+    fun showResetDialog() {
+        resetDialogDelegate.showDialog()
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): InfiniteGridViewModel
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/TileSpecRepository.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/TileSpecRepository.kt
index 24b80b8..d94e7cf 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/TileSpecRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/TileSpecRepository.kt
@@ -71,6 +71,9 @@
     /** Prepend the default list of tiles to the current set of tiles */
     suspend fun prependDefault(@UserIdInt userId: Int)
 
+    /** Reset the current set of tiles to the default list of tiles */
+    suspend fun resetToDefault(userId: Int)
+
     companion object {
         /** Position to indicate the end of the list */
         const val POSITION_AT_END = -1
@@ -148,22 +151,24 @@
 
     override suspend fun reconcileRestore(
         restoreData: RestoreData,
-        currentAutoAdded: Set<TileSpec>
+        currentAutoAdded: Set<TileSpec>,
     ) {
         userTileRepositories
             .get(restoreData.userId)
             ?.reconcileRestore(restoreData, currentAutoAdded)
     }
 
-    override suspend fun prependDefault(
-        userId: Int,
-    ) {
+    override suspend fun prependDefault(userId: Int) {
         if (retailModeRepository.inRetailMode) {
             return
         }
         userTileRepositories.get(userId)?.prependDefault()
     }
 
+    override suspend fun resetToDefault(userId: Int) {
+        userTileRepositories.get(userId)?.resetToDefault()
+    }
+
     companion object {
         private const val DELIMITER = TilesSettingConverter.DELIMITER
     }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/UserTileSpecRepository.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/UserTileSpecRepository.kt
index 1f9570a..b0ae1e1 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/UserTileSpecRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/UserTileSpecRepository.kt
@@ -119,14 +119,7 @@
                 .filter { it !is TileSpec.Invalid }
                 .joinToString(DELIMITER, transform = TileSpec::spec)
         withContext(backgroundDispatcher) {
-            secureSettings.putStringForUser(
-                SETTING,
-                toStore,
-                null,
-                false,
-                forUser,
-                true,
-            )
+            secureSettings.putStringForUser(SETTING, toStore, null, false, forUser, true)
         }
     }
 
@@ -172,13 +165,17 @@
         changeEvents.emit(PrependDefault(defaultTiles))
     }
 
+    suspend fun resetToDefault() {
+        changeEvents.emit(ResetToDefault(defaultTiles))
+    }
+
     sealed interface ChangeAction {
         fun apply(currentTiles: List<TileSpec>): List<TileSpec>
     }
 
     private data class AddTile(
         val tileSpec: TileSpec,
-        val position: Int = TileSpecRepository.POSITION_AT_END
+        val position: Int = TileSpecRepository.POSITION_AT_END,
     ) : ChangeAction {
         override fun apply(currentTiles: List<TileSpec>): List<TileSpec> {
             val tilesList = currentTiles.toMutableList()
@@ -199,9 +196,7 @@
         }
     }
 
-    private data class ChangeTiles(
-        val newTiles: List<TileSpec>,
-    ) : ChangeAction {
+    private data class ChangeTiles(val newTiles: List<TileSpec>) : ChangeAction {
         override fun apply(currentTiles: List<TileSpec>): List<TileSpec> {
             val new = newTiles.filter { it !is TileSpec.Invalid }
             return if (new.isNotEmpty()) new else currentTiles
@@ -214,6 +209,12 @@
         }
     }
 
+    private data class ResetToDefault(val defaultTiles: List<TileSpec>) : ChangeAction {
+        override fun apply(currentTiles: List<TileSpec>): List<TileSpec> {
+            return defaultTiles
+        }
+    }
+
     private data class RestoreTiles(
         val restoreData: RestoreData,
         val currentAutoAdded: Set<TileSpec>,
@@ -236,7 +237,7 @@
         fun reconcileTiles(
             currentTiles: List<TileSpec>,
             currentAutoAdded: Set<TileSpec>,
-            restoreData: RestoreData
+            restoreData: RestoreData,
         ): List<TileSpec> {
             val toRestore = restoreData.restoredTiles.toMutableList()
             val freshlyAutoAdded =
@@ -260,8 +261,6 @@
 
     @AssistedFactory
     interface Factory {
-        fun create(
-            userId: Int,
-        ): UserTileSpecRepository
+        fun create(userId: Int): UserTileSpecRepository
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt
index 4a96710..10097d6 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt
@@ -116,6 +116,9 @@
      */
     fun setTiles(specs: List<TileSpec>)
 
+    /** Requests that the list of tiles for the current user is changed to the default list. */
+    fun resetTiles()
+
     fun createTileSync(spec: TileSpec): QSTile?
 
     companion object {
@@ -222,7 +225,7 @@
                                         .TILE_NOT_PRESENT_IN_NEW_USER
                                 } else {
                                     QSPipelineLogger.TileDestroyedReason.TILE_REMOVED
-                                }
+                                },
                             )
                             (entry.value as TileOrNotInstalled.Tile).tile.destroy()
                         }
@@ -245,7 +248,7 @@
                                             tileSpec,
                                             specsToTiles.getValue(tileSpec),
                                             userChanged,
-                                            newUser
+                                            newUser,
                                         ) ?: createTile(tileSpec)
                                     } else {
                                         createTile(tileSpec)
@@ -268,7 +271,7 @@
                     _currentSpecsAndTiles.value = newResolvedTiles
                     logger.logTilesNotInstalled(
                         newTileMap.filter { it.value is TileOrNotInstalled.NotInstalled }.keys,
-                        newUser
+                        newUser,
                     )
                     if (newResolvedTiles.size < minTiles) {
                         // We ended up with not enough tiles (some may be not installed).
@@ -317,6 +320,10 @@
         }
     }
 
+    override fun resetTiles() {
+        scope.launch { tileSpecRepository.resetToDefault(currentUser.value) }
+    }
+
     override fun dump(pw: PrintWriter, args: Array<out String>) {
         pw.println("CurrentTileInteractorImpl:")
         pw.println("User: ${userId.value}")
@@ -384,7 +391,7 @@
                     !qsTile.isAvailable -> {
                         logger.logTileDestroyed(
                             tileSpec,
-                            QSPipelineLogger.TileDestroyedReason.EXISTING_TILE_NOT_AVAILABLE
+                            QSPipelineLogger.TileDestroyedReason.EXISTING_TILE_NOT_AVAILABLE,
                         )
                         qsTile.destroy()
                         null
@@ -409,7 +416,7 @@
                         qsTile.destroy()
                         logger.logTileDestroyed(
                             tileSpec,
-                            QSPipelineLogger.TileDestroyedReason.CUSTOM_TILE_USER_CHANGED
+                            QSPipelineLogger.TileDestroyedReason.CUSTOM_TILE_USER_CHANGED,
                         )
                         null
                     }
@@ -428,7 +435,7 @@
 private data class UserTilesAndComponents(
     val userId: Int,
     val tiles: List<TileSpec>,
-    val installedComponents: Set<ComponentName>
+    val installedComponents: Set<ComponentName>,
 )
 
 private data class DataWithUserChange(
@@ -439,9 +446,4 @@
 )
 
 private fun DataWithUserChange(data: UserTilesAndComponents, userChange: Boolean) =
-    DataWithUserChange(
-        data.userId,
-        data.tiles,
-        data.installedComponents,
-        userChange,
-    )
+    DataWithUserChange(data.userId, data.tiles, data.installedComponents, userChange)
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 3bbe624..cf2db6c 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
@@ -121,7 +121,7 @@
         state?.apply {
             this.state = tileState.activationState.legacyState
             val tileStateIcon = tileState.icon()
-            icon = tileStateIcon?.asQSTileIcon(tileState.iconRes) ?: ResourceIcon.get(ICON_RES_ID)
+            icon = tileStateIcon?.asQSTileIcon() ?: ResourceIcon.get(ICON_RES_ID)
             label = tileLabel
             secondaryLabel = tileState.secondaryLabel
             contentDescription = tileState.contentDescription
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 cc14e71..3e442582 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
@@ -78,14 +78,14 @@
         } else {
             return ModesTileModel(
                 isActivated = activeModes.isAnyActive(),
-                icon = Icon.Resource(ModesTile.ICON_RES_ID, null),
+                icon = context.getDrawable(ModesTile.ICON_RES_ID)!!.asIcon(),
                 iconResId = ModesTile.ICON_RES_ID,
                 activeModes = activeModes.modeNames,
             )
         }
     }
 
-    private data class TileIcon(val icon: Icon, val resId: Int?)
+    private data class TileIcon(val icon: Icon.Loaded, val resId: Int?)
 
     private fun getTileIcon(activeMode: ZenModeInfo?): TileIcon {
         return if (activeMode != null) {
@@ -96,7 +96,7 @@
                 TileIcon(activeMode.icon.drawable.asIcon(), null)
             }
         } else {
-            TileIcon(Icon.Resource(ModesTile.ICON_RES_ID, null), ModesTile.ICON_RES_ID)
+            TileIcon(context.getDrawable(ModesTile.ICON_RES_ID)!!.asIcon(), ModesTile.ICON_RES_ID)
         }
     }
 
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 9c31e32..db48123 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
@@ -21,12 +21,12 @@
 data class ModesTileModel(
     val isActivated: Boolean,
     val activeModes: List<String>,
-    val icon: Icon,
+    val icon: Icon.Loaded,
 
     /**
      * Resource id corresponding to [icon]. Will only be present if it's know to correspond to a
      * resource with a known id in SystemUI (such as resources from `android.R`,
      * `com.android.internal.R`, or `com.android.systemui.res` itself).
      */
-    val iconResId: Int? = null,
+    val iconResId: Int? = null
 )
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 537b56b..69da313 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
@@ -18,9 +18,7 @@
 
 import android.content.res.Resources
 import android.icu.text.MessageFormat
-import android.util.Log
 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
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
@@ -32,30 +30,14 @@
 
 class ModesTileMapper
 @Inject
-constructor(@Main private val resources: Resources, val theme: Resources.Theme) :
-    QSTileDataToStateMapper<ModesTileModel> {
+constructor(
+    @Main private val resources: Resources,
+    val theme: Resources.Theme,
+) : QSTileDataToStateMapper<ModesTileModel> {
     override fun map(config: QSTileConfig, data: ModesTileModel): QSTileState =
         QSTileState.build(resources, theme, config.uiConfig) {
-            val loadedIcon: Icon.Loaded =
-                when (val dataIcon = data.icon) {
-                    is Icon.Resource -> {
-                        if (data.iconResId != dataIcon.res) {
-                            Log.wtf(
-                                "ModesTileMapper",
-                                "Icon.Resource.res & iconResId are not identical",
-                            )
-                        }
-                        iconRes = dataIcon.res
-                        Icon.Loaded(resources.getDrawable(dataIcon.res, theme), null)
-                    }
-                    is Icon.Loaded -> {
-                        iconRes = data.iconResId
-                        dataIcon
-                    }
-                }
-
-            icon = { loadedIcon }
-
+            iconRes = data.iconResId
+            icon = { data.icon }
             activationState =
                 if (data.isActivated) {
                     QSTileState.ActivationState.ACTIVE
@@ -65,7 +47,10 @@
             secondaryLabel = getModesStatus(data, resources)
             contentDescription = "$label. $secondaryLabel"
             supportedActions =
-                setOf(QSTileState.UserAction.CLICK, QSTileState.UserAction.LONG_CLICK)
+                setOf(
+                    QSTileState.UserAction.CLICK,
+                    QSTileState.UserAction.LONG_CLICK,
+                )
             sideViewIcon = QSTileState.SideViewIcon.Chevron
             expandedAccessibilityClass = Button::class
         }
diff --git a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
index 559c263..ce9c441 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
@@ -87,7 +87,6 @@
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractor;
 import com.android.systemui.keyguard.KeyguardUnlockAnimationController;
 import com.android.systemui.keyguard.KeyguardWmStateRefactor;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
@@ -160,8 +159,6 @@
     private final Provider<SceneInteractor> mSceneInteractor;
     private final Provider<ShadeInteractor> mShadeInteractor;
 
-    private final KeyboardTouchpadEduStatsInteractor mKeyboardTouchpadEduStatsInteractor;
-
     private final Runnable mConnectionRunnable = () ->
             internalConnectToCurrentUser("runnable: startConnectionToCurrentUser");
     private final ComponentName mRecentsComponentName;
@@ -660,8 +657,7 @@
             AssistUtils assistUtils,
             DumpManager dumpManager,
             Optional<UnfoldTransitionProgressForwarder> unfoldTransitionProgressForwarder,
-            BroadcastDispatcher broadcastDispatcher,
-            KeyboardTouchpadEduStatsInteractor keyboardTouchpadEduStatsInteractor
+            BroadcastDispatcher broadcastDispatcher
     ) {
         // b/241601880: This component should only be running for primary users or
         // secondaryUsers when visibleBackgroundUsers are supported.
@@ -699,7 +695,6 @@
         mDisplayTracker = displayTracker;
         mUnfoldTransitionProgressForwarder = unfoldTransitionProgressForwarder;
         mBroadcastDispatcher = broadcastDispatcher;
-        mKeyboardTouchpadEduStatsInteractor = keyboardTouchpadEduStatsInteractor;
 
         if (!KeyguardWmStateRefactor.isEnabled()) {
             mSysuiUnlockAnimationController = sysuiUnlockAnimationController;
@@ -940,19 +935,6 @@
         return isEnabled() && !QuickStepContract.isLegacyMode(mNavBarMode);
     }
 
-    /**
-     * Updates contextual education stats when a gesture is triggered
-     * @param isTrackpadGesture indicates if the gesture is triggered by trackpad
-     * @param gestureType type of gesture triggered
-     */
-    public void updateContextualEduStats(boolean isTrackpadGesture, GestureType gestureType) {
-        if (isTrackpadGesture) {
-            mKeyboardTouchpadEduStatsInteractor.updateShortcutTriggerTime(gestureType);
-        } else {
-            mKeyboardTouchpadEduStatsInteractor.incrementSignalCount(gestureType);
-        }
-    }
-
     public boolean isEnabled() {
         return mIsEnabled;
     }
@@ -978,6 +960,17 @@
         }
     }
 
+    /**
+     * Updates contextual education stats when a gesture is triggered
+     * @param isTrackpadGesture indicates if the gesture is triggered by trackpad
+     * @param gestureType type of gesture triggered
+     */
+    public void updateContextualEduStats(boolean isTrackpadGesture, GestureType gestureType) {
+        for (int i = mConnectionCallbacks.size() - 1; i >= 0; --i) {
+            mConnectionCallbacks.get(i).updateContextualEduStats(isTrackpadGesture, gestureType);
+        }
+    }
+
     private void notifyHomeRotationEnabled(boolean enabled) {
         for (int i = mConnectionCallbacks.size() - 1; i >= 0; --i) {
             mConnectionCallbacks.get(i).onHomeRotationEnabled(enabled);
@@ -1207,6 +1200,9 @@
         /** Set override of home button long press duration, touch slop multiplier, and haptic. */
         default void setOverrideHomeButtonLongPress(
                 long override, float slopMultiplier, boolean haptic) {}
+        /** Updates contextual education stats when target gesture type is triggered. */
+        default void updateContextualEduStats(
+                boolean isTrackpadGesture, GestureType gestureType) {}
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
index 0ef5207..9455201 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
@@ -24,8 +24,8 @@
     data class FullScreen(val displayId: Int) : CaptureType
 
     /** Capture the contents of the task only. */
-    data class IsolatedTask(
-        val taskId: Int,
-        val taskBounds: Rect?,
-    ) : CaptureType
+    data class IsolatedTask(val taskId: Int, val taskBounds: Rect?) : CaptureType
+
+    data class RootTask(val parentTaskId: Int, val taskBounds: Rect?, val childTaskIds: List<Int>) :
+        CaptureType
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
index 039143a..e840668 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
@@ -26,6 +26,7 @@
 import android.util.Log
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
+import com.android.systemui.Flags.screenshotPolicySplitAndDesktopMode
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.screenshot.ImageCapture
 import com.android.systemui.screenshot.ScreenshotData
@@ -47,14 +48,17 @@
     private val capture: ImageCapture,
     /** Provides information about the tasks on a given display */
     private val displayTasks: DisplayContentRepository,
-    /** The list of policies to apply, in order of priority */
+    /** The legacy list of policy implementations to apply, in order of priority */
     private val policies: List<CapturePolicy>,
+    /** Implements the combined policy rules for all profile types. */
+    private val policy: ScreenshotPolicy,
     /** The owner to assign for screenshot when a focused task isn't visible */
     private val defaultOwner: UserHandle = myUserHandle(),
     /** The assigned component when no application has focus, or not visible */
     private val defaultComponent: ComponentName,
 ) : ScreenshotRequestProcessor {
     override suspend fun process(original: ScreenshotData): ScreenshotData {
+
         if (original.type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
             // The request contains an already captured screenshot, accept it as is.
             Log.i(TAG, "Screenshot bitmap provided. No modifications applied.")
@@ -62,6 +66,12 @@
         }
         val displayContent = displayTasks.getDisplayContent(original.displayId)
 
+        if (screenshotPolicySplitAndDesktopMode()) {
+            Log.i(TAG, "Applying screenshot policy....")
+            val type = policy.apply(displayContent, defaultComponent, defaultOwner)
+            return modify(original, type)
+        }
+
         // If policies yield explicit modifications, apply them and return the result
         Log.i(TAG, "Applying policy checks....")
         policies.map { policy ->
@@ -79,10 +89,8 @@
     }
 
     /** Produce a new [ScreenshotData] using [CaptureParameters] */
-    private suspend fun modify(
-        original: ScreenshotData,
-        updates: CaptureParameters,
-    ): ScreenshotData {
+    suspend fun modify(original: ScreenshotData, updates: CaptureParameters): ScreenshotData {
+        Log.d(TAG, "[modify] CaptureParameters = $updates")
         // Update and apply bitmap capture depending on the parameters.
         val updated =
             when (val type = updates.type) {
@@ -94,6 +102,14 @@
                         type.taskId,
                         type.taskBounds,
                     )
+                is CaptureType.RootTask ->
+                    replaceWithTaskSnapshot(
+                        original,
+                        updates.component,
+                        updates.owner,
+                        type.parentTaskId,
+                        type.taskBounds,
+                    )
                 is FullScreen ->
                     replaceWithScreenshot(
                         original,
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
index f768cfb..dd39f92 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
@@ -26,9 +26,11 @@
             childTaskIds[index],
             childTaskNames[index],
             childTaskBounds[index],
-            childTaskUserIds[index]
+            childTaskUserIds[index],
         )
     }
 }
 
 internal fun RootTaskInfo.hasChildTasks() = childTaskUserIds.isNotEmpty()
+
+internal fun RootTaskInfo.childTaskCount() = childTaskIds.size
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt
new file mode 100644
index 0000000..9967aff
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt
@@ -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 com.android.systemui.screenshot.policy
+
+import android.app.ActivityTaskManager.RootTaskInfo
+import android.app.WindowConfiguration
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WINDOWING_MODE_PINNED
+import android.content.ComponentName
+import android.os.UserHandle
+import android.util.Log
+import com.android.systemui.screenshot.data.model.DisplayContentModel
+import com.android.systemui.screenshot.data.model.ProfileType
+import com.android.systemui.screenshot.data.model.ProfileType.PRIVATE
+import com.android.systemui.screenshot.data.model.ProfileType.WORK
+import com.android.systemui.screenshot.data.repository.ProfileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
+import com.android.systemui.screenshot.policy.CaptureType.RootTask
+import javax.inject.Inject
+
+private const val TAG = "ScreenshotPolicy"
+
+/** Determines what to capture and which user owns the output. */
+class ScreenshotPolicy @Inject constructor(private val profileTypes: ProfileTypeRepository) {
+    /**
+     * Apply the policy to the content, resulting in [CaptureParameters].
+     *
+     * @param content the content of the display
+     * @param defaultComponent the component associated with the screenshot by default
+     * @param defaultOwner the user to own the screenshot by default
+     */
+    suspend fun apply(
+        content: DisplayContentModel,
+        defaultComponent: ComponentName,
+        defaultOwner: UserHandle,
+    ): CaptureParameters {
+        val defaultFullScreen by lazy {
+            CaptureParameters(
+                type = FullScreen(displayId = content.displayId),
+                component = defaultComponent,
+                owner = defaultOwner,
+            )
+        }
+
+        // When the systemUI notification shade is open, disregard tasks.
+        if (content.systemUiState.shadeExpanded) {
+            return defaultFullScreen
+        }
+
+        // find the first (top) RootTask which is visible and not Picture-in-Picture
+        val topRootTask =
+            content.rootTasks.firstOrNull {
+                it.isVisible && it.windowingMode != WindowConfiguration.WINDOWING_MODE_PINNED
+            } ?: return defaultFullScreen
+
+        Log.d(TAG, "topRootTask: $topRootTask")
+        val rootTaskOwners = topRootTask.childTaskUserIds.distinct()
+
+        // Special case: Only WORK in top root task which is full-screen or maximized freeform
+        if (
+            rootTaskOwners.size == 1 &&
+                profileTypes.getProfileType(rootTaskOwners.single()) == WORK &&
+                (topRootTask.isFullScreen() || topRootTask.isMaximizedFreeform())
+        ) {
+            val type =
+                if (topRootTask.childTaskCount() > 1) {
+                    RootTask(
+                        parentTaskId = topRootTask.taskId,
+                        taskBounds = topRootTask.bounds,
+                        childTaskIds = topRootTask.childTasksTopDown().map { it.id }.toList(),
+                    )
+                } else {
+                    IsolatedTask(
+                        taskId = topRootTask.childTasksTopDown().first().id,
+                        taskBounds = topRootTask.bounds,
+                    )
+                }
+            // Capture the RootTask (and all children)
+            return CaptureParameters(
+                type = type,
+                component = topRootTask.topActivity,
+                owner = UserHandle.of(rootTaskOwners.single()),
+            )
+        }
+
+        // In every other case the output will be a full screen capture regardless of content.
+        // For this reason, consider all owners of all visible content on the display (in all
+        // root tasks). This includes all root tasks in free-form mode.
+        val visibleChildTasks =
+            content.rootTasks.filter { it.isVisible }.flatMap { it.childTasksTopDown() }
+
+        val allVisibleProfileTypes =
+            visibleChildTasks
+                .map { it.userId }
+                .distinct()
+                .associate { profileTypes.getProfileType(it) to UserHandle.of(it) }
+
+        // If any visible content belongs to the private profile user -> private profile
+        // otherwise the personal user (including partial screen work content).
+        val ownerHandle =
+            allVisibleProfileTypes[PRIVATE]
+                ?: allVisibleProfileTypes[ProfileType.NONE]
+                ?: defaultOwner
+
+        // Attribute to the component of top-most task owned by this user (or fallback to default)
+        val topComponent =
+            visibleChildTasks.firstOrNull { it.userId == ownerHandle.identifier }?.componentName
+
+        return CaptureParameters(
+            type = FullScreen(content.displayId),
+            component = topComponent ?: topRootTask.topActivity ?: defaultComponent,
+            owner = ownerHandle,
+        )
+    }
+
+    private fun RootTaskInfo.isFullScreen(): Boolean =
+        configuration.windowConfiguration.windowingMode == WINDOWING_MODE_FULLSCREEN
+
+    private fun RootTaskInfo.isMaximizedFreeform(): Boolean {
+        val bounds = configuration.windowConfiguration.bounds
+        val maxBounds = configuration.windowConfiguration.maxBounds
+
+        if (
+            windowingMode != WINDOWING_MODE_FREEFORM ||
+                childTaskCount() != 1 ||
+                childTaskBounds[0] != bounds
+        ) {
+            return false
+        }
+
+        // Maximized floating windows fill maxBounds width
+        if (bounds.width() != maxBounds.width()) {
+            return false
+        }
+
+        // Maximized floating windows fill nearly all the height
+        return (bounds.height().toFloat() / maxBounds.height()) >= 0.89f
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
index 2cb9fe7..a9c6370 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
@@ -37,7 +37,6 @@
 
 @Module
 interface ScreenshotPolicyModule {
-
     @Binds
     @SysUISingleton
     fun bindProfileTypeRepository(impl: ProfileTypeRepositoryImpl): ProfileTypeRepository
@@ -67,6 +66,7 @@
             imageCapture: ImageCapture,
             displayContentRepo: DisplayContentRepository,
             policyListProvider: Provider<List<CapturePolicy>>,
+            standardPolicy: ScreenshotPolicy,
         ): ScreenshotRequestProcessor {
             return PolicyRequestProcessor(
                 background = background,
@@ -75,7 +75,8 @@
                 policies = policyListProvider.get(),
                 defaultOwner = Process.myUserHandle(),
                 defaultComponent =
-                    ComponentName(context.packageName, SystemUIService::class.java.toString())
+                    ComponentName(context.packageName, SystemUIService::class.java.toString()),
+                policy = standardPolicy,
             )
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
index 29450a2..cf90c0a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
@@ -28,7 +28,6 @@
 import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import javax.inject.Inject
-import kotlinx.coroutines.flow.first
 
 /**
  * Condition: When the top visible task (excluding PIP mode) belongs to a work user.
@@ -37,10 +36,8 @@
  */
 class WorkProfilePolicy
 @Inject
-constructor(
-    private val profileTypes: ProfileTypeRepository,
-    private val context: Context,
-) : CapturePolicy {
+constructor(private val profileTypes: ProfileTypeRepository, private val context: Context) :
+    CapturePolicy {
 
     override suspend fun check(content: DisplayContentModel): PolicyResult {
         // The systemUI notification shade isn't a work app, skip.
@@ -65,11 +62,7 @@
                 .map { it to it.childTasksTopDown().first() }
                 .firstOrNull { (_, child) ->
                     profileTypes.getProfileType(child.userId) == ProfileType.WORK
-                }
-                ?: return NotMatched(
-                    policy = NAME,
-                    reason = WORK_TASK_NOT_TOP,
-                )
+                } ?: return NotMatched(policy = NAME, reason = WORK_TASK_NOT_TOP)
 
         // If matched, return parameters needed to modify the request.
         return PolicyResult.Matched(
@@ -79,7 +72,7 @@
                 type = IsolatedTask(taskId = childTask.id, taskBounds = childTask.bounds),
                 component = childTask.componentName ?: rootTask.topActivity,
                 owner = UserHandle.of(childTask.userId),
-            )
+            ),
         )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
index 8d3f728..30f564f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationGroupingUtil.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar;
 
+import android.app.Flags;
 import android.app.Notification;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.Icon;
@@ -60,20 +61,6 @@
             return row.getEntry().getSbn().getNotification();
         }
     };
-    private static final IconComparator ICON_VISIBILITY_COMPARATOR = new IconComparator() {
-        public boolean compare(View parent, View child, Object parentData,
-                Object childData) {
-            return hasSameIcon(parentData, childData)
-                    && hasSameColor(parentData, childData);
-        }
-    };
-    private static final IconComparator GREY_COMPARATOR = new IconComparator() {
-        public boolean compare(View parent, View child, Object parentData,
-                Object childData) {
-            return !hasSameIcon(parentData, childData)
-                    || hasSameColor(parentData, childData);
-        }
-    };
     private static final ResultApplicator GREY_APPLICATOR = new ResultApplicator() {
         @Override
         public void apply(View parent, View view, boolean apply, boolean reset) {
@@ -90,34 +77,58 @@
 
     public NotificationGroupingUtil(ExpandableNotificationRow row) {
         mRow = row;
+
+        final IconComparator iconVisibilityComparator = new IconComparator(mRow) {
+            public boolean compare(View parent, View child, Object parentData,
+                    Object childData) {
+                return hasSameIcon(parentData, childData)
+                        && hasSameColor(parentData, childData);
+            }
+        };
+        final IconComparator greyComparator = new IconComparator(mRow) {
+            public boolean compare(View parent, View child, Object parentData,
+                    Object childData) {
+                if (Flags.notificationsRedesignAppIcons() && mRow.isShowingAppIcon()) {
+                    return false;
+                }
+                return !hasSameIcon(parentData, childData)
+                        || hasSameColor(parentData, childData);
+            }
+        };
+
         // To hide the icons if they are the same and the color is the same
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.icon,
                 ICON_EXTRACTOR,
-                ICON_VISIBILITY_COMPARATOR,
+                iconVisibilityComparator,
                 VISIBILITY_APPLICATOR));
-        // To grey them out the icons and expand button when the icons are not the same
+        // To grey out the icons when they are not the same, or they have the same color
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.status_bar_latest_event_content,
                 ICON_EXTRACTOR,
-                GREY_COMPARATOR,
+                greyComparator,
                 GREY_APPLICATOR));
+        // To show the large icon on the left side instead if all the small icons are the same
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.status_bar_latest_event_content,
                 ICON_EXTRACTOR,
-                ICON_VISIBILITY_COMPARATOR,
+                iconVisibilityComparator,
                 LEFT_ICON_APPLICATOR));
+        // To only show the work profile icon in the group header
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.profile_badge,
                 null /* Extractor */,
                 BADGE_COMPARATOR,
                 VISIBILITY_APPLICATOR));
+        // To hide the app name in group children
         mProcessors.add(new Processor(mRow,
                 com.android.internal.R.id.app_name_text,
                 null,
                 APP_NAME_COMPARATOR,
                 APP_NAME_APPLICATOR));
+        // To hide the header text if it's the same
         mProcessors.add(Processor.forTextView(mRow, com.android.internal.R.id.header_text));
+
         mDividers.add(com.android.internal.R.id.header_text_divider);
         mDividers.add(com.android.internal.R.id.header_text_secondary_divider);
         mDividers.add(com.android.internal.R.id.time_divider);
@@ -261,6 +272,7 @@
             mParentData = mExtractor == null ? null : mExtractor.extractData(mParentRow);
             mApply = !mComparator.isEmpty(mParentView);
         }
+
         public void compareToGroupParent(ExpandableNotificationRow row) {
             if (!mApply) {
                 return;
@@ -356,12 +368,21 @@
     }
 
     private abstract static class IconComparator implements ViewComparator {
+        private final ExpandableNotificationRow mRow;
+
+        IconComparator(ExpandableNotificationRow row) {
+            mRow = row;
+        }
+
         @Override
         public boolean compare(View parent, View child, Object parentData, Object childData) {
             return false;
         }
 
         protected boolean hasSameIcon(Object parentData, Object childData) {
+            if (Flags.notificationsRedesignAppIcons() && mRow.isShowingAppIcon()) {
+                return true;
+            }
             Icon parentIcon = getIcon((Notification) parentData);
             Icon childIcon = getIcon((Notification) childData);
             return parentIcon.sameAs(childIcon);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
index f6f4503..f65ae67 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
@@ -17,8 +17,10 @@
 package com.android.systemui.statusbar.dagger
 
 import android.content.Context
+import com.android.systemui.CameraProtectionLoader
 import com.android.systemui.CoreStartable
 import com.android.systemui.SysUICutoutProvider
+import com.android.systemui.SysUICutoutProviderImpl
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.log.LogBuffer
 import com.android.systemui.log.LogBufferFactory
@@ -114,6 +116,16 @@
 
         @Provides
         @SysUISingleton
+        fun sysUiCutoutProvider(
+            factory: SysUICutoutProviderImpl.Factory,
+            context: Context,
+            cameraProtectionLoader: CameraProtectionLoader,
+        ): SysUICutoutProvider {
+            return factory.create(context, cameraProtectionLoader)
+        }
+
+        @Provides
+        @SysUISingleton
         fun contentInsetsProvider(
             factory: StatusBarContentInsetsProviderImpl.Factory,
             context: Context,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/events/StatusBarEventsModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/events/StatusBarEventsModule.kt
index ed96482..415d990 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/events/StatusBarEventsModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/events/StatusBarEventsModule.kt
@@ -23,7 +23,7 @@
 import dagger.Module
 import dagger.Provides
 
-@Module
+@Module(includes = [SystemEventChipAnimationControllerModule::class])
 interface StatusBarEventsModule {
 
     companion object {
@@ -41,4 +41,4 @@
     fun bindSystemStatusAnimationScheduler(
         systemStatusAnimationSchedulerImpl: SystemStatusAnimationSchedulerImpl
     ): SystemStatusAnimationScheduler
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemEventChipAnimationController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemEventChipAnimationController.kt
index bf7e879..35816c2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemEventChipAnimationController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemEventChipAnimationController.kt
@@ -31,22 +31,46 @@
 import androidx.core.animation.AnimatorSet
 import androidx.core.animation.ValueAnimator
 import com.android.internal.annotations.VisibleForTesting
+import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.phone.StatusBarContentInsetsChangedListener
 import com.android.systemui.statusbar.phone.StatusBarContentInsetsProvider
+import com.android.systemui.statusbar.window.StatusBarWindowController
 import com.android.systemui.statusbar.window.StatusBarWindowControllerStore
 import com.android.systemui.util.animation.AnimationUtil.Companion.frames
-import javax.inject.Inject
+import dagger.Module
+import dagger.Provides
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import kotlin.math.roundToInt
 
-/**
- * Controls the view for system event animations.
- */
-class SystemEventChipAnimationController @Inject constructor(
-    private val context: Context,
-    private val statusBarWindowControllerStore: StatusBarWindowControllerStore,
-    private val contentInsetsProvider: StatusBarContentInsetsProvider,
-) : SystemStatusAnimationCallback {
+/** Controls the view for system event animations. */
+interface SystemEventChipAnimationController : SystemStatusAnimationCallback {
+
+    /**
+     * Give the chip controller a chance to inflate and configure the chip view before we start
+     * animating
+     */
+    fun prepareChipAnimation(viewCreator: ViewCreator)
+
+    fun init()
+
+    /** Announces [contentDescriptions] for accessibility. */
+    fun announceForAccessibility(contentDescriptions: String)
+
+    override fun onSystemEventAnimationBegin(): Animator
+
+    override fun onSystemEventAnimationFinish(hasPersistentDot: Boolean): Animator
+}
+
+class SystemEventChipAnimationControllerImpl
+@AssistedInject
+constructor(
+    @Assisted private val context: Context,
+    @Assisted private val statusBarWindowController: StatusBarWindowController,
+    @Assisted private val contentInsetsProvider: StatusBarContentInsetsProvider,
+) : SystemEventChipAnimationController {
 
     private lateinit var animationWindowView: FrameLayout
     private lateinit var themedContext: ContextThemeWrapper
@@ -57,25 +81,27 @@
     private var animationDirection = LEFT
 
     @VisibleForTesting var chipBounds = Rect()
-    private val chipWidth get() = chipBounds.width()
-    private val chipRight get() = chipBounds.right
-    private val chipLeft get() = chipBounds.left
-    private var chipMinWidth = context.resources.getDimensionPixelSize(
-            R.dimen.ongoing_appops_chip_min_animation_width)
+    private val chipWidth
+        get() = chipBounds.width()
 
-    private val dotSize = context.resources.getDimensionPixelSize(
-            R.dimen.ongoing_appops_dot_diameter)
+    private val chipRight
+        get() = chipBounds.right
+
+    private val chipLeft
+        get() = chipBounds.left
+
+    private var chipMinWidth =
+        context.resources.getDimensionPixelSize(R.dimen.ongoing_appops_chip_min_animation_width)
+
+    private val dotSize =
+        context.resources.getDimensionPixelSize(R.dimen.ongoing_appops_dot_diameter)
     // Use during animation so that multiple animators can update the drawing rect
     private var animRect = Rect()
 
     // TODO: move to dagger
     @VisibleForTesting var initialized = false
 
-    /**
-     * Give the chip controller a chance to inflate and configure the chip view before we start
-     * animating
-     */
-    fun prepareChipAnimation(viewCreator: ViewCreator) {
+    override fun prepareChipAnimation(viewCreator: ViewCreator) {
         if (!initialized) {
             init()
         }
@@ -83,47 +109,62 @@
 
         // Initialize the animated view
         val insets = contentInsetsProvider.getStatusBarContentInsetsForCurrentRotation()
-        currentAnimatedView = viewCreator(themedContext).also {
-            animationWindowView.addView(
+        currentAnimatedView =
+            viewCreator(themedContext).also {
+                animationWindowView.addView(
                     it.view,
                     layoutParamsDefault(
-                            if (animationWindowView.isLayoutRtl) insets.left
-                            else insets.right))
-            it.view.alpha = 0f
-            // For some reason, the window view's measured width is always 0 here, so use the
-            // parent (status bar)
-            it.view.measure(
+                        if (animationWindowView.isLayoutRtl) insets.left else insets.right
+                    ),
+                )
+                it.view.alpha = 0f
+                // For some reason, the window view's measured width is always 0 here, so use the
+                // parent (status bar)
+                it.view.measure(
                     View.MeasureSpec.makeMeasureSpec(
-                            (animationWindowView.parent as View).width, AT_MOST),
+                        (animationWindowView.parent as View).width,
+                        AT_MOST,
+                    ),
                     View.MeasureSpec.makeMeasureSpec(
-                            (animationWindowView.parent as View).height, AT_MOST))
+                        (animationWindowView.parent as View).height,
+                        AT_MOST,
+                    ),
+                )
 
-            updateChipBounds(it, contentInsetsProvider.getStatusBarContentAreaForCurrentRotation())
-        }
+                updateChipBounds(
+                    it,
+                    contentInsetsProvider.getStatusBarContentAreaForCurrentRotation(),
+                )
+            }
     }
 
     override fun onSystemEventAnimationBegin(): Animator {
         initializeAnimRect()
 
-        val alphaIn = ValueAnimator.ofFloat(0f, 1f).apply {
-            startDelay = 7.frames
-            duration = 5.frames
-            interpolator = null
-            addUpdateListener { currentAnimatedView?.view?.alpha = animatedValue as Float }
-        }
+        val alphaIn =
+            ValueAnimator.ofFloat(0f, 1f).apply {
+                startDelay = 7.frames
+                duration = 5.frames
+                interpolator = null
+                addUpdateListener { currentAnimatedView?.view?.alpha = animatedValue as Float }
+            }
         currentAnimatedView?.contentView?.alpha = 0f
-        val contentAlphaIn = ValueAnimator.ofFloat(0f, 1f).apply {
-            startDelay = 10.frames
-            duration = 10.frames
-            interpolator = null
-            addUpdateListener { currentAnimatedView?.contentView?.alpha = animatedValue as Float }
-        }
-        val moveIn = ValueAnimator.ofInt(chipMinWidth, chipWidth).apply {
-            startDelay = 7.frames
-            duration = 23.frames
-            interpolator = STATUS_BAR_X_MOVE_IN
-            addUpdateListener { updateAnimatedViewBoundsWidth(animatedValue as Int) }
-        }
+        val contentAlphaIn =
+            ValueAnimator.ofFloat(0f, 1f).apply {
+                startDelay = 10.frames
+                duration = 10.frames
+                interpolator = null
+                addUpdateListener {
+                    currentAnimatedView?.contentView?.alpha = animatedValue as Float
+                }
+            }
+        val moveIn =
+            ValueAnimator.ofInt(chipMinWidth, chipWidth).apply {
+                startDelay = 7.frames
+                duration = 23.frames
+                interpolator = STATUS_BAR_X_MOVE_IN
+                addUpdateListener { updateAnimatedViewBoundsWidth(animatedValue as Int) }
+            }
         val animSet = AnimatorSet()
         animSet.playTogether(alphaIn, contentAlphaIn, moveIn)
         return animSet
@@ -131,75 +172,80 @@
 
     override fun onSystemEventAnimationFinish(hasPersistentDot: Boolean): Animator {
         initializeAnimRect()
-        val finish = if (hasPersistentDot) {
-            createMoveOutAnimationForDot()
-        } else {
-            createMoveOutAnimationDefault()
-        }
-
-        finish.addListener(object : AnimatorListenerAdapter() {
-            override fun onAnimationEnd(animation: Animator) {
-                animationWindowView.removeView(currentAnimatedView!!.view)
+        val finish =
+            if (hasPersistentDot) {
+                createMoveOutAnimationForDot()
+            } else {
+                createMoveOutAnimationDefault()
             }
-        })
+
+        finish.addListener(
+            object : AnimatorListenerAdapter() {
+                override fun onAnimationEnd(animation: Animator) {
+                    animationWindowView.removeView(currentAnimatedView!!.view)
+                }
+            }
+        )
 
         return finish
     }
 
     private fun createMoveOutAnimationForDot(): Animator {
-        val width1 = ValueAnimator.ofInt(chipWidth, chipMinWidth).apply {
-            duration = 9.frames
-            interpolator = STATUS_CHIP_WIDTH_TO_DOT_KEYFRAME_1
-            addUpdateListener {
-                updateAnimatedViewBoundsWidth(animatedValue as Int)
+        val width1 =
+            ValueAnimator.ofInt(chipWidth, chipMinWidth).apply {
+                duration = 9.frames
+                interpolator = STATUS_CHIP_WIDTH_TO_DOT_KEYFRAME_1
+                addUpdateListener { updateAnimatedViewBoundsWidth(animatedValue as Int) }
             }
-        }
 
-        val width2 = ValueAnimator.ofInt(chipMinWidth, dotSize).apply {
-            startDelay = 9.frames
-            duration = 20.frames
-            interpolator = STATUS_CHIP_WIDTH_TO_DOT_KEYFRAME_2
-            addUpdateListener {
-                updateAnimatedViewBoundsWidth(animatedValue as Int)
+        val width2 =
+            ValueAnimator.ofInt(chipMinWidth, dotSize).apply {
+                startDelay = 9.frames
+                duration = 20.frames
+                interpolator = STATUS_CHIP_WIDTH_TO_DOT_KEYFRAME_2
+                addUpdateListener { updateAnimatedViewBoundsWidth(animatedValue as Int) }
             }
-        }
 
         val keyFrame1Height = dotSize * 2
         val chipVerticalCenter = chipBounds.top + chipBounds.height() / 2
-        val height1 = ValueAnimator.ofInt(chipBounds.height(), keyFrame1Height).apply {
-            startDelay = 8.frames
-            duration = 6.frames
-            interpolator = STATUS_CHIP_HEIGHT_TO_DOT_KEYFRAME_1
-            addUpdateListener {
-                updateAnimatedViewBoundsHeight(animatedValue as Int, chipVerticalCenter)
+        val height1 =
+            ValueAnimator.ofInt(chipBounds.height(), keyFrame1Height).apply {
+                startDelay = 8.frames
+                duration = 6.frames
+                interpolator = STATUS_CHIP_HEIGHT_TO_DOT_KEYFRAME_1
+                addUpdateListener {
+                    updateAnimatedViewBoundsHeight(animatedValue as Int, chipVerticalCenter)
+                }
             }
-        }
 
-        val height2 = ValueAnimator.ofInt(keyFrame1Height, dotSize).apply {
-            startDelay = 14.frames
-            duration = 15.frames
-            interpolator = STATUS_CHIP_HEIGHT_TO_DOT_KEYFRAME_2
-            addUpdateListener {
-                updateAnimatedViewBoundsHeight(animatedValue as Int, chipVerticalCenter)
+        val height2 =
+            ValueAnimator.ofInt(keyFrame1Height, dotSize).apply {
+                startDelay = 14.frames
+                duration = 15.frames
+                interpolator = STATUS_CHIP_HEIGHT_TO_DOT_KEYFRAME_2
+                addUpdateListener {
+                    updateAnimatedViewBoundsHeight(animatedValue as Int, chipVerticalCenter)
+                }
             }
-        }
 
         // Move the chip view to overlap exactly with the privacy dot. The chip displays by default
         // exactly adjacent to the dot, so we can just move over by the diameter of the dot itself
-        val moveOut = ValueAnimator.ofInt(0, dotSize).apply {
-            startDelay = 3.frames
-            duration = 11.frames
-            interpolator = STATUS_CHIP_MOVE_TO_DOT
-            addUpdateListener {
-                // If RTL, we can just invert the move
-                val amt = if (animationDirection == LEFT) {
-                        animatedValue as Int
-                } else {
-                    -(animatedValue as Int)
+        val moveOut =
+            ValueAnimator.ofInt(0, dotSize).apply {
+                startDelay = 3.frames
+                duration = 11.frames
+                interpolator = STATUS_CHIP_MOVE_TO_DOT
+                addUpdateListener {
+                    // If RTL, we can just invert the move
+                    val amt =
+                        if (animationDirection == LEFT) {
+                            animatedValue as Int
+                        } else {
+                            -(animatedValue as Int)
+                        }
+                    updateAnimatedBoundsX(amt)
                 }
-                updateAnimatedBoundsX(amt)
             }
-        }
 
         val animSet = AnimatorSet()
         animSet.playTogether(width1, width2, height1, height2, moveOut)
@@ -207,71 +253,80 @@
     }
 
     private fun createMoveOutAnimationDefault(): Animator {
-        val alphaOut = ValueAnimator.ofFloat(1f, 0f).apply {
-            startDelay = 6.frames
-            duration = 6.frames
-            interpolator = null
-            addUpdateListener { currentAnimatedView?.view?.alpha = animatedValue as Float }
-        }
+        val alphaOut =
+            ValueAnimator.ofFloat(1f, 0f).apply {
+                startDelay = 6.frames
+                duration = 6.frames
+                interpolator = null
+                addUpdateListener { currentAnimatedView?.view?.alpha = animatedValue as Float }
+            }
 
-        val contentAlphaOut = ValueAnimator.ofFloat(1f, 0f).apply {
-            duration = 5.frames
-            interpolator = null
-            addUpdateListener { currentAnimatedView?.contentView?.alpha = animatedValue as Float }
-        }
-
-        val moveOut = ValueAnimator.ofInt(chipWidth, chipMinWidth).apply {
-            duration = 23.frames
-            interpolator = STATUS_BAR_X_MOVE_OUT
-            addUpdateListener {
-                currentAnimatedView?.apply {
-                    updateAnimatedViewBoundsWidth(animatedValue as Int)
+        val contentAlphaOut =
+            ValueAnimator.ofFloat(1f, 0f).apply {
+                duration = 5.frames
+                interpolator = null
+                addUpdateListener {
+                    currentAnimatedView?.contentView?.alpha = animatedValue as Float
                 }
             }
-        }
+
+        val moveOut =
+            ValueAnimator.ofInt(chipWidth, chipMinWidth).apply {
+                duration = 23.frames
+                interpolator = STATUS_BAR_X_MOVE_OUT
+                addUpdateListener {
+                    currentAnimatedView?.apply {
+                        updateAnimatedViewBoundsWidth(animatedValue as Int)
+                    }
+                }
+            }
 
         val animSet = AnimatorSet()
         animSet.playTogether(alphaOut, contentAlphaOut, moveOut)
         return animSet
     }
 
-    fun init() {
+    override fun init() {
         initialized = true
         themedContext = ContextThemeWrapper(context, R.style.Theme_SystemUI_QuickSettings)
-        animationWindowView = LayoutInflater.from(themedContext)
-                .inflate(R.layout.system_event_animation_window, null) as FrameLayout
+        animationWindowView =
+            LayoutInflater.from(themedContext).inflate(R.layout.system_event_animation_window, null)
+                as FrameLayout
         // Matches status_bar.xml
         val height = themedContext.resources.getDimensionPixelSize(R.dimen.status_bar_height)
         val lp = FrameLayout.LayoutParams(MATCH_PARENT, height)
         lp.gravity = Gravity.END or Gravity.TOP
-        statusBarWindowControllerStore.defaultDisplay.addViewToWindow(animationWindowView, lp)
+        statusBarWindowController.addViewToWindow(animationWindowView, lp)
         animationWindowView.clipToPadding = false
         animationWindowView.clipChildren = false
 
         // Use contentInsetsProvider rather than configuration controller, since we only care
         // about status bar dimens
-        contentInsetsProvider.addCallback(object : StatusBarContentInsetsChangedListener {
-            override fun onStatusBarContentInsetsChanged() {
-                val newContentArea = contentInsetsProvider
-                    .getStatusBarContentAreaForCurrentRotation()
-                updateDimens(newContentArea)
+        contentInsetsProvider.addCallback(
+            object : StatusBarContentInsetsChangedListener {
+                override fun onStatusBarContentInsetsChanged() {
+                    val newContentArea =
+                        contentInsetsProvider.getStatusBarContentAreaForCurrentRotation()
+                    updateDimens(newContentArea)
 
-                // If we are currently animating, we have to re-solve for the chip bounds. If we're
-                // not animating then [prepareChipAnimation] will take care of it for us
-                currentAnimatedView?.let {
-                    updateChipBounds(it, newContentArea)
-                    // Since updateCurrentAnimatedView can only be called during an animation, we
-                    // have to create a dummy animator here to apply the new chip bounds
-                    val animator = ValueAnimator.ofInt(0, 1).setDuration(0)
-                    animator.addUpdateListener { updateCurrentAnimatedView() }
-                    animator.start()
+                    // If we are currently animating, we have to re-solve for the chip bounds. If
+                    // we're
+                    // not animating then [prepareChipAnimation] will take care of it for us
+                    currentAnimatedView?.let {
+                        updateChipBounds(it, newContentArea)
+                        // Since updateCurrentAnimatedView can only be called during an animation,
+                        // we
+                        // have to create a dummy animator here to apply the new chip bounds
+                        val animator = ValueAnimator.ofInt(0, 1).setDuration(0)
+                        animator.addUpdateListener { updateCurrentAnimatedView() }
+                        animator.start()
+                    }
                 }
             }
-        })
+        )
     }
 
-    /** Announces [contentDescriptions] for accessibility. */
-    fun announceForAccessibility(contentDescriptions: String) {
+    override fun announceForAccessibility(contentDescriptions: String) {
         currentAnimatedView?.view?.announceForAccessibility(contentDescriptions)
     }
 
@@ -283,9 +338,9 @@
     }
 
     /**
-     * Use the current status bar content area and the current chip's measured size to update
-     * the animation rect and chipBounds. This method can be called at any time and will update
-     * the current animation values properly during e.g. a rotation.
+     * Use the current status bar content area and the current chip's measured size to update the
+     * animation rect and chipBounds. This method can be called at any time and will update the
+     * current animation values properly during e.g. a rotation.
      */
     private fun updateChipBounds(chip: BackgroundAnimatableView, contentArea: Rect) {
         // decide which direction we're animating from, and then set some screen coordinates
@@ -309,14 +364,13 @@
     }
 
     private fun layoutParamsDefault(marginEnd: Int): FrameLayout.LayoutParams =
-            FrameLayout.LayoutParams(WRAP_CONTENT, WRAP_CONTENT).also {
-                it.gravity = Gravity.END or Gravity.CENTER_VERTICAL
-                it.marginEnd = marginEnd
-            }
+        FrameLayout.LayoutParams(WRAP_CONTENT, WRAP_CONTENT).also {
+            it.gravity = Gravity.END or Gravity.CENTER_VERTICAL
+            it.marginEnd = marginEnd
+        }
 
     private fun initializeAnimRect() = animRect.set(chipBounds)
 
-
     /**
      * To be called during an animation, sets the width and updates the current animated chip view
      */
@@ -324,7 +378,8 @@
         when (animationDirection) {
             LEFT -> {
                 animRect.set((chipRight - width), animRect.top, chipRight, animRect.bottom)
-            } else /* RIGHT */ -> {
+            }
+            else /* RIGHT */ -> {
                 animRect.set(chipLeft, animRect.top, (chipLeft + width), animRect.bottom)
             }
         }
@@ -337,44 +392,73 @@
      */
     private fun updateAnimatedViewBoundsHeight(height: Int, verticalCenter: Int) {
         animRect.set(
-                animRect.left,
-                verticalCenter - (height.toFloat() / 2).roundToInt(),
-                animRect.right,
-                verticalCenter + (height.toFloat() / 2).roundToInt())
+            animRect.left,
+            verticalCenter - (height.toFloat() / 2).roundToInt(),
+            animRect.right,
+            verticalCenter + (height.toFloat() / 2).roundToInt(),
+        )
 
         updateCurrentAnimatedView()
     }
 
-    /**
-     * To be called during an animation, updates the animation rect offset and updates the chip
-     */
+    /** To be called during an animation, updates the animation rect offset and updates the chip */
     private fun updateAnimatedBoundsX(translation: Int) {
         currentAnimatedView?.view?.translationX = translation.toFloat()
     }
 
-    /**
-     * To be called during an animation. Sets the chip rect to animRect
-     */
+    /** To be called during an animation. Sets the chip rect to animRect */
     private fun updateCurrentAnimatedView() {
         currentAnimatedView?.setBoundsForAnimation(
-                animRect.left, animRect.top, animRect.right, animRect.bottom
+            animRect.left,
+            animRect.top,
+            animRect.right,
+            animRect.bottom,
         )
     }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(
+            context: Context,
+            statusBarWindowController: StatusBarWindowController,
+            contentInsetsProvider: StatusBarContentInsetsProvider,
+        ): SystemEventChipAnimationControllerImpl
+    }
 }
 
-/**
- * Chips should provide a view that can be animated with something better than a fade-in
- */
+/** Chips should provide a view that can be animated with something better than a fade-in */
 interface BackgroundAnimatableView {
     val view: View // Since this can't extend View, add a view prop
         get() = this as View
+
     val contentView: View? // This will be alpha faded during appear and disappear animation
         get() = null
+
     val chipWidth: Int
         get() = view.measuredWidth
+
     fun setBoundsForAnimation(l: Int, t: Int, r: Int, b: Int)
 }
 
 // Animation directions
 private const val LEFT = 1
 private const val RIGHT = 2
+
+@Module
+object SystemEventChipAnimationControllerModule {
+
+    @Provides
+    @SysUISingleton
+    fun controller(
+        factory: SystemEventChipAnimationControllerImpl.Factory,
+        context: Context,
+        statusBarWindowControllerStore: StatusBarWindowControllerStore,
+        contentInsetsProvider: StatusBarContentInsetsProvider,
+    ): SystemEventChipAnimationController {
+        return factory.create(
+            context,
+            statusBarWindowControllerStore.defaultDisplay,
+            contentInsetsProvider,
+        )
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationScheduler.kt b/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationScheduler.kt
index ef90890..c6f3d7d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationScheduler.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationScheduler.kt
@@ -24,7 +24,7 @@
 import com.android.systemui.statusbar.policy.CallbackController
 
 interface SystemStatusAnimationScheduler :
-        CallbackController<SystemStatusAnimationCallback>, Dumpable {
+    CallbackController<SystemStatusAnimationCallback>, Dumpable {
 
     @SystemAnimationState fun getAnimationState(): Int
 
@@ -44,30 +44,36 @@
  */
 interface SystemStatusAnimationCallback {
     /** Implement this method to return an [Animator] or [AnimatorSet] that presents the chip */
-    fun onSystemEventAnimationBegin(): Animator? { return null }
+    fun onSystemEventAnimationBegin(): Animator? {
+        return null
+    }
+
     /** Implement this method to return an [Animator] or [AnimatorSet] that hides the chip */
-    fun onSystemEventAnimationFinish(hasPersistentDot: Boolean): Animator? { return null }
+    fun onSystemEventAnimationFinish(hasPersistentDot: Boolean): Animator? {
+        return null
+    }
 
     // Best method name, change my mind
     fun onSystemStatusAnimationTransitionToPersistentDot(contentDescription: String?): Animator? {
         return null
     }
-    fun onHidePersistentDot(): Animator? { return null }
+
+    fun onHidePersistentDot(): Animator? {
+        return null
+    }
 }
 
-
-/**
- * Animation state IntDef
- */
+/** Animation state IntDef */
 @Retention(AnnotationRetention.SOURCE)
 @IntDef(
-        value = [
+    value =
+        [
             IDLE,
             ANIMATION_QUEUED,
             ANIMATING_IN,
             RUNNING_CHIP_ANIM,
             ANIMATING_OUT,
-            SHOWING_PERSISTENT_DOT
+            SHOWING_PERSISTENT_DOT,
         ]
 )
 annotation class SystemAnimationState
@@ -110,4 +116,4 @@
 internal const val DISPLAY_LENGTH = 3000L
 internal const val DISAPPEAR_ANIMATION_DURATION = 500L
 
-internal const val MIN_UPTIME: Long = 5 * 1000
\ No newline at end of file
+internal const val MIN_UPTIME: Long = 5 * 1000
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImpl.kt
index e34f61d..5f9e426 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImpl.kt
@@ -69,7 +69,7 @@
     dumpManager: DumpManager,
     private val systemClock: SystemClock,
     @Application private val coroutineScope: CoroutineScope,
-    private val logger: SystemStatusAnimationSchedulerLogger?
+    private val logger: SystemStatusAnimationSchedulerLogger?,
 ) : SystemStatusAnimationScheduler {
 
     companion object {
@@ -122,9 +122,7 @@
                 }
         }
 
-        coroutineScope.launch {
-            animationState.collect { logger?.logAnimationStateUpdate(it) }
-        }
+        coroutineScope.launch { animationState.collect { logger?.logAnimationStateUpdate(it) } }
     }
 
     @SystemAnimationState override fun getAnimationState(): Int = animationState.value
@@ -195,7 +193,7 @@
         return DeviceConfig.getBoolean(
             DeviceConfig.NAMESPACE_PRIVACY,
             PROPERTY_ENABLE_IMMERSIVE_INDICATOR,
-            true
+            true,
         )
     }
 
@@ -262,7 +260,7 @@
 
     private fun announceForAccessibilityIfNeeded(event: StatusEvent) {
         val description = event.contentDescription ?: return
-        if (!event.shouldAnnounceAccessibilityEvent)  return
+        if (!event.shouldAnnounceAccessibilityEvent) return
         chipAnimationController.announceForAccessibility(description)
     }
 
@@ -356,9 +354,7 @@
         logger?.logTransitionToPersistentDotCallbackInvoked()
         val anims: List<Animator> =
             listeners.mapNotNull {
-                it.onSystemStatusAnimationTransitionToPersistentDot(
-                    event?.contentDescription
-                )
+                it.onSystemStatusAnimationTransitionToPersistentDot(event?.contentDescription)
             }
         if (anims.isNotEmpty()) {
             val aSet = AnimatorSet()
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCache.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCache.kt
new file mode 100644
index 0000000..2ee1dffd
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotifCollectionCache.kt
@@ -0,0 +1,188 @@
+/*
+ * 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.notification.collection
+
+import android.annotation.SuppressLint
+import com.android.internal.annotations.VisibleForTesting
+import com.android.systemui.Dumpable
+import com.android.systemui.util.asIndenting
+import com.android.systemui.util.printCollection
+import com.android.systemui.util.time.SystemClock
+import com.android.systemui.util.time.SystemClockImpl
+import com.android.systemui.util.withIncreasedIndent
+import java.io.PrintWriter
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicInteger
+
+/**
+ * A cache in which entries can "survive" getting purged [retainCount] times, given consecutive
+ * [purge] calls made at least [purgeTimeoutMillis] apart. See also [purge].
+ *
+ * This cache is safe for multithreaded usage, and is recommended for objects that take a while to
+ * resolve (such as drawables, or things that require binder calls). As such, [getOrFetch] is
+ * recommended to be run on a background thread, while [purge] can be done from any thread.
+ */
+@SuppressLint("DumpableNotRegistered") // this will be dumped by container classes
+class NotifCollectionCache<V>(
+    private val retainCount: Int = 1,
+    private val purgeTimeoutMillis: Long = 1000L,
+    private val systemClock: SystemClock = SystemClockImpl(),
+) : Dumpable {
+    @get:VisibleForTesting val cache = ConcurrentHashMap<String, CacheEntry>()
+
+    // Counters for cache hits and misses to be used to calculate and dump the hit ratio
+    @get:VisibleForTesting val misses = AtomicInteger(0)
+    @get:VisibleForTesting val hits = AtomicInteger(0)
+
+    init {
+        if (retainCount < 0) {
+            throw IllegalArgumentException("retainCount cannot be negative")
+        }
+    }
+
+    inner class CacheEntry(val key: String, val value: V) {
+        /**
+         * The "lives" represent how many times the entry will remain in the cache when purging it
+         * is attempted.
+         */
+        @get:VisibleForTesting var lives: Int = retainCount + 1
+        /**
+         * The last time this entry lost a "life". Starts at a negative value chosen so that the
+         * first purge is always considered "valid".
+         */
+        private var lastValidPurge: Long = -purgeTimeoutMillis
+
+        fun resetLives() {
+            // Lives/timeouts don't matter if retainCount is 0
+            if (retainCount == 0) {
+                return
+            }
+
+            synchronized(key) {
+                lives = retainCount + 1
+                lastValidPurge = -purgeTimeoutMillis
+            }
+            // Add it to the cache again just in case it was deleted before we could reset the lives
+            cache[key] = this
+        }
+
+        fun tryPurge(): Boolean {
+            // Lives/timeouts don't matter if retainCount is 0
+            if (retainCount == 0) {
+                return true
+            }
+
+            // Using uptimeMillis since it's guaranteed to be monotonic, as we don't want a
+            // timezone/clock change to break us
+            val now = systemClock.uptimeMillis()
+
+            // Cannot purge the same entry from two threads simultaneously
+            synchronized(key) {
+                if (now - lastValidPurge < purgeTimeoutMillis) {
+                    return false
+                }
+                lastValidPurge = now
+                return --lives <= 0
+            }
+        }
+    }
+
+    /**
+     * Get value from cache, or fetch it and add it to cache if not found. This can be called from
+     * any thread, but is usually expected to be called from the background.
+     *
+     * @param key key for the object to be obtained
+     * @param fetch method to fetch the object and add it to the cache if not present; note that
+     *   there is no guarantee that two [fetch] cannot run in parallel for the same [key] (if
+     *   [getOrFetch] is called simultaneously from different threads), so be mindful of potential
+     *   side effects
+     */
+    fun getOrFetch(key: String, fetch: (String) -> V): V {
+        val entry = cache[key]
+        if (entry != null) {
+            hits.incrementAndGet()
+            // Refresh lives on access
+            entry.resetLives()
+            return entry.value
+        }
+
+        misses.incrementAndGet()
+        val value = fetch(key)
+        cache[key] = CacheEntry(key, value)
+        return value
+    }
+
+    /**
+     * Clear entries that are NOT in [wantedKeys] if appropriate. This can be called from any
+     * thread.
+     *
+     * If retainCount > 0, a given entry will need to not be present in [wantedKeys] for
+     * ([retainCount] + 1) consecutive [purge] calls made within at least [purgeTimeoutMillis] of
+     * each other in order to be cleared. This count will be reset for any given entry 1) if
+     * [getOrFetch] is called for the entry or 2) if the entry is present in [wantedKeys] in a
+     * subsequent [purge] call. We prioritize keeping the entry if possible, so if [purge] is called
+     * simultaneously with [getOrFetch] on different threads for example, we will try to keep it in
+     * the cache, although it is not guaranteed. If avoiding cache misses is a concern, consider
+     * increasing the [retainCount] or [purgeTimeoutMillis].
+     *
+     * For example, say [retainCount] = 1 and [purgeTimeoutMillis] = 1000 and we start with entries
+     * (a, b, c) in the cache:
+     * ```kotlin
+     * purge((a, c)); // marks b for deletion
+     * Thread.sleep(500)
+     * purge((a, c)); // does nothing as it was called earlier than the min 1s
+     * Thread.sleep(500)
+     * purge((b, c)); // b is no longer marked for deletion, but now a is
+     * Thread.sleep(1000);
+     * purge((c));    // deletes a from the cache and marks b for deletion, etc.
+     * ```
+     */
+    fun purge(wantedKeys: List<String>) {
+        for ((key, entry) in cache) {
+            if (key in wantedKeys) {
+                entry.resetLives()
+            } else if (entry.tryPurge()) {
+                cache.remove(key)
+            }
+        }
+    }
+
+    /** Clear all entries from the cache. */
+    fun clear() {
+        cache.clear()
+    }
+
+    override fun dump(pwOrig: PrintWriter, args: Array<out String>) {
+        val pw = pwOrig.asIndenting()
+
+        pw.println("$TAG(retainCount = $retainCount, purgeTimeoutMillis = $purgeTimeoutMillis)")
+        pw.withIncreasedIndent {
+            pw.printCollection("keys present in cache", cache.keys.stream().sorted().toList())
+
+            val misses = misses.get()
+            val hits = hits.get()
+            pw.println(
+                "cache hit ratio = ${(hits.toFloat() / (hits + misses)) * 100}% " +
+                    "($hits hits, $misses misses)"
+            )
+        }
+    }
+
+    companion object {
+        const val TAG = "NotifCollectionCache"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index 38e6609..933f793 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -262,6 +262,11 @@
      */
     private boolean mIsHeadsUp;
 
+    /**
+     * Whether or not the notification is showing the app icon instead of the small icon.
+     */
+    private boolean mIsShowingAppIcon;
+
     private boolean mLastChronometerRunning = true;
     private ViewStub mChildrenContainerStub;
     private GroupMembershipManager mGroupMembershipManager;
@@ -816,6 +821,20 @@
         }
     }
 
+    /**
+     * Indicate that the notification is showing the app icon instead of the small icon.
+     */
+    public void setIsShowingAppIcon(boolean isShowingAppIcon) {
+        mIsShowingAppIcon = isShowingAppIcon;
+    }
+
+    /**
+     * Whether or not the notification is showing the app icon instead of the small icon.
+     */
+    public boolean isShowingAppIcon() {
+        return mIsShowingAppIcon;
+    }
+
     @Override
     public boolean showingPulsing() {
         return isHeadsUpState() && (isDozing() || (mOnKeyguard && isBypassEnabled()));
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt
index 79defd2..7b85bfd 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/icon/NotificationRowIconViewInflaterFactory.kt
@@ -21,6 +21,7 @@
 import android.util.AttributeSet
 import android.view.View
 import com.android.internal.widget.NotificationRowIconView
+import com.android.internal.widget.NotificationRowIconView.NotificationIconProvider
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.NotifRemoteViewsFactory
 import com.android.systemui.statusbar.notification.row.NotificationRowContentBinder
@@ -47,20 +48,27 @@
         return when (name) {
             NotificationRowIconView::class.java.name ->
                 NotificationRowIconView(context, attrs).also { view ->
-                    val sbn = row.entry.sbn
-                    view.setIconProvider(
-                        object : NotificationRowIconView.NotificationIconProvider {
-                            override fun shouldShowAppIcon(): Boolean {
-                                return iconStyleProvider.shouldShowAppIcon(row.entry.sbn, context)
-                            }
-
-                            override fun getAppIcon(): Drawable {
-                                return appIconProvider.getOrFetchAppIcon(sbn.packageName, context)
-                            }
-                        }
-                    )
+                    view.setIconProvider(createIconProvider(row, context))
                 }
             else -> null
         }
     }
+
+    private fun createIconProvider(
+        row: ExpandableNotificationRow,
+        context: Context,
+    ): NotificationIconProvider {
+        val sbn = row.entry.sbn
+        return object : NotificationIconProvider {
+            override fun shouldShowAppIcon(): Boolean {
+                val shouldShowAppIcon = iconStyleProvider.shouldShowAppIcon(row.entry.sbn, context)
+                row.setIsShowingAppIcon(shouldShowAppIcon)
+                return shouldShowAppIcon
+            }
+
+            override fun getAppIcon(): Drawable {
+                return appIconProvider.getOrFetchAppIcon(sbn.packageName, context)
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt
index b4411f1..f8aff69 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationConversationTemplateViewWrapper.kt
@@ -16,15 +16,18 @@
 
 package com.android.systemui.statusbar.notification.row.wrapper
 
+import android.app.Flags
 import android.content.Context
 import android.graphics.drawable.AnimatedImageDrawable
 import android.view.View
 import android.view.ViewGroup
+import android.view.ViewGroup.MarginLayoutParams
 import com.android.internal.widget.CachingIconView
 import com.android.internal.widget.ConversationLayout
 import com.android.internal.widget.MessagingGroup
 import com.android.internal.widget.MessagingImageMessage
 import com.android.internal.widget.MessagingLinearLayout
+import com.android.internal.widget.NotificationRowIconView
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.notification.NotificationFadeAware
 import com.android.systemui.statusbar.notification.NotificationUtils
@@ -32,23 +35,23 @@
 import com.android.systemui.statusbar.notification.row.wrapper.NotificationMessagingTemplateViewWrapper.setCustomImageMessageTransform
 import com.android.systemui.util.children
 
-/**
- * Wraps a notification containing a conversation template
- */
-class NotificationConversationTemplateViewWrapper constructor(
+/** Wraps a notification containing a conversation template */
+class NotificationConversationTemplateViewWrapper(
     ctx: Context,
     view: View,
-    row: ExpandableNotificationRow
+    row: ExpandableNotificationRow,
 ) : NotificationTemplateViewWrapper(ctx, view, row) {
 
-    private val minHeightWithActions: Int = NotificationUtils.getFontScaledHeight(
+    private val minHeightWithActions: Int =
+        NotificationUtils.getFontScaledHeight(
             ctx,
-            R.dimen.notification_messaging_actions_min_height
-    )
+            R.dimen.notification_messaging_actions_min_height,
+        )
     private val conversationLayout: ConversationLayout = view as ConversationLayout
 
     private lateinit var conversationIconContainer: View
     private lateinit var conversationIconView: CachingIconView
+    private lateinit var badgeIconView: NotificationRowIconView
     private lateinit var conversationBadgeBg: View
     private lateinit var expandBtn: View
     private lateinit var expandBtnContainer: View
@@ -68,10 +71,13 @@
         messageContainers = conversationLayout.messagingGroups
         with(conversationLayout) {
             conversationIconContainer =
-                    requireViewById(com.android.internal.R.id.conversation_icon_container)
+                requireViewById(com.android.internal.R.id.conversation_icon_container)
             conversationIconView = requireViewById(com.android.internal.R.id.conversation_icon)
+            if (Flags.notificationsRedesignAppIcons()) {
+                badgeIconView = requireViewById(com.android.internal.R.id.icon)
+            }
             conversationBadgeBg =
-                    requireViewById(com.android.internal.R.id.conversation_icon_badge_bg)
+                requireViewById(com.android.internal.R.id.conversation_icon_badge_bg)
             expandBtn = requireViewById(com.android.internal.R.id.expand_button)
             expandBtnContainer = requireViewById(com.android.internal.R.id.expand_button_container)
             importanceRing = requireViewById(com.android.internal.R.id.conversation_icon_badge_ring)
@@ -80,7 +86,7 @@
             facePileTop = findViewById(com.android.internal.R.id.conversation_face_pile_top)
             facePileBottom = findViewById(com.android.internal.R.id.conversation_face_pile_bottom)
             facePileBottomBg =
-                    findViewById(com.android.internal.R.id.conversation_face_pile_bottom_background)
+                findViewById(com.android.internal.R.id.conversation_face_pile_bottom_background)
         }
     }
 
@@ -88,6 +94,13 @@
         // Reinspect the notification. Before the super call, because the super call also updates
         // the transformation types and we need to have our values set by then.
         resolveViews()
+        if (Flags.notificationsRedesignAppIcons() && row.isShowingAppIcon) {
+            // Override the margins to be 2dp instead of 4dp according to the new design if we're
+            // showing the app icon.
+            val lp = badgeIconView.layoutParams as MarginLayoutParams
+            lp.setMargins(2, 2, 2, 2)
+            badgeIconView.layoutParams = lp
+        }
         super.onContentUpdated(row)
     }
 
@@ -96,56 +109,50 @@
         super.updateTransformedTypes()
 
         mTransformationHelper.addTransformedView(TRANSFORMING_VIEW_TITLE, conversationTitleView)
-        addTransformedViews(
-                messagingLinearLayout,
-                appName
-        )
+        addTransformedViews(messagingLinearLayout, appName)
 
         setCustomImageMessageTransform(mTransformationHelper, imageMessageContainer)
 
         addViewsTransformingToSimilar(
-                conversationIconView,
-                conversationBadgeBg,
-                expandBtn,
-                importanceRing,
-                facePileTop,
-                facePileBottom,
-                facePileBottomBg
+            conversationIconView,
+            conversationBadgeBg,
+            expandBtn,
+            importanceRing,
+            facePileTop,
+            facePileBottom,
+            facePileBottomBg,
         )
     }
 
     override fun getShelfTransformationTarget(): View? =
-            if (conversationLayout.isImportantConversation)
-                if (conversationIconView.visibility != View.GONE)
-                    conversationIconView
-                else
-                    // A notification with a fallback icon was set to important. Currently
-                    // the transformation doesn't work for these and needs to be fixed.
-                    // In the meantime those are using the icon.
-                    super.getShelfTransformationTarget()
+        if (conversationLayout.isImportantConversation)
+            if (conversationIconView.visibility != View.GONE) conversationIconView
             else
-                super.getShelfTransformationTarget()
+            // A notification with a fallback icon was set to important. Currently
+            // the transformation doesn't work for these and needs to be fixed.
+            // In the meantime those are using the icon.
+            super.getShelfTransformationTarget()
+        else super.getShelfTransformationTarget()
 
     override fun setRemoteInputVisible(visible: Boolean) =
-            conversationLayout.showHistoricMessages(visible)
+        conversationLayout.showHistoricMessages(visible)
 
     override fun updateExpandability(
         expandable: Boolean,
         onClickListener: View.OnClickListener,
-        requestLayout: Boolean
+        requestLayout: Boolean,
     ) = conversationLayout.updateExpandability(expandable, onClickListener)
 
     override fun disallowSingleClick(x: Float, y: Float): Boolean {
-        val isOnExpandButton = expandBtnContainer.visibility == View.VISIBLE &&
-                isOnView(expandBtnContainer, x, y)
+        val isOnExpandButton =
+            expandBtnContainer.visibility == View.VISIBLE && isOnView(expandBtnContainer, x, y)
         return isOnExpandButton || super.disallowSingleClick(x, y)
     }
 
     override fun getMinLayoutHeight(): Int =
-            if (mActionsContainer != null && mActionsContainer.visibility != View.GONE)
-                minHeightWithActions
-            else
-                super.getMinLayoutHeight()
+        if (mActionsContainer != null && mActionsContainer.visibility != View.GONE)
+            minHeightWithActions
+        else super.getMinLayoutHeight()
 
     override fun setNotificationFaded(faded: Boolean) {
         // Do not call super
@@ -157,16 +164,17 @@
     override fun setAnimationsRunning(running: Boolean) {
         // We apply to both the child message containers in a conversation group,
         // and the top level image message container.
-        val containers = messageContainers.asSequence().map { it.messageContainer } +
+        val containers =
+            messageContainers.asSequence().map { it.messageContainer } +
                 sequenceOf(imageMessageContainer)
         val drawables =
-                containers
-                        .flatMap { it.children }
-                        .mapNotNull { child ->
-                            (child as? MessagingImageMessage)?.let { imageMessage ->
-                                imageMessage.drawable as? AnimatedImageDrawable
-                            }
-                        }
+            containers
+                .flatMap { it.children }
+                .mapNotNull { child ->
+                    (child as? MessagingImageMessage)?.let { imageMessage ->
+                        imageMessage.drawable as? AnimatedImageDrawable
+                    }
+                }
         drawables.toSet().forEach {
             when {
                 running -> it.start()
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/StatusBarSystemEventAnimator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/StatusBarSystemEventAnimator.kt
index e73063b..1f9ea08 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/StatusBarSystemEventAnimator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/fragment/StatusBarSystemEventAnimator.kt
@@ -16,11 +16,11 @@
 
 package com.android.systemui.statusbar.phone.fragment
 
+import android.content.res.Resources
+import android.view.View
 import androidx.core.animation.Animator
 import androidx.core.animation.AnimatorSet
 import androidx.core.animation.ValueAnimator
-import android.content.res.Resources
-import android.view.View
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.events.STATUS_BAR_X_MOVE_IN
 import com.android.systemui.statusbar.events.STATUS_BAR_X_MOVE_OUT
@@ -33,16 +33,15 @@
  * An implementation of [StatusBarSystemEventDefaultAnimator], applying the onAlphaChanged and
  * onTranslationXChanged callbacks directly to the provided animatedView.
  */
-class StatusBarSystemEventAnimator @JvmOverloads constructor(
-        val animatedView: View,
-        resources: Resources,
-        isAnimationRunning: Boolean = false
-) : StatusBarSystemEventDefaultAnimator(
+class StatusBarSystemEventAnimator
+@JvmOverloads
+constructor(val animatedView: View, resources: Resources, isAnimationRunning: Boolean = false) :
+    StatusBarSystemEventDefaultAnimator(
         resources = resources,
         onAlphaChanged = animatedView::setAlpha,
         onTranslationXChanged = animatedView::setTranslationX,
-        isAnimationRunning = isAnimationRunning
-)
+        isAnimationRunning = isAnimationRunning,
+    )
 
 /**
  * Tied directly to [SystemStatusAnimationScheduler]. Any StatusBar-like thing (keyguard, collapsed
@@ -53,34 +52,39 @@
  * this class could be used directly as the animation callback, it's probably best to forward calls
  * to it so that it can be recreated at any moment without needing to remove/add callback.
  */
-
-open class StatusBarSystemEventDefaultAnimator @JvmOverloads constructor(
-        resources: Resources,
-        private val onAlphaChanged: (Float) -> Unit,
-        private val onTranslationXChanged: (Float) -> Unit,
-        var isAnimationRunning: Boolean = false
+open class StatusBarSystemEventDefaultAnimator
+@JvmOverloads
+constructor(
+    resources: Resources,
+    private val onAlphaChanged: (Float) -> Unit,
+    private val onTranslationXChanged: (Float) -> Unit,
+    var isAnimationRunning: Boolean = false,
 ) : SystemStatusAnimationCallback {
-    private val translationXIn: Int = resources.getDimensionPixelSize(
-            R.dimen.ongoing_appops_chip_animation_in_status_bar_translation_x)
-    private val translationXOut: Int = resources.getDimensionPixelSize(
-            R.dimen.ongoing_appops_chip_animation_out_status_bar_translation_x)
+    private val translationXIn: Int =
+        resources.getDimensionPixelSize(
+            R.dimen.ongoing_appops_chip_animation_in_status_bar_translation_x
+        )
+    private val translationXOut: Int =
+        resources.getDimensionPixelSize(
+            R.dimen.ongoing_appops_chip_animation_out_status_bar_translation_x
+        )
 
     override fun onSystemEventAnimationBegin(): Animator {
         isAnimationRunning = true
-        val moveOut = ValueAnimator.ofFloat(0f, 1f).apply {
-            duration = 23.frames
-            interpolator = STATUS_BAR_X_MOVE_OUT
-            addUpdateListener {
-                onTranslationXChanged(-(translationXIn * animatedValue as Float))
+        val moveOut =
+            ValueAnimator.ofFloat(0f, 1f).apply {
+                duration = 23.frames
+                interpolator = STATUS_BAR_X_MOVE_OUT
+                addUpdateListener {
+                    onTranslationXChanged(-(translationXIn * animatedValue as Float))
+                }
             }
-        }
-        val alphaOut = ValueAnimator.ofFloat(1f, 0f).apply {
-            duration = 8.frames
-            interpolator = null
-            addUpdateListener {
-                onAlphaChanged(animatedValue as Float)
+        val alphaOut =
+            ValueAnimator.ofFloat(1f, 0f).apply {
+                duration = 8.frames
+                interpolator = null
+                addUpdateListener { onAlphaChanged(animatedValue as Float) }
             }
-        }
 
         val animSet = AnimatorSet()
         animSet.playTogether(moveOut, alphaOut)
@@ -89,22 +93,22 @@
 
     override fun onSystemEventAnimationFinish(hasPersistentDot: Boolean): Animator {
         onTranslationXChanged(translationXOut.toFloat())
-        val moveIn = ValueAnimator.ofFloat(1f, 0f).apply {
-            duration = 23.frames
-            startDelay = 7.frames
-            interpolator = STATUS_BAR_X_MOVE_IN
-            addUpdateListener {
-                onTranslationXChanged(translationXOut * animatedValue as Float)
+        val moveIn =
+            ValueAnimator.ofFloat(1f, 0f).apply {
+                duration = 23.frames
+                startDelay = 7.frames
+                interpolator = STATUS_BAR_X_MOVE_IN
+                addUpdateListener {
+                    onTranslationXChanged(translationXOut * animatedValue as Float)
+                }
             }
-        }
-        val alphaIn = ValueAnimator.ofFloat(0f, 1f).apply {
-            duration = 5.frames
-            startDelay = 11.frames
-            interpolator = null
-            addUpdateListener {
-                onAlphaChanged(animatedValue as Float)
+        val alphaIn =
+            ValueAnimator.ofFloat(0f, 1f).apply {
+                duration = 5.frames
+                startDelay = 11.frames
+                interpolator = null
+                addUpdateListener { onAlphaChanged(animatedValue as Float) }
             }
-        }
 
         val animatorSet = AnimatorSet()
         animatorSet.playTogether(moveIn, alphaIn)
@@ -112,4 +116,4 @@
         animatorSet.doOnCancel { isAnimationRunning = false }
         return animatorSet
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/CameraProtectionLoaderImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/CameraProtectionLoaderImplTest.kt
index d2a17c2..ad58a01 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/CameraProtectionLoaderImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/CameraProtectionLoaderImplTest.kt
@@ -37,14 +37,14 @@
         overrideResource(R.string.config_protectedPhysicalCameraId, OUTER_CAMERA_PHYSICAL_ID)
         overrideResource(
             R.string.config_frontBuiltInDisplayCutoutProtection,
-            OUTER_CAMERA_PROTECTION_PATH
+            OUTER_CAMERA_PROTECTION_PATH,
         )
         overrideResource(R.string.config_protectedScreenUniqueId, OUTER_SCREEN_UNIQUE_ID)
         overrideResource(R.string.config_protectedInnerCameraId, INNER_CAMERA_LOGICAL_ID)
         overrideResource(R.string.config_protectedInnerPhysicalCameraId, INNER_CAMERA_PHYSICAL_ID)
         overrideResource(
             R.string.config_innerBuiltInDisplayCutoutProtection,
-            INNER_CAMERA_PROTECTION_PATH
+            INNER_CAMERA_PROTECTION_PATH,
         )
         overrideResource(R.string.config_protectedInnerScreenUniqueId, INNER_SCREEN_UNIQUE_ID)
     }
@@ -107,7 +107,7 @@
         private const val OUTER_CAMERA_PHYSICAL_ID = "11"
         private const val OUTER_CAMERA_PROTECTION_PATH = "M 0,0 H 10,10 V 10,10 H 0,10 Z"
         private val OUTER_CAMERA_PROTECTION_BOUNDS =
-            Rect(/* left = */ 0, /* top = */ 0, /* right = */ 10, /* bottom = */ 10)
+            Rect(/* left= */ 0, /* top= */ 0, /* right= */ 10, /* bottom= */ 10)
         private const val OUTER_SCREEN_UNIQUE_ID = "111"
         private val OUTER_CAMERA_PROTECTION_INFO =
             TestableProtectionInfo(
@@ -121,7 +121,7 @@
         private const val INNER_CAMERA_PHYSICAL_ID = "22"
         private const val INNER_CAMERA_PROTECTION_PATH = "M 0,0 H 20,20 V 20,20 H 0,20 Z"
         private val INNER_CAMERA_PROTECTION_BOUNDS =
-            Rect(/* left = */ 0, /* top = */ 0, /* right = */ 20, /* bottom = */ 20)
+            Rect(/* left= */ 0, /* top= */ 0, /* right= */ 20, /* bottom= */ 20)
         private const val INNER_SCREEN_UNIQUE_ID = "222"
         private val INNER_CAMERA_PROTECTION_INFO =
             TestableProtectionInfo(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/FaceScanningProviderFactoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/FaceScanningProviderFactoryTest.kt
index bc12aaa..a01feca 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/FaceScanningProviderFactoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/FaceScanningProviderFactoryTest.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.biometrics.AuthController
 import com.android.systemui.biometrics.data.repository.FakeFacePropertyRepository
 import com.android.systemui.decor.FaceScanningProviderFactory
+import com.android.systemui.decor.FaceScanningProviderFactoryImpl
 import com.android.systemui.log.ScreenDecorationsLogger
 import com.android.systemui.log.logcatLogBuffer
 import com.android.systemui.plugins.statusbar.StatusBarStateController
@@ -69,20 +70,20 @@
                 dmGlobal,
                 displayId,
                 displayInfo,
-                DisplayAdjustments.DEFAULT_DISPLAY_ADJUSTMENTS
+                DisplayAdjustments.DEFAULT_DISPLAY_ADJUSTMENTS,
             )
         whenever(dmGlobal.getDisplayInfo(eq(displayId))).thenReturn(displayInfo)
         val displayContext = context.createDisplayContext(display) as SysuiTestableContext
         displayContext.orCreateTestableResources.addOverride(
             R.array.config_displayUniqueIdArray,
-            arrayOf(displayId)
+            arrayOf(displayId),
         )
         displayContext.orCreateTestableResources.addOverride(
             R.bool.config_fillMainBuiltInDisplayCutout,
-            true
+            true,
         )
         underTest =
-            FaceScanningProviderFactory(
+            FaceScanningProviderFactoryImpl(
                 authController,
                 displayContext,
                 statusBarStateController,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/SysUICutoutProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/SysUICutoutProviderTest.kt
index 61c7e1d..ef33210 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/SysUICutoutProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/SysUICutoutProviderTest.kt
@@ -42,7 +42,7 @@
     fun cutoutInfoForCurrentDisplay_noCutout_returnsNull() {
         val noCutoutDisplay = createDisplay(cutout = null)
         val noCutoutDisplayContext = context.createDisplayContext(noCutoutDisplay)
-        val provider = SysUICutoutProvider(noCutoutDisplayContext, fakeProtectionLoader)
+        val provider = SysUICutoutProviderImpl(noCutoutDisplayContext, fakeProtectionLoader)
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()
 
@@ -53,7 +53,7 @@
     fun cutoutInfoForCurrentDisplay_returnsCutout() {
         val cutoutDisplay = createDisplay()
         val cutoutDisplayContext = context.createDisplayContext(cutoutDisplay)
-        val provider = SysUICutoutProvider(cutoutDisplayContext, fakeProtectionLoader)
+        val provider = SysUICutoutProviderImpl(cutoutDisplayContext, fakeProtectionLoader)
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
@@ -64,7 +64,7 @@
     fun cutoutInfoForCurrentDisplay_noAssociatedProtection_returnsNoProtection() {
         val cutoutDisplay = createDisplay()
         val cutoutDisplayContext = context.createDisplayContext(cutoutDisplay)
-        val provider = SysUICutoutProvider(cutoutDisplayContext, fakeProtectionLoader)
+        val provider = SysUICutoutProviderImpl(cutoutDisplayContext, fakeProtectionLoader)
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
@@ -75,7 +75,7 @@
     fun cutoutInfoForCurrentDisplay_outerDisplay_protectionAssociated_returnsProtection() {
         fakeProtectionLoader.addOuterCameraProtection(displayUniqueId = OUTER_DISPLAY_UNIQUE_ID)
         val outerDisplayContext = context.createDisplayContext(OUTER_DISPLAY)
-        val provider = SysUICutoutProvider(outerDisplayContext, fakeProtectionLoader)
+        val provider = SysUICutoutProviderImpl(outerDisplayContext, fakeProtectionLoader)
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
@@ -86,7 +86,7 @@
     fun cutoutInfoForCurrentDisplay_outerDisplay_protectionNotAvailable_returnsNullProtection() {
         fakeProtectionLoader.clearProtectionInfoList()
         val outerDisplayContext = context.createDisplayContext(OUTER_DISPLAY)
-        val provider = SysUICutoutProvider(outerDisplayContext, fakeProtectionLoader)
+        val provider = SysUICutoutProviderImpl(outerDisplayContext, fakeProtectionLoader)
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
@@ -97,7 +97,7 @@
     fun cutoutInfoForCurrentDisplay_displayWithNullId_protectionsWithNoId_returnsNullProtection() {
         fakeProtectionLoader.addOuterCameraProtection(displayUniqueId = "")
         val displayContext = context.createDisplayContext(createDisplay(uniqueId = null))
-        val provider = SysUICutoutProvider(displayContext, fakeProtectionLoader)
+        val provider = SysUICutoutProviderImpl(displayContext, fakeProtectionLoader)
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
@@ -108,7 +108,7 @@
     fun cutoutInfoForCurrentDisplay_displayWithEmptyId_protectionsWithNoId_returnsNullProtection() {
         fakeProtectionLoader.addOuterCameraProtection(displayUniqueId = "")
         val displayContext = context.createDisplayContext(createDisplay(uniqueId = ""))
-        val provider = SysUICutoutProvider(displayContext, fakeProtectionLoader)
+        val provider = SysUICutoutProviderImpl(displayContext, fakeProtectionLoader)
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
@@ -123,15 +123,13 @@
                 displayHeight = 1000,
                 rotation = Surface.ROTATION_0,
                 protectionBounds =
-                    Rect(/* left = */ 440, /* top = */ 10, /* right = */ 490, /* bottom = */ 110)
+                    Rect(/* left= */ 440, /* top= */ 10, /* right= */ 490, /* bottom= */ 110),
             )
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
         assertThat(sysUICutout.cameraProtection!!.bounds)
-            .isEqualTo(
-                Rect(/* left = */ 440, /* top = */ 10, /* right = */ 490, /* bottom = */ 110)
-            )
+            .isEqualTo(Rect(/* left= */ 440, /* top= */ 10, /* right= */ 490, /* bottom= */ 110))
     }
 
     @Test
@@ -142,13 +140,13 @@
                 displayHeight = 1000,
                 rotation = Surface.ROTATION_90,
                 protectionBounds =
-                    Rect(/* left = */ 440, /* top = */ 10, /* right = */ 490, /* bottom = */ 110)
+                    Rect(/* left= */ 440, /* top= */ 10, /* right= */ 490, /* bottom= */ 110),
             )
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
         assertThat(sysUICutout.cameraProtection!!.bounds)
-            .isEqualTo(Rect(/* left = */ 10, /* top = */ 10, /* right = */ 110, /* bottom = */ 60))
+            .isEqualTo(Rect(/* left= */ 10, /* top= */ 10, /* right= */ 110, /* bottom= */ 60))
     }
 
     @Test
@@ -156,7 +154,7 @@
         val displayNaturalWidth = 500
         val displayNaturalHeight = 1000
         val originalProtectionBounds =
-            Rect(/* left = */ 440, /* top = */ 10, /* right = */ 490, /* bottom = */ 110)
+            Rect(/* left= */ 440, /* top= */ 10, /* right= */ 490, /* bottom= */ 110)
         // Safe copy as we don't know at which layer the mutation could happen
         val originalProtectionBoundsCopy = Rect(originalProtectionBounds)
         val display =
@@ -168,10 +166,10 @@
             )
         fakeProtectionLoader.addOuterCameraProtection(
             displayUniqueId = OUTER_DISPLAY_UNIQUE_ID,
-            bounds = originalProtectionBounds
+            bounds = originalProtectionBounds,
         )
         val provider =
-            SysUICutoutProvider(context.createDisplayContext(display), fakeProtectionLoader)
+            SysUICutoutProviderImpl(context.createDisplayContext(display), fakeProtectionLoader)
 
         // Here we get the rotated bounds once
         provider.cutoutInfoForCurrentDisplayAndRotation()
@@ -194,13 +192,13 @@
                 displayHeight = 1000,
                 rotation = Surface.ROTATION_180,
                 protectionBounds =
-                    Rect(/* left = */ 440, /* top = */ 10, /* right = */ 490, /* bottom = */ 110)
+                    Rect(/* left= */ 440, /* top= */ 10, /* right= */ 490, /* bottom= */ 110),
             )
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
         assertThat(sysUICutout.cameraProtection!!.bounds)
-            .isEqualTo(Rect(/* left = */ 10, /* top = */ 890, /* right = */ 60, /* bottom = */ 990))
+            .isEqualTo(Rect(/* left= */ 10, /* top= */ 890, /* right= */ 60, /* bottom= */ 990))
     }
 
     @Test
@@ -211,15 +209,13 @@
                 displayHeight = 1000,
                 rotation = Surface.ROTATION_270,
                 protectionBounds =
-                    Rect(/* left = */ 440, /* top = */ 10, /* right = */ 490, /* bottom = */ 110)
+                    Rect(/* left= */ 440, /* top= */ 10, /* right= */ 490, /* bottom= */ 110),
             )
 
         val sysUICutout = provider.cutoutInfoForCurrentDisplayAndRotation()!!
 
         assertThat(sysUICutout.cameraProtection!!.bounds)
-            .isEqualTo(
-                Rect(/* left = */ 890, /* top = */ 440, /* right = */ 990, /* bottom = */ 490)
-            )
+            .isEqualTo(Rect(/* left= */ 890, /* top= */ 440, /* right= */ 990, /* bottom= */ 490))
     }
 
     private fun setUpProviderWithCameraProtection(
@@ -245,9 +241,9 @@
             )
         fakeProtectionLoader.addOuterCameraProtection(
             displayUniqueId = OUTER_DISPLAY_UNIQUE_ID,
-            bounds = protectionBounds
+            bounds = protectionBounds,
         )
-        return SysUICutoutProvider(context.createDisplayContext(display), fakeProtectionLoader)
+        return SysUICutoutProviderImpl(context.createDisplayContext(display), fakeProtectionLoader)
     }
 
     companion object {
@@ -259,7 +255,7 @@
             height: Int = 1000,
             @Rotation rotation: Int = Surface.ROTATION_0,
             uniqueId: String? = "uniqueId",
-            cutout: DisplayCutout? = mock<DisplayCutout>()
+            cutout: DisplayCutout? = mock<DisplayCutout>(),
         ) =
             mock<Display> {
                 whenever(this.getDisplayInfo(any())).thenAnswer {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt
index fd550b0..6e36d42b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt
@@ -88,7 +88,7 @@
                     },
                 layout = BouncerSceneLayout.BESIDE_USER_SWITCHER,
                 modifier = Modifier.fillMaxSize().testTag("BouncerContent"),
-                dialogFactory = bouncerDialogFactory
+                dialogFactory = bouncerDialogFactory,
             )
         }
     }
@@ -110,11 +110,19 @@
                             }
                         }
                     ) {
-                        feature(hasTestTag("UserSwitcher"), positionInRoot, "userSwitcher_pos")
-                        feature(hasTestTag("UserSwitcher"), alpha, "userSwitcher_alpha")
+                        feature(
+                            hasTestTag("com.android.systemui:id/UserSwitcher"),
+                            positionInRoot,
+                            "userSwitcher_pos",
+                        )
+                        feature(
+                            hasTestTag("com.android.systemui:id/UserSwitcher"),
+                            alpha,
+                            "userSwitcher_alpha",
+                        )
                         feature(hasTestTag("FoldAware"), positionInRoot, "foldAware_pos")
                         feature(hasTestTag("FoldAware"), alpha, "foldAware_alpha")
-                    }
+                    },
                 )
 
             assertThat(motion).timeSeriesMatchesGolden()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java
index 8731853..63ec78fd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/dialog/MediaOutputAdapterTest.java
@@ -35,6 +35,8 @@
 import android.app.WallpaperColors;
 import android.graphics.Bitmap;
 import android.graphics.drawable.Icon;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.testing.TestableLooper;
 import android.view.View;
 import android.widget.LinearLayout;
@@ -44,6 +46,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
+import com.android.media.flags.Flags;
 import com.android.settingslib.media.LocalMediaManager;
 import com.android.settingslib.media.MediaDevice;
 import com.android.systemui.SysuiTestCase;
@@ -738,4 +741,68 @@
 
         assertThat(mMediaOutputAdapter.getItemCount()).isEqualTo(updatedList.size());
     }
+
+    @DisableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagDisabled_InputDeviceMutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(true /* isInputDevice */, true /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.media_output_icon_volume_off);
+    }
+
+    @DisableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagDisabled_OutputDeviceMutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(false /* isInputDevice */, true /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.media_output_icon_volume_off);
+    }
+
+    @DisableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagDisabled_InputDeviceUnmutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(true /* isInputDevice */, false /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.media_output_icon_volume);
+    }
+
+    @DisableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagDisabled_OutputDeviceUnmutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(false /* isInputDevice */, false /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.media_output_icon_volume);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagEnabled_InputDeviceMutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(true /* isInputDevice */, true /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.ic_mic_off);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagEnabled_OutputDeviceMutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(false /* isInputDevice */, true /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.media_output_icon_volume_off);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagEnabled_InputDeviceUnmutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(true /* isInputDevice */, false /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.ic_mic_26dp);
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_AUDIO_INPUT_DEVICE_ROUTING_AND_VOLUME_CONTROL)
+    @Test
+    public void getDrawableId_FlagEnabled_OutputDeviceUnmutedIcon() {
+        assertThat(
+                mViewHolder.getDrawableId(false /* isInputDevice */, false /* isMutedVolumeIcon */))
+                .isEqualTo(R.drawable.media_output_icon_volume);
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt
index 8d060e9..8a6df1c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt
@@ -67,6 +67,8 @@
             onRemoveTile = {},
             onSetTiles = onSetTiles,
             onResize = { _, _ -> },
+            onStopEditing = {},
+            onReset = null,
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt
index ee1c0e9..d9c1d99 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt
@@ -66,6 +66,8 @@
             onRemoveTile = {},
             onSetTiles = {},
             onResize = onResize,
+            onStopEditing = {},
+            onReset = null,
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
index 4959224..3bfde68 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
@@ -35,7 +35,6 @@
 import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractor
 import com.android.systemui.keyguard.KeyguardUnlockAnimationController
 import com.android.systemui.keyguard.WakefulnessLifecycle
 import com.android.systemui.keyguard.ui.view.InWindowLauncherUnlockAnimationManager
@@ -122,9 +121,6 @@
         Optional<UnfoldTransitionProgressForwarder>
     @Mock private lateinit var broadcastDispatcher: BroadcastDispatcher
 
-    @Mock
-    private lateinit var keyboardTouchpadEduStatsInteractor: KeyboardTouchpadEduStatsInteractor
-
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
@@ -293,7 +289,6 @@
             dumpManager,
             unfoldTransitionProgressForwarder,
             broadcastDispatcher,
-            keyboardTouchpadEduStatsInteractor,
         )
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
index 0d4cb4c..7709a65 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
@@ -20,30 +20,141 @@
 import android.graphics.Insets
 import android.graphics.Rect
 import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import android.view.Display.DEFAULT_DISPLAY
 import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_CHORD
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.systemui.Flags
+import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.screenshot.ImageCapture
 import com.android.systemui.screenshot.ScreenshotData
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.emptyDisplayContent
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.launcherOnly
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.repository.DisplayContentRepository
+import com.android.systemui.screenshot.data.repository.profileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
 import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
 import com.android.systemui.screenshot.policy.TestUserIds.WORK
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @RunWith(AndroidJUnit4::class)
 class PolicyRequestProcessorTest {
+    private val kosmos = Kosmos()
+
+    private val screenshotRequest =
+        ScreenshotData(
+            TAKE_SCREENSHOT_FULLSCREEN,
+            SCREENSHOT_KEY_CHORD,
+            UserHandle.CURRENT,
+            topComponent = null,
+            originalScreenBounds = FULL_SCREEN,
+            taskId = -1,
+            originalInsets = Insets.NONE,
+            bitmap = null,
+            displayId = DEFAULT_DISPLAY,
+        )
+
+    val defaultComponent = ComponentName("default", "Component")
+    val defaultOwner = UserHandle.of(PERSONAL)
+
+    @get:Rule val setFlagsRule: SetFlagsRule = SetFlagsRule()
+
+    /** Tests applying CaptureParameters with 'IsolatedTask' CaptureType */
+    @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
+    fun testProcess_newPolicy_isolatedTask() = runTest {
+        val taskImage = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(task = taskImage),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
+                policies = emptyList(),
+                defaultOwner = defaultOwner,
+                defaultComponent = defaultComponent,
+                displayTasks = { emptyDisplayContent },
+            )
+
+        val result =
+            requestProcessor.modify(
+                screenshotRequest,
+                CaptureParameters(
+                    IsolatedTask(taskId = TASK_ID, taskBounds = null),
+                    ComponentName.unflattenFromString(FILES),
+                    UserHandle.of(WORK),
+                ),
+            )
+
+        assertWithMessage("The screenshot bitmap").that(result.bitmap).isSameInstanceAs(taskImage)
+
+        assertWithMessage("The assigned owner of the screenshot")
+            .that(result.userHandle)
+            .isEqualTo(UserHandle.of(WORK))
+
+        assertWithMessage("The topComponent of the screenshot")
+            .that(result.topComponent)
+            .isEqualTo(ComponentName.unflattenFromString(FILES))
+
+        assertWithMessage("Task ID").that(result.taskId).isEqualTo(TASK_ID)
+    }
+
+    /** Tests applying CaptureParameters with 'FullScreen' CaptureType */
+    @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
+    fun testProcess_newPolicy_fullScreen() = runTest {
+        val screenImage = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(display = screenImage),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
+                policies = emptyList(),
+                defaultOwner = defaultOwner,
+                defaultComponent = defaultComponent,
+                displayTasks = { emptyDisplayContent },
+            )
+
+        val result =
+            requestProcessor.modify(
+                screenshotRequest,
+                CaptureParameters(FullScreen(displayId = 0), defaultComponent, defaultOwner),
+            )
+
+        assertWithMessage("The result bitmap").that(result.bitmap).isSameInstanceAs(screenImage)
+
+        assertWithMessage("The assigned owner of the screenshot")
+            .that(result.userHandle)
+            .isEqualTo(defaultOwner)
+
+        assertWithMessage("The topComponent of the screenshot")
+            .that(result.topComponent)
+            .isEqualTo(defaultComponent)
+
+        assertWithMessage("Task ID").that(result.taskId).isEqualTo(-1)
+    }
+
     /** Tests behavior when no policies are applied */
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
     fun testProcess_defaultOwner_whenNoPolicyApplied() {
         val fullScreenWork = DisplayContentRepository {
             singleFullScreen(TaskSpec(taskId = TASK_ID, name = FILES, userId = WORK))
@@ -67,6 +178,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultOwner = UserHandle.of(PERSONAL),
                 defaultComponent = ComponentName("default", "Component"),
@@ -95,6 +207,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(display = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = DisplayContentRepository { launcherOnly() },
@@ -118,7 +231,7 @@
                 reason = "",
                 parameters =
                     CaptureParameters(
-                        CaptureType.IsolatedTask(taskId = 0, taskBounds = null),
+                        IsolatedTask(taskId = 0, taskBounds = null),
                         null,
                         UserHandle.CURRENT,
                     ),
@@ -130,6 +243,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(task = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = listOf(captureTaskPolicy),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = fullScreenWork,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImplTest.kt
index 5d8a8fd..4e7de81 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemStatusAnimationSchedulerImplTest.kt
@@ -88,30 +88,26 @@
             .thenReturn(statusBarWindowController)
         systemClock = FakeSystemClock()
         chipAnimationController =
-            SystemEventChipAnimationController(
+            SystemEventChipAnimationControllerImpl(
                 mContext,
-                statusBarWindowControllerStore,
+                statusBarWindowController,
                 statusBarContentInsetProvider,
             )
 
         // StatusBarContentInsetProvider is mocked. Ensure that it returns some mocked values.
         whenever(statusBarContentInsetProvider.getStatusBarContentInsetsForCurrentRotation())
-            .thenReturn(
-                Insets.of(/* left = */ 10, /* top = */ 10, /* right = */ 10, /* bottom = */ 0)
-            )
+            .thenReturn(Insets.of(/* left= */ 10, /* top= */ 10, /* right= */ 10, /* bottom= */ 0))
         whenever(statusBarContentInsetProvider.getStatusBarContentAreaForCurrentRotation())
-            .thenReturn(
-                Rect(/* left = */ 10, /* top = */ 10, /* right = */ 990, /* bottom = */ 100)
-            )
+            .thenReturn(Rect(/* left= */ 10, /* top= */ 10, /* right= */ 990, /* bottom= */ 100))
 
         // StatusBarWindowController is mocked. The addViewToWindow function needs to be mocked to
         // ensure that the chip view is added to a parent view
         whenever(statusBarWindowController.addViewToWindow(any(), any())).then {
             val statusbarFake = FrameLayout(mContext)
-            statusbarFake.layout(/* l = */ 0, /* t = */ 0, /* r = */ 1000, /* b = */ 100)
+            statusbarFake.layout(/* l= */ 0, /* t= */ 0, /* r= */ 1000, /* b= */ 100)
             statusbarFake.addView(
                 it.arguments[0] as View,
-                it.arguments[1] as FrameLayout.LayoutParams
+                it.arguments[1] as FrameLayout.LayoutParams,
             )
         }
     }
@@ -386,7 +382,7 @@
         scheduleFakeEventWithView(
             accessibilityDesc,
             mockAnimatableView,
-            shouldAnnounceAccessibilityEvent = true
+            shouldAnnounceAccessibilityEvent = true,
         )
         fastForwardAnimationToState(ANIMATING_OUT)
 
@@ -405,7 +401,7 @@
         scheduleFakeEventWithView(
             accessibilityDesc,
             mockAnimatableView,
-            shouldAnnounceAccessibilityEvent = true
+            shouldAnnounceAccessibilityEvent = true,
         )
         fastForwardAnimationToState(ANIMATING_OUT)
 
@@ -424,7 +420,7 @@
         scheduleFakeEventWithView(
             accessibilityDesc,
             mockAnimatableView,
-            shouldAnnounceAccessibilityEvent = false
+            shouldAnnounceAccessibilityEvent = false,
         )
         fastForwardAnimationToState(ANIMATING_OUT)
 
@@ -637,13 +633,13 @@
     private fun scheduleFakeEventWithView(
         desc: String?,
         view: BackgroundAnimatableView,
-        shouldAnnounceAccessibilityEvent: Boolean
+        shouldAnnounceAccessibilityEvent: Boolean,
     ) {
         val fakeEvent =
             FakeStatusEvent(
                 viewCreator = { view },
                 contentDescription = desc,
-                shouldAnnounceAccessibilityEvent = shouldAnnounceAccessibilityEvent
+                shouldAnnounceAccessibilityEvent = shouldAnnounceAccessibilityEvent,
             )
         systemStatusAnimationScheduler.onStatusEvent(fakeEvent)
     }
@@ -668,7 +664,7 @@
                 dumpManager,
                 systemClock,
                 CoroutineScope(StandardTestDispatcher(testScope.testScheduler)),
-                logger
+                logger,
             )
         // add a mock listener
         systemStatusAnimationScheduler.addCallback(listener)
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
index 2d275f9..3fd2503 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.recents.OverviewProxyService
 import com.android.systemui.touchpad.data.repository.touchpadRepository
 import com.android.systemui.user.data.repository.userRepository
 import org.mockito.kotlin.mock
@@ -38,27 +39,13 @@
                     testDispatcher,
                     keyboardRepository,
                     touchpadRepository,
-                    userRepository
-                ),
-            clock = fakeEduClock
-        )
-    }
-
-var Kosmos.mockEduInputManager by Kosmos.Fixture { mock<InputManager>() }
-
-var Kosmos.keyboardTouchpadEduStatsInteractor by
-    Kosmos.Fixture {
-        KeyboardTouchpadEduStatsInteractorImpl(
-            backgroundScope = testScope.backgroundScope,
-            contextualEducationInteractor = contextualEducationInteractor,
-            inputDeviceRepository =
-                UserInputDeviceRepository(
-                    testDispatcher,
-                    keyboardRepository,
-                    touchpadRepository,
-                    userRepository
+                    userRepository,
                 ),
             tutorialSchedulerRepository,
-            fakeEduClock
+            mockOverviewProxyService,
+            clock = fakeEduClock,
         )
     }
+
+var Kosmos.mockOverviewProxyService by Kosmos.Fixture { mock<OverviewProxyService>() }
+var Kosmos.mockEduInputManager by Kosmos.Fixture { mock<InputManager>() }
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 c2a03d4..fbfaba6 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
@@ -33,7 +33,6 @@
 import com.android.systemui.keyboard.shortcut.data.source.SystemShortcutsSource
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperCategoriesInteractor
 import com.android.systemui.keyboard.shortcut.domain.interactor.ShortcutHelperStateInteractor
-import com.android.systemui.keyboard.shortcut.ui.ShortcutHelperActivityStarter
 import com.android.systemui.keyboard.shortcut.ui.viewmodel.ShortcutHelperViewModel
 import com.android.systemui.keyguard.data.repository.fakeCommandQueue
 import com.android.systemui.kosmos.Kosmos
@@ -45,12 +44,7 @@
 import com.android.systemui.settings.fakeUserTracker
 
 var Kosmos.shortcutHelperAppCategoriesShortcutsSource: KeyboardShortcutGroupsSource by
-    Kosmos.Fixture {
-        AppCategoriesShortcutsSource(
-            windowManager,
-            testDispatcher,
-        )
-    }
+    Kosmos.Fixture { AppCategoriesShortcutsSource(windowManager, testDispatcher) }
 
 var Kosmos.shortcutHelperSystemShortcutsSource: KeyboardShortcutGroupsSource by
     Kosmos.Fixture { SystemShortcutsSource(mainResources) }
@@ -65,7 +59,7 @@
             broadcastDispatcher,
             fakeInputManager.inputManager,
             testScope,
-            testDispatcher
+            testDispatcher,
         )
     }
 
@@ -109,7 +103,7 @@
             displayTracker,
             testScope,
             sysUiState,
-            shortcutHelperStateRepository
+            shortcutHelperStateRepository,
         )
     }
 
@@ -124,18 +118,6 @@
             applicationCoroutineScope,
             testDispatcher,
             shortcutHelperStateInteractor,
-            shortcutHelperCategoriesInteractor
-        )
-    }
-
-val Kosmos.fakeShortcutHelperStartActivity by Kosmos.Fixture { FakeShortcutHelperStartActivity() }
-
-val Kosmos.shortcutHelperActivityStarter by
-    Kosmos.Fixture {
-        ShortcutHelperActivityStarter(
-            applicationContext,
-            applicationCoroutineScope,
-            shortcutHelperViewModel,
-            fakeShortcutHelperStartActivity,
+            shortcutHelperCategoriesInteractor,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelKosmos.kt
index c218ff6..dff5625 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelKosmos.kt
@@ -20,6 +20,7 @@
 import androidx.lifecycle.LifecycleCoroutineScope
 import com.android.systemui.common.ui.domain.interactor.configurationInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.qs.footerActionsController
 import com.android.systemui.qs.footerActionsViewModelFactory
@@ -30,7 +31,9 @@
 import com.android.systemui.shade.transition.largeScreenShadeInterpolator
 import com.android.systemui.statusbar.disableflags.data.repository.disableFlagsRepository
 import com.android.systemui.statusbar.sysuiStatusBarStateController
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 
+@OptIn(ExperimentalCoroutinesApi::class)
 val Kosmos.qsFragmentComposeViewModelFactory by
     Kosmos.Fixture {
         object : QSFragmentComposeViewModel.Factory {
@@ -45,6 +48,7 @@
                     sysuiStatusBarStateController,
                     deviceEntryInteractor,
                     disableFlagsRepository,
+                    keyguardTransitionInteractor,
                     largeScreenShadeInterpolator,
                     configurationInteractor,
                     largeScreenHeaderHelper,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractorFactoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractorFactoryKosmos.kt
new file mode 100644
index 0000000..a5fe8cf
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/DynamicIconTilesInteractorFactoryKosmos.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.qs.panels.domain.interactor
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
+
+val Kosmos.dynamicIconTilesInteractorFactory by
+    Kosmos.Fixture {
+        object : DynamicIconTilesInteractor.Factory {
+            override fun create(): DynamicIconTilesInteractor {
+                return DynamicIconTilesInteractor(iconTilesInteractor, currentTilesInteractor)
+            }
+        }
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
index b4317ad..b6b0a41 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/InfiniteGridLayoutKosmos.kt
@@ -19,10 +19,7 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.qs.panels.ui.compose.infinitegrid.InfiniteGridLayout
 import com.android.systemui.qs.panels.ui.viewmodel.iconTilesViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.qsColumnsViewModel
-import com.android.systemui.qs.panels.ui.viewmodel.tileSquishinessViewModel
+import com.android.systemui.qs.panels.ui.viewmodel.infiniteGridViewModelFactory
 
 val Kosmos.infiniteGridLayout by
-    Kosmos.Fixture {
-        InfiniteGridLayout(iconTilesViewModel, qsColumnsViewModel, tileSquishinessViewModel)
-    }
+    Kosmos.Fixture { InfiniteGridLayout(iconTilesViewModel, infiniteGridViewModelFactory) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractorKosmos.kt
similarity index 61%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractorKosmos.kt
index 3190171..70bf9bb 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/SizedTilesResetInteractorKosmos.kt
@@ -14,15 +14,13 @@
  * limitations under the License.
  */
 
-package com.android.systemui.keyboard.shortcut
+package com.android.systemui.qs.panels.domain.interactor
 
-import android.content.Intent
+import com.android.internal.logging.uiEventLogger
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
 
-class FakeShortcutHelperStartActivity : (Intent) -> Unit {
-
-    val startIntents = mutableListOf<Intent>()
-
-    override fun invoke(intent: Intent) {
-        startIntents += intent
+val Kosmos.sizedTilesResetInteractor by
+    Kosmos.Fixture {
+        SizedTilesResetInteractor(currentTilesInteractor, iconTilesInteractor, uiEventLogger)
     }
-}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/dialog/QSResetDialogDelegateKosmos.kt
similarity index 61%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/dialog/QSResetDialogDelegateKosmos.kt
index 3190171..c58d55e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/FakeShortcutHelperStartActivity.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/dialog/QSResetDialogDelegateKosmos.kt
@@ -14,15 +14,11 @@
  * limitations under the License.
  */
 
-package com.android.systemui.keyboard.shortcut
+package com.android.systemui.qs.panels.ui.dialog
 
-import android.content.Intent
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.panels.domain.interactor.sizedTilesResetInteractor
+import com.android.systemui.statusbar.phone.systemUIDialogFactory
 
-class FakeShortcutHelperStartActivity : (Intent) -> Unit {
-
-    val startIntents = mutableListOf<Intent>()
-
-    override fun invoke(intent: Intent) {
-        startIntents += intent
-    }
-}
+val Kosmos.qsResetDialogDelegateKosmos by
+    Kosmos.Fixture { QSResetDialogDelegate(systemUIDialogFactory, sizedTilesResetInteractor) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/DynamicIconTilesViewModelKosmosFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/DynamicIconTilesViewModelKosmosFactory.kt
new file mode 100644
index 0000000..d185287
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/DynamicIconTilesViewModelKosmosFactory.kt
@@ -0,0 +1,32 @@
+/*
+ * 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.viewmodel
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.panels.domain.interactor.dynamicIconTilesInteractorFactory
+
+val Kosmos.dynamicIconTilesViewModelFactory by
+    Kosmos.Fixture {
+        object : DynamicIconTilesViewModel.Factory {
+            override fun create(): DynamicIconTilesViewModel {
+                return DynamicIconTilesViewModel(
+                    dynamicIconTilesInteractorFactory,
+                    iconTilesViewModel,
+                )
+            }
+        }
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/InfiniteGridViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/InfiniteGridViewModelKosmos.kt
new file mode 100644
index 0000000..7613ea31
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/InfiniteGridViewModelKosmos.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.qs.panels.ui.viewmodel
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.panels.ui.dialog.qsResetDialogDelegateKosmos
+
+val Kosmos.infiniteGridViewModelFactory by
+    Kosmos.Fixture {
+        object : InfiniteGridViewModel.Factory {
+            override fun create(): InfiniteGridViewModel {
+                return InfiniteGridViewModel(
+                    dynamicIconTilesViewModelFactory,
+                    qsColumnsViewModel,
+                    tileSquishinessViewModel,
+                    qsResetDialogDelegateKosmos,
+                )
+            }
+        }
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeTileSpecRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeTileSpecRepository.kt
index a9cce69..1c69eab 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeTileSpecRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeTileSpecRepository.kt
@@ -63,7 +63,7 @@
 
     override suspend fun reconcileRestore(
         restoreData: RestoreData,
-        currentAutoAdded: Set<TileSpec>
+        currentAutoAdded: Set<TileSpec>,
     ) {
         with(getFlow(restoreData.userId)) {
             value = UserTileSpecRepository.reconcileTiles(value, currentAutoAdded, restoreData)
@@ -73,4 +73,8 @@
     override suspend fun prependDefault(userId: Int) {
         with(getFlow(userId)) { value = defaultTilesRepository.defaultTiles + value }
     }
+
+    override suspend fun resetToDefault(userId: Int) {
+        with(getFlow(userId)) { value = defaultTilesRepository.defaultTiles }
+    }
 }
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
index e2d73d1..9a145cb 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
@@ -137,9 +137,6 @@
 
     private static RavenwoodConfig sConfig;
     private static RavenwoodSystemProperties sProps;
-    // TODO: use the real UiAutomation class instead of a mock
-    private static UiAutomation sMockUiAutomation;
-    private static Set<String> sAdoptedPermissions = Collections.emptySet();
     private static boolean sInitialized = false;
 
     /**
@@ -187,7 +184,6 @@
                 "androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner");
 
         assertMockitoVersion();
-        sMockUiAutomation = createMockUiAutomation();
     }
 
     /**
@@ -273,7 +269,7 @@
 
         // Prepare other fields.
         config.mInstrumentation = new Instrumentation();
-        config.mInstrumentation.basicInit(instContext, targetContext, sMockUiAutomation);
+        config.mInstrumentation.basicInit(instContext, targetContext, createMockUiAutomation());
         InstrumentationRegistry.registerInstance(config.mInstrumentation, Bundle.EMPTY);
 
         RavenwoodSystemServer.init(config);
@@ -318,7 +314,6 @@
             ((RavenwoodContext) config.mTargetContext).cleanUp();
             config.mTargetContext = null;
         }
-        sMockUiAutomation.dropShellPermissionIdentity();
 
         Looper.getMainLooper().quit();
         Looper.clearMainLooperForTest();
@@ -421,28 +416,30 @@
                 () -> Class.forName("org.mockito.Matchers"));
     }
 
+    // TODO: use the real UiAutomation class instead of a mock
     private static UiAutomation createMockUiAutomation() {
+        final Set[] adoptedPermission = { Collections.emptySet() };
         var mock = mock(UiAutomation.class, inv -> {
             HostTestUtils.onThrowMethodCalled();
             return null;
         });
         doAnswer(inv -> {
-            sAdoptedPermissions = UiAutomation.ALL_PERMISSIONS;
+            adoptedPermission[0] = UiAutomation.ALL_PERMISSIONS;
             return null;
         }).when(mock).adoptShellPermissionIdentity();
         doAnswer(inv -> {
             if (inv.getArgument(0) == null) {
-                sAdoptedPermissions = UiAutomation.ALL_PERMISSIONS;
+                adoptedPermission[0] = UiAutomation.ALL_PERMISSIONS;
             } else {
-                sAdoptedPermissions = (Set) Set.of(inv.getArguments());
+                adoptedPermission[0] = Set.of(inv.getArguments());
             }
             return null;
         }).when(mock).adoptShellPermissionIdentity(any());
         doAnswer(inv -> {
-            sAdoptedPermissions = Collections.emptySet();
+            adoptedPermission[0] = Collections.emptySet();
             return null;
         }).when(mock).dropShellPermissionIdentity();
-        doAnswer(inv -> sAdoptedPermissions).when(mock).getAdoptedShellPermissions();
+        doAnswer(inv -> adoptedPermission[0]).when(mock).getAdoptedShellPermissions();
         return mock;
     }
 
diff --git a/services/accessibility/accessibility.aconfig b/services/accessibility/accessibility.aconfig
index 7057cc3..cb4e994 100644
--- a/services/accessibility/accessibility.aconfig
+++ b/services/accessibility/accessibility.aconfig
@@ -192,6 +192,16 @@
 }
 
 flag {
+    name: "package_monitor_dedicated_thread"
+    namespace: "accessibility"
+    description: "Runs the A11yManagerService PackageMonitor on a dedicated thread"
+    bug: "348138695"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "manager_package_monitor_logic_fix"
     namespace: "accessibility"
     description: "Corrects the return values of the HandleForceStop function"
diff --git a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
index 73b7b35..3441d94 100644
--- a/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
+++ b/services/accessibility/java/com/android/server/accessibility/AbstractAccessibilityServiceConnection.java
@@ -35,7 +35,6 @@
 import static android.view.accessibility.AccessibilityNodeInfo.ACTION_LONG_CLICK;
 
 import static com.android.server.pm.UserManagerService.enforceCurrentUserIfVisibleBackgroundEnabled;
-import static com.android.window.flags.Flags.deleteCaptureDisplay;
 
 import android.accessibilityservice.AccessibilityGestureEvent;
 import android.accessibilityservice.AccessibilityService;
@@ -62,7 +61,6 @@
 import android.graphics.Region;
 import android.hardware.HardwareBuffer;
 import android.hardware.display.DisplayManager;
-import android.hardware.display.DisplayManagerInternal;
 import android.hardware.usb.UsbDevice;
 import android.os.Binder;
 import android.os.Build;
@@ -104,7 +102,6 @@
 import com.android.internal.os.SomeArgs;
 import com.android.internal.util.DumpUtils;
 import com.android.internal.util.function.pooled.PooledLambda;
-import com.android.server.LocalServices;
 import com.android.server.accessibility.AccessibilityWindowManager.RemoteAccessibilityConnection;
 import com.android.server.accessibility.magnification.MagnificationProcessor;
 import com.android.server.wm.WindowManagerInternal;
@@ -1513,68 +1510,31 @@
             return;
         }
         final long identity = Binder.clearCallingIdentity();
-        if (deleteCaptureDisplay()) {
-            try {
-                ScreenCapture.ScreenCaptureListener screenCaptureListener = new
-                        ScreenCapture.ScreenCaptureListener(
-                        (screenshotBuffer, result) -> {
-                            if (screenshotBuffer != null && result == 0) {
-                                sendScreenshotSuccess(screenshotBuffer, callback);
-                            } else {
-                                sendScreenshotFailure(
-                                        AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY,
-                                        callback);
-                            }
+        try {
+            ScreenCapture.ScreenCaptureListener screenCaptureListener = new
+                    ScreenCapture.ScreenCaptureListener(
+                    (screenshotBuffer, result) -> {
+                        if (screenshotBuffer != null && result == 0) {
+                            sendScreenshotSuccess(screenshotBuffer, callback);
+                        } else {
+                            sendScreenshotFailure(
+                                    AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY,
+                                    callback);
                         }
-                );
-                mWindowManagerService.captureDisplay(displayId, null, screenCaptureListener);
-            } catch (Exception e) {
-                sendScreenshotFailure(AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY,
-                        callback);
-            } finally {
-                Binder.restoreCallingIdentity(identity);
-            }
-        } else {
-            try {
-                mMainHandler.post(PooledLambda.obtainRunnable((nonArg) -> {
-                    final ScreenshotHardwareBuffer screenshotBuffer = LocalServices
-                            .getService(DisplayManagerInternal.class).userScreenshot(displayId);
-                    if (screenshotBuffer != null) {
-                        sendScreenshotSuccess(screenshotBuffer, callback);
-                    } else {
-                        sendScreenshotFailure(
-                                AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY,
-                                callback);
                     }
-                }, null).recycleOnUse());
-            } finally {
-                Binder.restoreCallingIdentity(identity);
-            }
+            );
+            mWindowManagerService.captureDisplay(displayId, null, screenCaptureListener);
+        } catch (Exception e) {
+            sendScreenshotFailure(AccessibilityService.ERROR_TAKE_SCREENSHOT_INVALID_DISPLAY,
+                    callback);
+        } finally {
+            Binder.restoreCallingIdentity(identity);
         }
     }
 
     private void sendScreenshotSuccess(ScreenshotHardwareBuffer screenshotBuffer,
             RemoteCallback callback) {
-        if (deleteCaptureDisplay()) {
-            mMainHandler.post(PooledLambda.obtainRunnable((nonArg) -> {
-                final HardwareBuffer hardwareBuffer = screenshotBuffer.getHardwareBuffer();
-                final ParcelableColorSpace colorSpace =
-                        new ParcelableColorSpace(screenshotBuffer.getColorSpace());
-
-                final Bundle payload = new Bundle();
-                payload.putInt(KEY_ACCESSIBILITY_SCREENSHOT_STATUS,
-                        AccessibilityService.TAKE_SCREENSHOT_SUCCESS);
-                payload.putParcelable(KEY_ACCESSIBILITY_SCREENSHOT_HARDWAREBUFFER,
-                        hardwareBuffer);
-                payload.putParcelable(KEY_ACCESSIBILITY_SCREENSHOT_COLORSPACE, colorSpace);
-                payload.putLong(KEY_ACCESSIBILITY_SCREENSHOT_TIMESTAMP,
-                        SystemClock.uptimeMillis());
-
-                // Send back the result.
-                callback.sendResult(payload);
-                hardwareBuffer.close();
-            }, null).recycleOnUse());
-        } else {
+        mMainHandler.post(PooledLambda.obtainRunnable((nonArg) -> {
             final HardwareBuffer hardwareBuffer = screenshotBuffer.getHardwareBuffer();
             final ParcelableColorSpace colorSpace =
                     new ParcelableColorSpace(screenshotBuffer.getColorSpace());
@@ -1591,7 +1551,7 @@
             // Send back the result.
             callback.sendResult(payload);
             hardwareBuffer.close();
-        }
+        }, null).recycleOnUse());
     }
 
     private void sendScreenshotFailure(@AccessibilityService.ScreenshotErrorCode int errorCode,
diff --git a/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java b/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java
index ec8908b..c6fe497 100644
--- a/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java
+++ b/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java
@@ -116,6 +116,7 @@
 import android.os.Build;
 import android.os.Bundle;
 import android.os.Handler;
+import android.os.HandlerThread;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
@@ -901,7 +902,19 @@
     private void registerBroadcastReceivers() {
         // package changes
         mPackageMonitor = new ManagerPackageMonitor(this);
-        mPackageMonitor.register(mContext, null,  UserHandle.ALL, true);
+        final Looper packageMonitorLooper;
+        if (Flags.packageMonitorDedicatedThread()) {
+            // Use a dedicated thread because the default BackgroundThread used by PackageMonitor
+            // is shared by other components and can get busy, causing a delay and eventual ANR when
+            // responding to broadcasts sent to this PackageMonitor.
+            HandlerThread packageMonitorThread = new HandlerThread(LOG_TAG + " PackageMonitor",
+                    Process.THREAD_PRIORITY_BACKGROUND);
+            packageMonitorThread.start();
+            packageMonitorLooper = packageMonitorThread.getLooper();
+        } else {
+            packageMonitorLooper = null;
+        }
+        mPackageMonitor.register(mContext, packageMonitorLooper,  UserHandle.ALL, true);
 
         // user change and unlock
         IntentFilter intentFilter = new IntentFilter();
diff --git a/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java b/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
index 74908a4..3608360 100644
--- a/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
+++ b/services/companion/java/com/android/server/companion/transport/CompanionTransportManager.java
@@ -57,8 +57,11 @@
     /** Association id -> Transport */
     @GuardedBy("mTransports")
     private final SparseArray<Transport> mTransports = new SparseArray<>();
+
+    // Use mTransports to synchronize both mTransports and mTransportsListeners to avoid deadlock
+    // between threads that access both
     @NonNull
-    @GuardedBy("mTransportsListeners")
+    @GuardedBy("mTransports")
     private final RemoteCallbackList<IOnTransportsChangedListener> mTransportsListeners =
             new RemoteCallbackList<>();
 
@@ -95,7 +98,7 @@
      */
     public void addListener(IOnTransportsChangedListener listener) {
         Slog.i(TAG, "Registering OnTransportsChangedListener");
-        synchronized (mTransportsListeners) {
+        synchronized (mTransports) {
             mTransportsListeners.register(listener);
             mTransportsListeners.broadcast(listener1 -> {
                 // callback to the current listener with all the associations of the transports
@@ -114,7 +117,7 @@
      * Remove the listener for receiving callbacks when any of the transports is changed
      */
     public void removeListener(IOnTransportsChangedListener listener) {
-        synchronized (mTransportsListeners) {
+        synchronized (mTransports) {
             mTransportsListeners.unregister(listener);
         }
     }
@@ -204,7 +207,7 @@
     }
 
     private void notifyOnTransportsChanged() {
-        synchronized (mTransportsListeners) {
+        synchronized (mTransports) {
             mTransportsListeners.broadcast(listener -> {
                 try {
                     listener.onTransportsChanged(getAssociationsWithTransport());
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index 6405ebb..217ef20 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -2757,8 +2757,11 @@
         if (isolated) {
             if (mIsolatedAppBindArgs == null) {
                 mIsolatedAppBindArgs = new ArrayMap<>(1);
+                // See b/79378449 about the following exemption.
                 addServiceToMap(mIsolatedAppBindArgs, "package");
-                addServiceToMap(mIsolatedAppBindArgs, "permissionmgr");
+                if (!android.server.Flags.removeJavaServiceManagerCache()) {
+                    addServiceToMap(mIsolatedAppBindArgs, "permissionmgr");
+                }
             }
             return mIsolatedAppBindArgs;
         }
@@ -2769,27 +2772,33 @@
             // Add common services.
             // IMPORTANT: Before adding services here, make sure ephemeral apps can access them too.
             // Enable the check in ApplicationThread.bindApplication() to make sure.
+            if (!android.server.Flags.removeJavaServiceManagerCache()) {
+                addServiceToMap(mAppBindArgs, "permissionmgr");
+                addServiceToMap(mAppBindArgs, Context.ALARM_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.DISPLAY_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.NETWORKMANAGEMENT_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.CONNECTIVITY_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.ACCESSIBILITY_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.INPUT_METHOD_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.INPUT_SERVICE);
+                addServiceToMap(mAppBindArgs, "graphicsstats");
+                addServiceToMap(mAppBindArgs, Context.APP_OPS_SERVICE);
+                addServiceToMap(mAppBindArgs, "content");
+                addServiceToMap(mAppBindArgs, Context.JOB_SCHEDULER_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.NOTIFICATION_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.VIBRATOR_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.ACCOUNT_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.POWER_SERVICE);
+                addServiceToMap(mAppBindArgs, Context.USER_SERVICE);
+                addServiceToMap(mAppBindArgs, "mount");
+                addServiceToMap(mAppBindArgs, Context.PLATFORM_COMPAT_SERVICE);
+            }
+            // See b/79378449
+            // Getting the window service and package service binder from servicemanager
+            // is blocked for Apps. However they are necessary for apps.
+            // TODO: remove exception
             addServiceToMap(mAppBindArgs, "package");
-            addServiceToMap(mAppBindArgs, "permissionmgr");
             addServiceToMap(mAppBindArgs, Context.WINDOW_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.ALARM_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.DISPLAY_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.NETWORKMANAGEMENT_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.CONNECTIVITY_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.ACCESSIBILITY_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.INPUT_METHOD_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.INPUT_SERVICE);
-            addServiceToMap(mAppBindArgs, "graphicsstats");
-            addServiceToMap(mAppBindArgs, Context.APP_OPS_SERVICE);
-            addServiceToMap(mAppBindArgs, "content");
-            addServiceToMap(mAppBindArgs, Context.JOB_SCHEDULER_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.NOTIFICATION_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.VIBRATOR_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.ACCOUNT_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.POWER_SERVICE);
-            addServiceToMap(mAppBindArgs, Context.USER_SERVICE);
-            addServiceToMap(mAppBindArgs, "mount");
-            addServiceToMap(mAppBindArgs, Context.PLATFORM_COMPAT_SERVICE);
         }
         return mAppBindArgs;
     }
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 592d89e..c47cad9 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -196,6 +196,7 @@
     private final PowerAttributor mPowerAttributor;
 
     private volatile boolean mMonitorEnabled = true;
+    private boolean mRailsStatsCollectionEnabled = true;
 
     private native void getRailEnergyPowerStats(RailStats railStats);
     private CharsetDecoder mDecoderStat = StandardCharsets.UTF_8
@@ -312,8 +313,17 @@
         }
     }
 
+    public void setRailsStatsCollectionEnabled(boolean railsStatsCollectionEnabled) {
+        mRailsStatsCollectionEnabled = railsStatsCollectionEnabled;
+    }
+
     @Override
     public void fillRailDataStats(RailStats railStats) {
+        if (!mRailsStatsCollectionEnabled) {
+            railStats.setRailStatsAvailability(false);
+            return;
+        }
+
         if (DBG) Slog.d(TAG, "begin getRailEnergyPowerStats");
         try {
             getRailEnergyPowerStats(railStats);
@@ -423,7 +433,7 @@
         mStats = new BatteryStatsImpl(mBatteryStatsConfig, Clock.SYSTEM_CLOCK, mMonotonicClock,
                 systemDir, mHandler, this, this, mUserManagerUserInfoProvider, mPowerProfile,
                 mCpuScalingPolicies, mPowerStatsUidResolver);
-        mWorker = new BatteryExternalStatsWorker(context, mStats);
+        mWorker = new BatteryExternalStatsWorker(context, mStats, mHandler);
         mStats.setExternalStatsSyncLocked(mWorker);
         mStats.setRadioScanningTimeoutLocked(mContext.getResources().getInteger(
                 com.android.internal.R.integer.config_radioScanningTimeout) * 1000L);
@@ -436,9 +446,12 @@
                 mCpuScalingPolicies, () -> mStats.getBatteryCapacity(),
                 mPowerStatsUidResolver);
         mPowerStatsScheduler = createPowerStatsScheduler(mContext);
+
+        int accumulatedBatteryUsageStatsSpanSize = mContext.getResources().getInteger(
+                com.android.internal.R.integer.config_accumulatedBatteryUsageStatsSpanSize);
         mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context,
                 mPowerAttributor, mPowerProfile, mCpuScalingPolicies,
-                mPowerStatsStore, Clock.SYSTEM_CLOCK);
+                mPowerStatsStore, accumulatedBatteryUsageStatsSpanSize, Clock.SYSTEM_CLOCK);
         mDumpHelper = new BatteryStatsDumpHelperImpl(mBatteryUsageStatsProvider);
         mCpuWakeupStats = new CpuWakeupStats(context, R.xml.irq_device_map, mHandler);
         mConfigFile = new AtomicFile(new File(systemDir, "battery_usage_stats_config"));
@@ -506,7 +519,7 @@
 
     public void systemServicesReady() {
         mStats.saveBatteryUsageStatsOnReset(mBatteryUsageStatsProvider, mPowerStatsStore,
-                Flags.accumulateBatteryUsageStats());
+                isBatteryUsageStatsAccumulationSupported());
 
         MultiStatePowerAttributor attributor = (MultiStatePowerAttributor) mPowerAttributor;
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CPU,
@@ -588,6 +601,12 @@
                 BatteryConsumer.POWER_COMPONENT_CAMERA,
                 Flags.streamlinedMiscBatteryStats());
 
+        // Currently unimplemented.
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MEMORY,
+                Flags.streamlinedMiscBatteryStats());
+        attributor.setPowerComponentSupported(BatteryConsumer.POWER_COMPONENT_MEMORY,
+                Flags.streamlinedMiscBatteryStats());
+
         // By convention POWER_COMPONENT_ANY represents custom Energy Consumers
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_ANY,
                 Flags.streamlinedMiscBatteryStats());
@@ -631,6 +650,13 @@
         registerStatsCallbacks();
     }
 
+    private static boolean isBatteryUsageStatsAccumulationSupported() {
+        return Flags.accumulateBatteryUsageStats()
+                && Flags.streamlinedBatteryStats()
+                && Flags.streamlinedConnectivityBatteryStats()
+                && Flags.streamlinedMiscBatteryStats();
+    }
+
     /**
      * Notifies BatteryStatsService that the system server is ready.
      */
@@ -776,7 +802,8 @@
 
     private void syncStats(String reason, int flags) {
         mStats.collectPowerStatsSamples();
-        awaitUninterruptibly(mWorker.scheduleSync(reason, flags));
+        mWorker.scheduleSync(reason, flags);
+        awaitCompletion();
     }
 
     private void awaitCompletion() {
@@ -1135,7 +1162,7 @@
                             .includeVirtualUids()
                             .setMinConsumedPowerThreshold(minConsumedPowerThreshold);
 
-                    if (Flags.accumulateBatteryUsageStats()) {
+                    if (isBatteryUsageStatsAccumulationSupported()) {
                         query.accumulated();
                     }
 
@@ -3054,7 +3081,7 @@
         if (Flags.streamlinedBatteryStats()) {
             pw.println("  --sample: collect and dump a sample of stats for debugging purpose");
         }
-        if (Flags.accumulateBatteryUsageStats()) {
+        if (isBatteryUsageStatsAccumulationSupported()) {
             pw.println("  --accumulated: continuously accumulated since setup or reset-all");
         }
         pw.println("  <package.name>: optional name of package to filter output by.");
@@ -3670,24 +3697,12 @@
                     android.Manifest.permission.BATTERY_STATS, null);
         }
 
-        Future future;
         if (shouldCollectExternalStats()) {
-            future = mWorker.scheduleSync("get-health-stats-for-uids",
+            mWorker.scheduleSync("get-health-stats-for-uids",
                     BatteryExternalStatsWorker.UPDATE_ALL);
-        } else {
-            future = null;
         }
 
         mHandler.post(() -> {
-            if (future != null) {
-                try {
-                    // Worker uses a separate thread pool, so waiting here won't cause a deadlock
-                    future.get();
-                } catch (InterruptedException | ExecutionException e) {
-                    Slog.e(TAG, "Sync failed", e);
-                }
-            }
-
             final long ident = Binder.clearCallingIdentity();
             int i = -1;
             try {
diff --git a/services/core/java/com/android/server/audio/AudioService.java b/services/core/java/com/android/server/audio/AudioService.java
index 37a2fba..78a1fa7 100644
--- a/services/core/java/com/android/server/audio/AudioService.java
+++ b/services/core/java/com/android/server/audio/AudioService.java
@@ -3567,8 +3567,10 @@
      * @see AudioManager#addOnDevicesForAttributesChangedListener(
      *      AudioAttributes, Executor, OnDevicesForAttributesChangedListener)
      */
+    @android.annotation.EnforcePermission(anyOf = { MODIFY_AUDIO_ROUTING, QUERY_AUDIO_STATE })
     public void addOnDevicesForAttributesChangedListener(AudioAttributes attributes,
             IDevicesForAttributesCallback callback) {
+        super.addOnDevicesForAttributesChangedListener_enforcePermission();
         mAudioSystem.addOnDevicesForAttributesChangedListener(
                 attributes, false /* forVolume */, callback);
     }
diff --git a/services/core/java/com/android/server/content/OWNERS b/services/core/java/com/android/server/content/OWNERS
index b6a9fe86..5642382 100644
--- a/services/core/java/com/android/server/content/OWNERS
+++ b/services/core/java/com/android/server/content/OWNERS
@@ -1 +1,3 @@
-include /services/core/java/com/android/server/am/OWNERS
\ No newline at end of file
+include /services/core/java/com/android/server/am/OWNERS
+
+per-file Sync* = file:/apex/jobscheduler/JOB_OWNERS
\ No newline at end of file
diff --git a/services/core/java/com/android/server/display/color/ColorDisplayService.java b/services/core/java/com/android/server/display/color/ColorDisplayService.java
index dc59e66..7892639 100644
--- a/services/core/java/com/android/server/display/color/ColorDisplayService.java
+++ b/services/core/java/com/android/server/display/color/ColorDisplayService.java
@@ -78,6 +78,7 @@
 
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.annotations.WeaklyReferencedCallback;
 import com.android.internal.util.DumpUtils;
 import com.android.server.DisplayThread;
 import com.android.server.LocalServices;
@@ -1795,6 +1796,7 @@
     /**
      * Interface for applying transforms to a given AppWindow.
      */
+    @WeaklyReferencedCallback
     public interface ColorTransformController {
 
         /**
diff --git a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
index 07343f4..c0aa4cc 100644
--- a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
+++ b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
@@ -237,6 +237,11 @@
             Flags::enableHasArrSupport
     );
 
+    private final FlagState mAutoBrightnessModeBedtimeWearFlagState = new FlagState(
+            Flags.FLAG_AUTO_BRIGHTNESS_MODE_BEDTIME_WEAR,
+            Flags::autoBrightnessModeBedtimeWear
+    );
+
     /**
      * @return {@code true} if 'port' is allowed in display layout configuration file.
      */
@@ -503,6 +508,15 @@
     public boolean hasArrSupportFlag() {
         return mHasArrSupport.isEnabled();
     }
+
+    /**
+     * @return {@code true} if bedtime mode specific auto-brightness curve should be loaded and be
+     * applied when bedtime mode is enabled.
+     */
+    public boolean isAutoBrightnessModeBedtimeWearEnabled() {
+        return mAutoBrightnessModeBedtimeWearFlagState.isEnabled();
+    }
+
     /**
      * dumps all flagstates
      * @param pw printWriter
@@ -553,6 +567,7 @@
         pw.println(" " + mBlockAutobrightnessChangesOnStylusUsage);
         pw.println(" " + mIsUserRefreshRateForExternalDisplayEnabled);
         pw.println(" " + mHasArrSupport);
+        pw.println(" " + mAutoBrightnessModeBedtimeWearFlagState);
     }
 
     private static class FlagState {
diff --git a/services/core/java/com/android/server/display/feature/display_flags.aconfig b/services/core/java/com/android/server/display/feature/display_flags.aconfig
index ddb2969..36cadf5 100644
--- a/services/core/java/com/android/server/display/feature/display_flags.aconfig
+++ b/services/core/java/com/android/server/display/feature/display_flags.aconfig
@@ -422,3 +422,11 @@
     bug: "361433651"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "auto_brightness_mode_bedtime_wear"
+    namespace: "wear_frameworks"
+    description: "Feature flag for loading and applying auto-brightness curve while wear bedtime mode enabled."
+    bug: "350617205"
+    is_fixed_read_only: true
+}
diff --git a/services/core/java/com/android/server/net/OWNERS b/services/core/java/com/android/server/net/OWNERS
index bbc7c01..4596a44 100644
--- a/services/core/java/com/android/server/net/OWNERS
+++ b/services/core/java/com/android/server/net/OWNERS
@@ -2,7 +2,5 @@
 file:platform/packages/modules/Connectivity:main:/OWNERS_core_networking
 per-file NetworkPolicyManagerService.java=jackyu@google.com, sarahchin@google.com
 
-jsharkey@android.com
 sudheersai@google.com
-yamasani@google.com
 suprabh@google.com
diff --git a/services/core/java/com/android/server/notification/ManagedServices.java b/services/core/java/com/android/server/notification/ManagedServices.java
index 93482e7..122836e 100644
--- a/services/core/java/com/android/server/notification/ManagedServices.java
+++ b/services/core/java/com/android/server/notification/ManagedServices.java
@@ -21,6 +21,9 @@
 import static android.content.Context.BIND_AUTO_CREATE;
 import static android.content.Context.BIND_FOREGROUND_SERVICE;
 import static android.content.Context.DEVICE_POLICY_SERVICE;
+import static android.content.pm.PackageManager.MATCH_DIRECT_BOOT_AWARE;
+import static android.content.pm.PackageManager.MATCH_DIRECT_BOOT_UNAWARE;
+import static android.content.pm.PackageManager.MATCH_INSTANT;
 import static android.os.UserHandle.USER_ALL;
 import static android.os.UserHandle.USER_SYSTEM;
 import static android.service.notification.NotificationListenerService.META_DATA_DEFAULT_AUTOBIND;
@@ -106,7 +109,8 @@
     protected final String TAG = getClass().getSimpleName().replace('$', '.');
     protected final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
 
-    private static final int ON_BINDING_DIED_REBIND_DELAY_MS = 10000;
+    protected static final int ON_BINDING_DIED_REBIND_DELAY_MS = 10000;
+    protected static final int ON_BINDING_DIED_REBIND_MSG = 1234;
     protected static final String ENABLED_SERVICES_SEPARATOR = ":";
     private static final String DB_VERSION_1 = "1";
     private static final String DB_VERSION_2 = "2";
@@ -875,7 +879,21 @@
             String approvedItem = getApprovedValue(pkgOrComponent);
 
             if (approvedItem != null) {
+                final ComponentName component = ComponentName.unflattenFromString(approvedItem);
                 if (enabled) {
+                    if (Flags.notificationNlsRebind()) {
+                        if (component != null && !isValidService(component, userId)) {
+                            // Only fail if package is available
+                            // If not, it will be validated again in onPackagesChanged
+                            final PackageManager pm = mContext.getPackageManager();
+                            if (pm.isPackageAvailable(component.getPackageName())) {
+                                Slog.w(TAG, "Skip allowing " + mConfig.caption
+                                        + " " + pkgOrComponent + " (userSet: " + userSet
+                                        + ") for invalid service");
+                                return;
+                            }
+                        }
+                    }
                     approved.add(approvedItem);
                 } else {
                     approved.remove(approvedItem);
@@ -973,7 +991,7 @@
                 || isPackageOrComponentAllowed(component.getPackageName(), userId))) {
             return false;
         }
-        return componentHasBindPermission(component, userId);
+        return isValidService(component, userId);
     }
 
     private boolean componentHasBindPermission(ComponentName component, int userId) {
@@ -1220,12 +1238,21 @@
         if (!TextUtils.isEmpty(packageName)) {
             queryIntent.setPackage(packageName);
         }
+
+        if (Flags.notificationNlsRebind()) {
+            // Expand the package query
+            extraFlags |= MATCH_DIRECT_BOOT_AWARE | MATCH_DIRECT_BOOT_UNAWARE;
+            extraFlags |= MATCH_INSTANT;
+        }
+
         List<ResolveInfo> installedServices = pm.queryIntentServicesAsUser(
                 queryIntent,
                 PackageManager.GET_SERVICES | PackageManager.GET_META_DATA | extraFlags,
                 userId);
-        if (DEBUG)
-            Slog.v(TAG, mConfig.serviceInterface + " services: " + installedServices);
+        if (DEBUG) {
+            Slog.v(TAG, mConfig.serviceInterface + " pkg: " + packageName + " services: "
+                    + installedServices);
+        }
         if (installedServices != null) {
             for (int i = 0, count = installedServices.size(); i < count; i++) {
                 ResolveInfo resolveInfo = installedServices.get(i);
@@ -1325,11 +1352,12 @@
                     if (TextUtils.equals(getPackageName(approvedPackageOrComponent), packageName)) {
                         final ComponentName component = ComponentName.unflattenFromString(
                                 approvedPackageOrComponent);
-                        if (component != null && !componentHasBindPermission(component, userId)) {
+                        if (component != null && !isValidService(component, userId)) {
                             approved.removeAt(j);
                             if (DEBUG) {
                                 Slog.v(TAG, "Removing " + approvedPackageOrComponent
-                                        + " from approved list; no bind permission found "
+                                        + " from approved list; no bind permission or "
+                                        + "service interface filter found "
                                         + mConfig.bindPermission);
                             }
                         }
@@ -1348,6 +1376,15 @@
         }
     }
 
+    protected boolean isValidService(ComponentName component, int userId) {
+        if (Flags.notificationNlsRebind()) {
+            return componentHasBindPermission(component, userId) && queryPackageForServices(
+                    component.getPackageName(), userId).contains(component);
+        } else {
+            return componentHasBindPermission(component, userId);
+        }
+    }
+
     protected boolean isValidEntry(String packageOrComponent, int userId) {
         return hasMatchingServices(packageOrComponent, userId);
     }
@@ -1505,23 +1542,27 @@
      * Called when user switched to unbind all services from other users.
      */
     @VisibleForTesting
-    void unbindOtherUserServices(int currentUser) {
+    void unbindOtherUserServices(int switchedToUser) {
         TimingsTraceAndSlog t = new TimingsTraceAndSlog();
-        t.traceBegin("ManagedServices.unbindOtherUserServices_current" + currentUser);
-        unbindServicesImpl(currentUser, true /* allExceptUser */);
+        t.traceBegin("ManagedServices.unbindOtherUserServices_current" + switchedToUser);
+        unbindServicesImpl(switchedToUser, true /* allExceptUser */);
         t.traceEnd();
     }
 
-    void unbindUserServices(int user) {
+    void unbindUserServices(int removedUser) {
         TimingsTraceAndSlog t = new TimingsTraceAndSlog();
-        t.traceBegin("ManagedServices.unbindUserServices" + user);
-        unbindServicesImpl(user, false /* allExceptUser */);
+        t.traceBegin("ManagedServices.unbindUserServices" + removedUser);
+        unbindServicesImpl(removedUser, false /* allExceptUser */);
         t.traceEnd();
     }
 
     void unbindServicesImpl(int user, boolean allExceptUser) {
         final SparseArray<Set<ComponentName>> componentsToUnbind = new SparseArray<>();
         synchronized (mMutex) {
+            if (Flags.notificationNlsRebind()) {
+                // Remove enqueued rebinds to avoid rebinding services for a switched user
+                mHandler.removeMessages(ON_BINDING_DIED_REBIND_MSG);
+            }
             final Set<ManagedServiceInfo> removableBoundServices = getRemovableConnectedServices();
             for (ManagedServiceInfo info : removableBoundServices) {
                 if ((allExceptUser && (info.userid != user))
@@ -1716,6 +1757,7 @@
                             mServicesRebinding.add(servicesBindingTag);
                             mHandler.postDelayed(() ->
                                     reregisterService(name, userid),
+                                    ON_BINDING_DIED_REBIND_MSG,
                                     ON_BINDING_DIED_REBIND_DELAY_MS);
                         } else {
                             Slog.v(TAG, getCaption() + " not rebinding in user " + userid
diff --git a/services/core/java/com/android/server/notification/flags.aconfig b/services/core/java/com/android/server/notification/flags.aconfig
index f79d9ef..c479acf 100644
--- a/services/core/java/com/android/server/notification/flags.aconfig
+++ b/services/core/java/com/android/server/notification/flags.aconfig
@@ -194,3 +194,13 @@
   description: "Enables sound uri with vibration source in notification channel"
   bug: "351975435"
 }
+
+flag {
+  name: "notification_nls_rebind"
+  namespace: "systemui"
+  description: "Check for NLS service intent filter when rebinding services"
+  bug: "347674739"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/services/core/java/com/android/server/pm/UserManagerService.java b/services/core/java/com/android/server/pm/UserManagerService.java
index daf413b..6c03214 100644
--- a/services/core/java/com/android/server/pm/UserManagerService.java
+++ b/services/core/java/com/android/server/pm/UserManagerService.java
@@ -195,6 +195,7 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -5789,6 +5790,9 @@
             }
 
             userInfo.partial = false;
+            if (android.multiuser.Flags.invalidateCacheOnUsersChangedReadOnly()) {
+                UserManager.invalidateCacheOnUserListChange();
+            }
             synchronized (mPackagesLock) {
                 writeUserLP(userData);
             }
@@ -6261,14 +6265,16 @@
         Slog.i(LOG_TAG, "removeUser u" + userId);
         checkCreateUsersPermission("Only the system can remove users");
 
-        final String restriction = getUserRemovalRestriction(userId);
-        if (getUserRestrictions(UserHandle.getCallingUserId()).getBoolean(restriction, false)) {
-            Slog.w(LOG_TAG, "Cannot remove user. " + restriction + " is enabled.");
+        final Optional<String> restrictionOptional = getUserRemovalRestrictionOptional(userId);
+        if (!restrictionOptional.isEmpty()
+                && getUserRestrictions(UserHandle.getCallingUserId())
+                        .getBoolean(restrictionOptional.get(), false)) {
+            Slog.w(LOG_TAG, "Cannot remove user. " + restrictionOptional.get() + " is enabled.");
             return false;
         }
         if (mCurrentBootPhase < SystemService.PHASE_ACTIVITY_MANAGER_READY) {
             Slog.w(LOG_TAG, "Cannot remove user, removeUser is called too early during boot. "
-                + "ActivityManager is not ready yet.");
+                            + "ActivityManager is not ready yet.");
             return false;
         }
         return removeUserWithProfilesUnchecked(userId);
@@ -6335,18 +6341,30 @@
     }
 
     /**
-     * Returns the string name of the restriction to check for user removal. The restriction name
-     * varies depending on whether the user is a managed profile.
+     * Returns an optional string name of the restriction to check for user removal. The restriction
+     * name varies depending on whether the user is a managed profile.
+     *
+     * <p>If the flag android.multiuser.ignore_restrictions_when_deleting_private_profile is enabled
+     * and the user is a private profile (i.e. has no removal restrictions) the method will return
+     * {@code Optional.empty()}.
      */
-    private String getUserRemovalRestriction(@UserIdInt int userId) {
+    private Optional<String> getUserRemovalRestrictionOptional(@UserIdInt int userId) {
+        final boolean isPrivateProfile;
         final boolean isManagedProfile;
         final UserInfo userInfo;
         synchronized (mUsersLock) {
             userInfo = getUserInfoLU(userId);
         }
+        isPrivateProfile = userInfo != null && userInfo.isPrivateProfile();
         isManagedProfile = userInfo != null && userInfo.isManagedProfile();
-        return isManagedProfile
-                ? UserManager.DISALLOW_REMOVE_MANAGED_PROFILE : UserManager.DISALLOW_REMOVE_USER;
+        if (android.multiuser.Flags.ignoreRestrictionsWhenDeletingPrivateProfile()
+                && isPrivateProfile) {
+            return Optional.empty();
+        }
+        return Optional.of(
+                isManagedProfile
+                        ? UserManager.DISALLOW_REMOVE_MANAGED_PROFILE
+                        : UserManager.DISALLOW_REMOVE_USER);
     }
 
     private boolean removeUserUnchecked(@UserIdInt int userId) {
@@ -6367,6 +6385,9 @@
                 // on next startup, in case the runtime stops now before stopping and
                 // removing the user completely.
                 userData.info.partial = true;
+                if (android.multiuser.Flags.invalidateCacheOnUsersChangedReadOnly()) {
+                    UserManager.invalidateCacheOnUserListChange();
+                }
                 // Mark it as disabled, so that it isn't returned any more when
                 // profiles are queried.
                 userData.info.flags |= UserInfo.FLAG_DISABLED;
@@ -6455,9 +6476,13 @@
         checkCreateUsersPermission("Only the system can remove users");
 
         if (!overrideDevicePolicy) {
-            final String restriction = getUserRemovalRestriction(userId);
-            if (getUserRestrictions(UserHandle.getCallingUserId()).getBoolean(restriction, false)) {
-                Slog.w(LOG_TAG, "Cannot remove user. " + restriction + " is enabled.");
+            final Optional<String> restrictionOptional = getUserRemovalRestrictionOptional(userId);
+            if (!restrictionOptional.isEmpty()
+                    && getUserRestrictions(UserHandle.getCallingUserId())
+                            .getBoolean(restrictionOptional.get(), false)) {
+                Slog.w(
+                        LOG_TAG,
+                        "Cannot remove user. " + restrictionOptional.get() + " is enabled.");
                 return UserManager.REMOVE_RESULT_ERROR_USER_RESTRICTION;
             }
         }
diff --git a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
index 8311034..f90da64 100644
--- a/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/power/stats/BatteryExternalStatsWorker.java
@@ -27,12 +27,11 @@
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 import android.os.Bundle;
+import android.os.Handler;
 import android.os.OutcomeReceiver;
 import android.os.Parcelable;
-import android.os.Process;
 import android.os.SynchronousResultReceiver;
 import android.os.SystemClock;
-import android.os.ThreadLocalWorkSource;
 import android.os.connectivity.WifiActivityEnergyInfo;
 import android.power.PowerStatsInternal;
 import android.telephony.ModemActivityInfo;
@@ -50,11 +49,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
-import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
-import java.util.concurrent.RejectedExecutionException;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
@@ -85,29 +80,23 @@
     // stop running.
     public static final int UID_FINAL_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS = 10_000;
 
-    private final ScheduledExecutorService mExecutorService =
-            Executors.newSingleThreadScheduledExecutor(
-                    (ThreadFactory) r -> {
-                        Thread t = new Thread(
-                                () -> {
-                                    ThreadLocalWorkSource.setUid(Process.myUid());
-                                    r.run();
-                                },
-                                "batterystats-worker");
-                        t.setPriority(Thread.NORM_PRIORITY);
-                        return t;
-                    });
+    // Various types of sync, passed to Handler
+    private static final int SYNC_UPDATE = 1;
+    private static final int SYNC_WAKELOCK_CHANGE = 2;
+    private static final int SYNC_BATTERY_LEVEL_CHANGE = 3;
+    private static final int SYNC_PROCESS_STATE_CHANGE = 4;
+    private static final int SYNC_USER_REMOVAL = 5;
+
+    private final Handler mHandler;
 
     @GuardedBy("mStats")
     private final BatteryStatsImpl mStats;
 
     @GuardedBy("this")
+    @ExternalUpdateFlag
     private int mUpdateFlags = 0;
 
     @GuardedBy("this")
-    private Future<?> mCurrentFuture = null;
-
-    @GuardedBy("this")
     private String mCurrentReason = null;
 
     @GuardedBy("this")
@@ -125,15 +114,6 @@
     @GuardedBy("this")
     private boolean mUseLatestStates = true;
 
-    @GuardedBy("this")
-    private Future<?> mWakelockChangesUpdate;
-
-    @GuardedBy("this")
-    private Future<?> mBatteryLevelSync;
-
-    @GuardedBy("this")
-    private Future<?> mProcessStateSync;
-
     // If both mStats and mWorkerLock need to be synchronized, mWorkerLock must be acquired first.
     private final Object mWorkerLock = new Object();
 
@@ -190,14 +170,15 @@
         }
     }
 
-    public BatteryExternalStatsWorker(Context context, BatteryStatsImpl stats) {
-        this(new Injector(context), stats);
+    public BatteryExternalStatsWorker(Context context, BatteryStatsImpl stats, Handler handler) {
+        this(new Injector(context), stats, handler);
     }
 
     @VisibleForTesting
-    BatteryExternalStatsWorker(Injector injector, BatteryStatsImpl stats) {
+    BatteryExternalStatsWorker(Injector injector, BatteryStatsImpl stats, Handler handler) {
         mInjector = injector;
         mStats = stats;
+        mHandler = handler;
     }
 
     public void systemServicesReady() {
@@ -249,20 +230,20 @@
     }
 
     @Override
-    public synchronized Future<?> scheduleSync(String reason, int flags) {
-        return scheduleSyncLocked(reason, flags);
+    public synchronized void scheduleSync(String reason, @ExternalUpdateFlag int flags) {
+        scheduleSyncLocked(reason, flags);
     }
 
     @Override
-    public synchronized Future<?> scheduleCpuSyncDueToRemovedUid(int uid) {
-        return scheduleSyncLocked("remove-uid", UPDATE_CPU);
+    public synchronized void scheduleCpuSyncDueToRemovedUid(int uid) {
+        scheduleSyncLocked("remove-uid", UPDATE_CPU);
     }
 
     @Override
-    public Future<?> scheduleSyncDueToScreenStateChange(int flags, boolean onBattery,
+    public void scheduleSyncDueToScreenStateChange(@ExternalUpdateFlag int flags, boolean onBattery,
             boolean onBatteryScreenOff, int screenState, int[] perDisplayScreenStates) {
         synchronized (BatteryExternalStatsWorker.this) {
-            if (mCurrentFuture == null || (mUpdateFlags & UPDATE_CPU) == 0) {
+            if (!mHandler.hasMessages(SYNC_UPDATE) || (mUpdateFlags & UPDATE_CPU) == 0) {
                 mOnBattery = onBattery;
                 mOnBatteryScreenOff = onBatteryScreenOff;
                 mUseLatestStates = false;
@@ -270,91 +251,70 @@
             // always update screen state
             mScreenState = screenState;
             mPerDisplayScreenStates = perDisplayScreenStates;
-            return scheduleSyncLocked("screen-state", flags);
+            scheduleSyncLocked("screen-state", flags);
         }
     }
 
     @Override
-    public Future<?> scheduleCpuSyncDueToWakelockChange(long delayMillis) {
+    public void scheduleCpuSyncDueToWakelockChange(long delayMillis) {
         synchronized (BatteryExternalStatsWorker.this) {
-            mWakelockChangesUpdate = scheduleDelayedSyncLocked(mWakelockChangesUpdate,
+            scheduleDelayedSyncLocked(SYNC_WAKELOCK_CHANGE,
                     () -> {
                         scheduleSync("wakelock-change", UPDATE_CPU);
                         scheduleRunnable(() -> mStats.postBatteryNeedsCpuUpdateMsg());
                     },
                     delayMillis);
-            return mWakelockChangesUpdate;
         }
     }
 
     @Override
     public void cancelCpuSyncDueToWakelockChange() {
-        synchronized (BatteryExternalStatsWorker.this) {
-            if (mWakelockChangesUpdate != null) {
-                mWakelockChangesUpdate.cancel(false);
-                mWakelockChangesUpdate = null;
-            }
-        }
+        mHandler.removeMessages(SYNC_WAKELOCK_CHANGE);
     }
 
     @Override
-    public Future<?> scheduleSyncDueToBatteryLevelChange(long delayMillis) {
+    public void scheduleSyncDueToBatteryLevelChange(long delayMillis) {
         synchronized (BatteryExternalStatsWorker.this) {
-            mBatteryLevelSync = scheduleDelayedSyncLocked(mBatteryLevelSync,
+            scheduleDelayedSyncLocked(SYNC_BATTERY_LEVEL_CHANGE,
                     () -> scheduleSync("battery-level", UPDATE_ALL),
                     delayMillis);
-            return mBatteryLevelSync;
         }
     }
 
     @GuardedBy("this")
     private void cancelSyncDueToBatteryLevelChangeLocked() {
-        if (mBatteryLevelSync != null) {
-            mBatteryLevelSync.cancel(false);
-            mBatteryLevelSync = null;
-        }
+        mHandler.removeMessages(SYNC_BATTERY_LEVEL_CHANGE);
     }
 
     @Override
     public void scheduleSyncDueToProcessStateChange(int flags, long delayMillis) {
         synchronized (BatteryExternalStatsWorker.this) {
-            mProcessStateSync = scheduleDelayedSyncLocked(mProcessStateSync,
+            scheduleDelayedSyncLocked(SYNC_PROCESS_STATE_CHANGE,
                     () -> scheduleSync("procstate-change", flags),
                     delayMillis);
         }
     }
 
     public void cancelSyncDueToProcessStateChange() {
-        synchronized (BatteryExternalStatsWorker.this) {
-            if (mProcessStateSync != null) {
-                mProcessStateSync.cancel(false);
-                mProcessStateSync = null;
-            }
-        }
+        mHandler.removeMessages(SYNC_PROCESS_STATE_CHANGE);
     }
 
     @Override
-    public Future<?> scheduleCleanupDueToRemovedUser(int userId) {
-        synchronized (BatteryExternalStatsWorker.this) {
-            try {
-                // Initial quick clean-up after a user removal
-                mExecutorService.schedule(() -> {
-                    synchronized (mStats) {
-                        mStats.clearRemovedUserUidsLocked(userId);
-                    }
-                }, UID_QUICK_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS, TimeUnit.MILLISECONDS);
-
-                // Final clean-up after a user removal, to take care of UIDs that were running
-                // longer than expected
-                return mExecutorService.schedule(() -> {
-                    synchronized (mStats) {
-                        mStats.clearRemovedUserUidsLocked(userId);
-                    }
-                }, UID_FINAL_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS, TimeUnit.MILLISECONDS);
-            } catch (RejectedExecutionException e) {
-                return CompletableFuture.failedFuture(e);
+    public void scheduleCleanupDueToRemovedUser(int userId) {
+        // Initial quick clean-up after a user removal
+        mHandler.postDelayed(() -> {
+            synchronized (mStats) {
+                mStats.clearRemovedUserUidsLocked(userId);
             }
-        }
+        }, SYNC_USER_REMOVAL, UID_QUICK_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS);
+
+        // Final clean-up after a user removal, to take care of UIDs that were running
+        // longer than expected
+        mHandler.postDelayed(() -> {
+            synchronized (mStats) {
+                mStats.clearRemovedUserUidsLocked(userId);
+            }
+        }, SYNC_USER_REMOVAL, UID_FINAL_REMOVAL_AFTER_USER_REMOVAL_DELAY_MILLIS);
     }
 
     /**
@@ -368,42 +328,27 @@
      *         cancel it if needed
      */
     @GuardedBy("this")
-    private Future<?> scheduleDelayedSyncLocked(Future<?> lastScheduledSync, Runnable syncRunnable,
+    private void scheduleDelayedSyncLocked(int what, Runnable syncRunnable,
             long delayMillis) {
-        if (mExecutorService.isShutdown()) {
-            return CompletableFuture.failedFuture(new IllegalStateException("worker shutdown"));
-        }
-
-        if (lastScheduledSync != null) {
+        if (mHandler.hasMessages(what)) {
             // If there's already a scheduled task, leave it as is if we're trying to
             // re-schedule it again with a delay, otherwise cancel and re-schedule it.
             if (delayMillis == 0) {
-                lastScheduledSync.cancel(false);
+                mHandler.removeMessages(what);
             } else {
-                return lastScheduledSync;
+                return;
             }
         }
 
-        try {
-            return mExecutorService.schedule(syncRunnable, delayMillis, TimeUnit.MILLISECONDS);
-        } catch (RejectedExecutionException e) {
-            return CompletableFuture.failedFuture(e);
-        }
+        mHandler.postDelayed(syncRunnable, what, delayMillis);
     }
 
-    public synchronized Future<?> scheduleWrite() {
-        if (mExecutorService.isShutdown()) {
-            return CompletableFuture.failedFuture(new IllegalStateException("worker shutdown"));
-        }
-
+    /**
+     * Schedule and async writing of battery stats to disk
+     */
+    public synchronized void scheduleWrite() {
         scheduleSyncLocked("write", UPDATE_ALL);
-        // Since we use a single threaded executor, we can assume the next scheduled task's
-        // Future finishes after the sync.
-        try {
-            return mExecutorService.submit(mWriteTask);
-        } catch (RejectedExecutionException e) {
-            return CompletableFuture.failedFuture(e);
-        }
+        mHandler.post(mWriteTask);
     }
 
     /**
@@ -411,34 +356,25 @@
      * within the task, never wait on the resulting Future. This will result in a deadlock.
      */
     public synchronized void scheduleRunnable(Runnable runnable) {
-        try {
-            mExecutorService.submit(runnable);
-        } catch (RejectedExecutionException e) {
-            Slog.e(TAG, "Couldn't schedule " + runnable, e);
-        }
+        mHandler.post(runnable);
     }
 
     public void shutdown() {
-        mExecutorService.shutdownNow();
+        mHandler.removeMessages(SYNC_UPDATE);
+        mHandler.removeMessages(SYNC_WAKELOCK_CHANGE);
+        mHandler.removeMessages(SYNC_BATTERY_LEVEL_CHANGE);
+        mHandler.removeMessages(SYNC_PROCESS_STATE_CHANGE);
+        mHandler.removeMessages(SYNC_USER_REMOVAL);
     }
 
     @GuardedBy("this")
-    private Future<?> scheduleSyncLocked(String reason, int flags) {
-        if (mExecutorService.isShutdown()) {
-            return CompletableFuture.failedFuture(new IllegalStateException("worker shutdown"));
-        }
-
-        if (mCurrentFuture == null) {
+    private void scheduleSyncLocked(String reason, @ExternalUpdateFlag int flags) {
+        if (!mHandler.hasMessages(SYNC_UPDATE)) {
             mUpdateFlags = flags;
             mCurrentReason = reason;
-            try {
-                mCurrentFuture = mExecutorService.submit(mSyncTask);
-            } catch (RejectedExecutionException e) {
-                return CompletableFuture.failedFuture(e);
-            }
+            mHandler.postDelayed(mSyncTask, SYNC_UPDATE, 0);
         }
         mUpdateFlags |= flags;
-        return mCurrentFuture;
     }
 
     public long getLastCollectionTimeStamp() {
@@ -468,7 +404,6 @@
                 useLatestStates = mUseLatestStates;
                 mUpdateFlags = 0;
                 mCurrentReason = null;
-                mCurrentFuture = null;
                 mUseLatestStates = true;
                 if ((updateFlags & UPDATE_ALL) == UPDATE_ALL) {
                     cancelSyncDueToBatteryLevelChangeLocked();
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 391071f..c04158f 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -176,7 +176,6 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.concurrent.Executor;
-import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantLock;
@@ -949,19 +948,38 @@
         public @interface ExternalUpdateFlag {
         }
 
-        Future<?> scheduleSync(String reason, int flags);
-        Future<?> scheduleCpuSyncDueToRemovedUid(int uid);
+        /**
+         * Schedules a sync of kernel metrics in accordance with the specified flags.
+         */
+        void scheduleSync(String reason, @ExternalUpdateFlag int flags);
+
+        /**
+         * Schedules a CPU stats sync after a UID removal.
+         */
+        void scheduleCpuSyncDueToRemovedUid(int uid);
 
         /**
          * Schedule a sync because of a screen state change.
          */
-        Future<?> scheduleSyncDueToScreenStateChange(int flags, boolean onBattery,
+        void scheduleSyncDueToScreenStateChange(@ExternalUpdateFlag int flags, boolean onBattery,
                 boolean onBatteryScreenOff, int screenState, int[] perDisplayScreenStates);
-        Future<?> scheduleCpuSyncDueToWakelockChange(long delayMillis);
+
+        /**
+         * Schedules a sync after a wakelock state change
+         */
+        void scheduleCpuSyncDueToWakelockChange(long delayMillis);
+
+        /**
+         * Canceles any pending sync due to a wakelock state change
+         */
         void cancelCpuSyncDueToWakelockChange();
-        Future<?> scheduleSyncDueToBatteryLevelChange(long delayMillis);
+
+        /**
+         * Schedules a sync caused by the battery level change
+         */
+        void scheduleSyncDueToBatteryLevelChange(long delayMillis);
         /** Schedule removal of UIDs corresponding to a removed user */
-        Future<?> scheduleCleanupDueToRemovedUser(int userId);
+        void scheduleCleanupDueToRemovedUser(int userId);
         /** Schedule a sync because of a process state change */
         void scheduleSyncDueToProcessStateChange(int flags, long delayMillis);
     }
@@ -12263,14 +12281,8 @@
             // start time
             long monotonicStartTime =
                     mMonotonicClock.monotonicTime() - batteryUsageStats.getStatsDuration();
-            mHandler.post(() -> {
-                mPowerStatsStore.storeBatteryUsageStats(monotonicStartTime, batteryUsageStats);
-                try {
-                    batteryUsageStats.close();
-                } catch (IOException e) {
-                    Log.e(TAG, "Cannot close BatteryUsageStats", e);
-                }
-            });
+            commitMonotonicClock();
+            mPowerStatsStore.storeBatteryUsageStatsAsync(monotonicStartTime, batteryUsageStats);
         }
     }
 
@@ -15391,6 +15403,10 @@
         mMaxLearnedBatteryCapacityUah = Math.max(mMaxLearnedBatteryCapacityUah, chargeFullUah);
 
         mBatteryTimeToFullSeconds = chargeTimeToFullSeconds;
+
+        if (mAccumulateBatteryUsageStats) {
+            mBatteryUsageStatsProvider.accumulateBatteryUsageStatsAsync(this, mHandler);
+        }
     }
 
     public static boolean isOnBattery(int plugType, int status) {
@@ -17699,6 +17715,13 @@
         }
     }
 
+    /**
+     * Persists the monotonic clock associated with battery stats.
+     */
+    public void commitMonotonicClock() {
+        mMonotonicClock.write();
+    }
+
     @GuardedBy("this")
     public void prepareForDumpLocked() {
         // Need to retrieve current kernel wake lock stats before printing.
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 b466dd2..265f1dfc 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -23,17 +23,16 @@
 import android.os.BatteryStats;
 import android.os.BatteryUsageStats;
 import android.os.BatteryUsageStatsQuery;
+import android.os.Handler;
 import android.os.Process;
-import android.os.UidBatteryConsumer;
 import android.util.Log;
-import android.util.Slog;
 import android.util.SparseArray;
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.CpuScalingPolicies;
+import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerProfile;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -49,20 +48,31 @@
     private final PowerStatsStore mPowerStatsStore;
     private final PowerProfile mPowerProfile;
     private final CpuScalingPolicies mCpuScalingPolicies;
+    private final int mAccumulatedBatteryUsageStatsSpanSize;
     private final Clock mClock;
     private final Object mLock = new Object();
     private List<PowerCalculator> mPowerCalculators;
     private UserPowerCalculator mUserPowerCalculator;
+    private long mLastAccumulationMonotonicHistorySize;
+
+    private static class AccumulatedBatteryUsageStats {
+        public BatteryUsageStats.Builder builder;
+        public long startWallClockTime;
+        public long startMonotonicTime;
+        public long endMonotonicTime;
+    }
 
     public BatteryUsageStatsProvider(@NonNull Context context,
             @NonNull PowerAttributor powerAttributor,
             @NonNull PowerProfile powerProfile, @NonNull CpuScalingPolicies cpuScalingPolicies,
-            @NonNull PowerStatsStore powerStatsStore, @NonNull Clock clock) {
+            @NonNull PowerStatsStore powerStatsStore, int accumulatedBatteryUsageStatsSpanSize,
+            @NonNull Clock clock) {
         mContext = context;
         mPowerAttributor = powerAttributor;
         mPowerStatsStore = powerStatsStore;
         mPowerProfile = powerProfile;
         mCpuScalingPolicies = cpuScalingPolicies;
+        mAccumulatedBatteryUsageStatsSpanSize = accumulatedBatteryUsageStatsSpanSize;
         mClock = clock;
         mUserPowerCalculator = new UserPowerCalculator();
 
@@ -85,7 +95,10 @@
                     mPowerCalculators.add(
                             new CpuPowerCalculator(mCpuScalingPolicies, mPowerProfile));
                 }
-                mPowerCalculators.add(new MemoryPowerCalculator(mPowerProfile));
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_MEMORY)) {
+                    mPowerCalculators.add(new MemoryPowerCalculator(mPowerProfile));
+                }
                 if (!mPowerAttributor.isPowerComponentSupported(
                         BatteryConsumer.POWER_COMPONENT_WAKELOCK)) {
                     mPowerCalculators.add(new WakelockPowerCalculator(mPowerProfile));
@@ -141,7 +154,11 @@
                         BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY)) {
                     mPowerCalculators.add(new AmbientDisplayPowerCalculator(mPowerProfile));
                 }
-                mPowerCalculators.add(new IdlePowerCalculator(mPowerProfile));
+                // IDLE power attribution is covered by WakelockPowerStatsProcessor
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_WAKELOCK)) {
+                    mPowerCalculators.add(new IdlePowerCalculator(mPowerProfile));
+                }
                 if (!mPowerAttributor.isPowerComponentSupported(
                         BatteryConsumer.POWER_COMPONENT_ANY)) {
                     mPowerCalculators.add(new CustomEnergyConsumerPowerCalculator(mPowerProfile));
@@ -159,53 +176,45 @@
     }
 
     /**
-     * Compute BatteryUsageStats for the period since the last accumulated stats were stored,
-     * add them to the accumulated stats and save the result.
+     * Conditionally runs a battery usage stats accumulation on the supplied handler.
+     */
+    public void accumulateBatteryUsageStatsAsync(BatteryStatsImpl stats, Handler handler) {
+        synchronized (this) {
+            long historySize = stats.getHistory().getMonotonicHistorySize();
+            if (historySize - mLastAccumulationMonotonicHistorySize
+                    < mAccumulatedBatteryUsageStatsSpanSize) {
+                return;
+            }
+            mLastAccumulationMonotonicHistorySize = historySize;
+        }
+
+        handler.post(() -> accumulateBatteryUsageStats(stats));
+    }
+
+    /**
+     * Computes BatteryUsageStats for the period since the last accumulated stats were stored,
+     * adds them to the accumulated stats and saves the result.
      */
     public void accumulateBatteryUsageStats(BatteryStatsImpl stats) {
-        BatteryUsageStats.Builder accumulatedBatteryUsageStatsBuilder = null;
+        AccumulatedBatteryUsageStats accumulatedStats = loadAccumulatedBatteryUsageStats();
 
-        PowerStatsSpan powerStatsSpan = mPowerStatsStore.loadPowerStatsSpan(
-                AccumulatedBatteryUsageStatsSection.ID,
-                AccumulatedBatteryUsageStatsSection.TYPE);
-        if (powerStatsSpan != null) {
-            List<PowerStatsSpan.Section> sections = powerStatsSpan.getSections();
-            for (int i = sections.size() - 1; i >= 0; i--) {
-                PowerStatsSpan.Section section = sections.get(i);
-                if (AccumulatedBatteryUsageStatsSection.TYPE.equals(section.getType())) {
-                    accumulatedBatteryUsageStatsBuilder =
-                            ((AccumulatedBatteryUsageStatsSection) section)
-                                    .getBatteryUsageStatsBuilder();
-                    break;
-                }
-            }
-        }
+        final BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
+                .setMaxStatsAgeMs(0)
+                .includeProcessStateData()
+                .includePowerStateData()
+                .includeScreenStateData()
+                .build();
+        updateAccumulatedBatteryUsageStats(accumulatedStats, stats, query);
 
-        // TODO(b/366493365): add the current batteryusagestats directly into the "accumulated"
-        // builder to avoid allocating a second CursorWindow
-        BatteryUsageStats.Builder currentBatteryUsageStatsBuilder =
-                getCurrentBatteryUsageStatsBuilder(stats,
-                        new BatteryUsageStatsQuery.Builder()
-                                .setMaxStatsAgeMs(0)
-                                .includeProcessStateData()
-                                .includePowerStateData()
-                                .includeScreenStateData()
-                                .build(),
-                        mClock.currentTimeMillis());
-
-        if (accumulatedBatteryUsageStatsBuilder == null) {
-            accumulatedBatteryUsageStatsBuilder = currentBatteryUsageStatsBuilder;
-        } else {
-            accumulatedBatteryUsageStatsBuilder.add(currentBatteryUsageStatsBuilder.build());
-            currentBatteryUsageStatsBuilder.discard();
-        }
-
-        powerStatsSpan = new PowerStatsSpan(AccumulatedBatteryUsageStatsSection.ID);
+        PowerStatsSpan powerStatsSpan = new PowerStatsSpan(AccumulatedBatteryUsageStatsSection.ID);
         powerStatsSpan.addSection(
-                new AccumulatedBatteryUsageStatsSection(accumulatedBatteryUsageStatsBuilder));
-
+                new AccumulatedBatteryUsageStatsSection(accumulatedStats.builder));
+        powerStatsSpan.addTimeFrame(accumulatedStats.startMonotonicTime,
+                accumulatedStats.startWallClockTime,
+                accumulatedStats.endMonotonicTime - accumulatedStats.startMonotonicTime);
+        stats.commitMonotonicClock();
         mPowerStatsStore.storePowerStatsSpanAsync(powerStatsSpan,
-                accumulatedBatteryUsageStatsBuilder::discard);
+                accumulatedStats.builder::discard);
     }
 
     /**
@@ -252,68 +261,73 @@
             BatteryUsageStatsQuery query, long currentTimeMs) {
         if ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_ACCUMULATED) != 0) {
-            return getAccumulatedBatteryUsageStats(stats, query);
-        } else if (query.getToTimestamp() == 0) {
-            return getCurrentBatteryUsageStats(stats, query, currentTimeMs);
+            return getAccumulatedBatteryUsageStats(stats, query, currentTimeMs);
+        } else if (query.getAggregatedToTimestamp() == 0) {
+            BatteryUsageStats.Builder builder = computeBatteryUsageStats(stats, query,
+                    query.getMonotonicStartTime(),
+                    query.getMonotonicEndTime(), currentTimeMs);
+            return builder.build();
         } else {
             return getAggregatedBatteryUsageStats(stats, query);
         }
     }
 
     private BatteryUsageStats getAccumulatedBatteryUsageStats(BatteryStatsImpl stats,
-            BatteryUsageStatsQuery query) {
+            BatteryUsageStatsQuery query, long currentTimeMs) {
+        AccumulatedBatteryUsageStats accumulatedStats = loadAccumulatedBatteryUsageStats();
+        updateAccumulatedBatteryUsageStats(accumulatedStats, stats, query);
+        return accumulatedStats.builder.build();
+    }
+
+    private AccumulatedBatteryUsageStats loadAccumulatedBatteryUsageStats() {
+        AccumulatedBatteryUsageStats stats = new AccumulatedBatteryUsageStats();
+        stats.startWallClockTime = 0;
+        stats.startMonotonicTime = MonotonicClock.UNDEFINED;
+        stats.endMonotonicTime = MonotonicClock.UNDEFINED;
         PowerStatsSpan powerStatsSpan = mPowerStatsStore.loadPowerStatsSpan(
                 AccumulatedBatteryUsageStatsSection.ID,
                 AccumulatedBatteryUsageStatsSection.TYPE);
-
-        BatteryUsageStats.Builder accumulatedBatteryUsageStatsBuilder = null;
         if (powerStatsSpan != null) {
             List<PowerStatsSpan.Section> sections = powerStatsSpan.getSections();
-            if (sections.size() == 1) {
-                accumulatedBatteryUsageStatsBuilder =
-                        ((AccumulatedBatteryUsageStatsSection) sections.get(0))
-                                .getBatteryUsageStatsBuilder();
-            } else {
-                Slog.wtf(TAG, "Unexpected number of sections for type "
-                        + AccumulatedBatteryUsageStatsSection.TYPE);
+            for (int i = sections.size() - 1; i >= 0; i--) {
+                PowerStatsSpan.Section section = sections.get(i);
+                if (AccumulatedBatteryUsageStatsSection.TYPE.equals(section.getType())) {
+                    stats.builder = ((AccumulatedBatteryUsageStatsSection) section)
+                            .getBatteryUsageStatsBuilder();
+                    stats.startWallClockTime = powerStatsSpan.getMetadata().getStartTime();
+                    stats.startMonotonicTime = powerStatsSpan.getMetadata().getStartMonotonicTime();
+                    stats.endMonotonicTime = powerStatsSpan.getMetadata().getEndMonotonicTime();
+                    break;
+                }
             }
         }
+        return stats;
+    }
 
-        BatteryUsageStats currentBatteryUsageStats = getCurrentBatteryUsageStats(stats, query,
+    private void updateAccumulatedBatteryUsageStats(AccumulatedBatteryUsageStats accumulatedStats,
+            BatteryStatsImpl stats, BatteryUsageStatsQuery query) {
+        // TODO(b/366493365): add the current batteryusagestats directly into
+        //  `accumulatedStats.builder` to avoid allocating a second CursorWindow
+        BatteryUsageStats.Builder remainingBatteryUsageStats = computeBatteryUsageStats(stats,
+                query, accumulatedStats.endMonotonicTime, query.getMonotonicEndTime(),
                 mClock.currentTimeMillis());
 
-        BatteryUsageStats result;
-        if (accumulatedBatteryUsageStatsBuilder == null) {
-            result = currentBatteryUsageStats;
+        if (accumulatedStats.builder == null) {
+            accumulatedStats.builder = remainingBatteryUsageStats;
+            accumulatedStats.startWallClockTime = stats.getStartClockTime();
+            accumulatedStats.startMonotonicTime = stats.getMonotonicStartTime();
+            accumulatedStats.endMonotonicTime = accumulatedStats.startMonotonicTime
+                    + accumulatedStats.builder.getStatsDuration();
         } else {
-            accumulatedBatteryUsageStatsBuilder.add(currentBatteryUsageStats);
-            try {
-                currentBatteryUsageStats.close();
-            } catch (IOException ex) {
-                Slog.e(TAG, "Closing BatteryUsageStats", ex);
-            }
-            result = accumulatedBatteryUsageStatsBuilder.build();
+            accumulatedStats.builder.add(remainingBatteryUsageStats.build());
+            accumulatedStats.endMonotonicTime += remainingBatteryUsageStats.getStatsDuration();
+            remainingBatteryUsageStats.discard();
         }
-
-        return result;
     }
 
-    private BatteryUsageStats getCurrentBatteryUsageStats(BatteryStatsImpl stats,
-            BatteryUsageStatsQuery query, long currentTimeMs) {
-        BatteryUsageStats.Builder builder = getCurrentBatteryUsageStatsBuilder(stats, query,
-                currentTimeMs);
-        BatteryUsageStats batteryUsageStats = builder.build();
-        if (batteryUsageStats.isProcessStateDataIncluded()) {
-            verify(batteryUsageStats);
-        }
-        return batteryUsageStats;
-    }
-
-    private BatteryUsageStats.Builder getCurrentBatteryUsageStatsBuilder(BatteryStatsImpl stats,
-            BatteryUsageStatsQuery query, long currentTimeMs) {
-        final long realtimeUs = mClock.elapsedRealtime() * 1000;
-        final long uptimeUs = mClock.uptimeMillis() * 1000;
-
+    private BatteryUsageStats.Builder computeBatteryUsageStats(BatteryStatsImpl stats,
+            BatteryUsageStatsQuery query, long monotonicStartTime, long monotonicEndTime,
+            long currentTimeMs) {
         final boolean includePowerModels = (query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_POWER_MODELS) != 0;
         final boolean includeProcessStateData = ((query.getFlags()
@@ -324,11 +338,8 @@
         final double minConsumedPowerThreshold = query.getMinConsumedPowerThreshold();
 
         String[] customEnergyConsumerNames;
-        long monotonicStartTime, monotonicEndTime;
         synchronized (stats) {
             customEnergyConsumerNames = stats.getCustomEnergyConsumerNames();
-            monotonicStartTime = stats.getMonotonicStartTime();
-            monotonicEndTime = stats.getMonotonicEndTime();
         }
 
         final BatteryUsageStats.Builder batteryUsageStatsBuilder = new BatteryUsageStats.Builder(
@@ -337,12 +348,31 @@
                 query.isPowerStateDataNeeded(), minConsumedPowerThreshold);
 
         synchronized (stats) {
-            // TODO(b/188068523): use a monotonic clock to ensure resilience of order and duration
-            // of batteryUsageStats sessions to wall-clock adjustments
-            batteryUsageStatsBuilder.setStatsStartTimestamp(stats.getStartClockTime());
-            batteryUsageStatsBuilder.setStatsEndTimestamp(currentTimeMs);
             final List<PowerCalculator> powerCalculators = getPowerCalculators();
             if (!powerCalculators.isEmpty()) {
+                if (monotonicStartTime != MonotonicClock.UNDEFINED
+                        || monotonicEndTime != MonotonicClock.UNDEFINED) {
+                    throw new IllegalStateException("BatteryUsageStatsQuery specifies a time "
+                            + "range that is incompatible with PowerCalculators: "
+                            + powerCalculators);
+                }
+            }
+
+            if (monotonicStartTime == MonotonicClock.UNDEFINED) {
+                monotonicStartTime = stats.getMonotonicStartTime();
+            }
+            batteryUsageStatsBuilder.setStatsStartTimestamp(stats.getStartClockTime()
+                    + (monotonicStartTime - stats.getMonotonicStartTime()));
+            if (monotonicEndTime != MonotonicClock.UNDEFINED) {
+                batteryUsageStatsBuilder.setStatsEndTimestamp(stats.getStartClockTime()
+                        + (monotonicEndTime - stats.getMonotonicStartTime()));
+            } else {
+                batteryUsageStatsBuilder.setStatsEndTimestamp(currentTimeMs);
+            }
+
+            if (!powerCalculators.isEmpty()) {
+                final long realtimeUs = mClock.elapsedRealtime() * 1000;
+                final long uptimeUs = mClock.uptimeMillis() * 1000;
                 final int[] powerComponents = query.getPowerComponents();
                 SparseArray<? extends BatteryStats.Uid> uidStats = stats.getUidStats();
                 for (int i = uidStats.size() - 1; i >= 0; i--) {
@@ -381,8 +411,7 @@
                 monotonicStartTime, monotonicEndTime);
 
         // Combine apps by the user if necessary
-        mUserPowerCalculator.calculate(batteryUsageStatsBuilder, stats, realtimeUs, uptimeUs,
-                query);
+        mUserPowerCalculator.calculate(batteryUsageStatsBuilder, stats, 0, 0, query);
 
         populateGeneralInfo(batteryUsageStatsBuilder, stats);
         return batteryUsageStatsBuilder;
@@ -402,48 +431,6 @@
         }
     }
 
-    // STOPSHIP(b/229906525): remove verification before shipping
-    private static boolean sErrorReported;
-
-    private void verify(BatteryUsageStats stats) {
-        if (sErrorReported) {
-            return;
-        }
-
-        final double precision = 2.0;   // Allow rounding errors up to 2 mAh
-        final int[] components =
-                {BatteryConsumer.POWER_COMPONENT_CPU,
-                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
-                        BatteryConsumer.POWER_COMPONENT_WIFI,
-                        BatteryConsumer.POWER_COMPONENT_BLUETOOTH};
-        final int[] states =
-                {BatteryConsumer.PROCESS_STATE_FOREGROUND,
-                        BatteryConsumer.PROCESS_STATE_BACKGROUND,
-                        BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE,
-                        BatteryConsumer.PROCESS_STATE_CACHED};
-        for (UidBatteryConsumer ubc : stats.getUidBatteryConsumers()) {
-            for (int component : components) {
-                double consumedPower = ubc.getConsumedPower(ubc.getKey(component));
-                double sumStates = 0;
-                for (int state : states) {
-                    sumStates += ubc.getConsumedPower(ubc.getKey(component, state));
-                }
-                if (sumStates > consumedPower + precision) {
-                    String error = "Sum of states exceeds total. UID = " + ubc.getUid() + " "
-                            + BatteryConsumer.powerComponentIdToString(component)
-                            + " total = " + consumedPower + " states = " + sumStates;
-                    if (!sErrorReported) {
-                        Slog.wtf(TAG, error);
-                        sErrorReported = true;
-                    } else {
-                        Slog.e(TAG, error);
-                    }
-                    return;
-                }
-            }
-        }
-    }
-
     private BatteryUsageStats getAggregatedBatteryUsageStats(BatteryStatsImpl stats,
             BatteryUsageStatsQuery query) {
         final boolean includePowerModels = (query.getFlags()
@@ -489,8 +476,10 @@
             // Per BatteryUsageStatsQuery API, the "from" timestamp is *exclusive*,
             // while the "to" timestamp is *inclusive*.
             boolean isInRange =
-                    (query.getFromTimestamp() == 0 || minTime > query.getFromTimestamp())
-                            && (query.getToTimestamp() == 0 || maxTime <= query.getToTimestamp());
+                    (query.getAggregatedFromTimestamp() == 0
+                            || minTime > query.getAggregatedFromTimestamp())
+                            && (query.getAggregatedToTimestamp() == 0
+                            || maxTime <= query.getAggregatedToTimestamp());
             if (!isInRange) {
                 continue;
             }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsSpan.java b/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
index fc0611f..5105272 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
@@ -25,6 +25,7 @@
 import android.util.TimeUtils;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.MonotonicClock;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
 
@@ -147,6 +148,40 @@
             mTimeFrames.add(timeFrame);
         }
 
+        long getStartTime() {
+            long startTime = Long.MAX_VALUE;
+            for (int i = 0; i < mTimeFrames.size(); i++) {
+                TimeFrame timeFrame = mTimeFrames.get(i);
+                if (timeFrame.startTime < startTime) {
+                    startTime = timeFrame.startTime;
+                }
+            }
+            return startTime != Long.MAX_VALUE ? startTime : 0;
+        }
+
+        long getStartMonotonicTime() {
+            long startTime = Long.MAX_VALUE;
+            for (int i = 0; i < mTimeFrames.size(); i++) {
+                TimeFrame timeFrame = mTimeFrames.get(i);
+                if (timeFrame.startMonotonicTime < startTime) {
+                    startTime = timeFrame.startMonotonicTime;
+                }
+            }
+            return startTime != Long.MAX_VALUE ? startTime : MonotonicClock.UNDEFINED;
+        }
+
+        long getEndMonotonicTime() {
+            long maxTime = Long.MIN_VALUE;
+            for (int i = 0; i < mTimeFrames.size(); i++) {
+                TimeFrame timeFrame = mTimeFrames.get(i);
+                long endTime = timeFrame.startMonotonicTime + timeFrame.duration;
+                if (endTime > maxTime) {
+                    maxTime = endTime;
+                }
+            }
+            return maxTime != Long.MIN_VALUE ? maxTime : MonotonicClock.UNDEFINED;
+        }
+
         void addSection(String sectionType) {
             // The number of sections per span is small, so there is no need to use a Set
             if (!mSections.contains(sectionType)) {
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsStore.java b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
index 5a6f973..3673617 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsStore.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
@@ -149,7 +149,6 @@
      * Saves the specified span in the store.
      */
     public void storePowerStatsSpan(PowerStatsSpan span) {
-        maybeClearLegacyStore();
         lockStoreDirectory();
         try {
             if (!mStoreDir.exists()) {
@@ -203,13 +202,23 @@
      * Stores a {@link PowerStatsSpan} containing a single section for the supplied
      * battery usage stats.
      */
-    public void storeBatteryUsageStats(long monotonicStartTime,
+    public void storeBatteryUsageStatsAsync(long monotonicStartTime,
             BatteryUsageStats batteryUsageStats) {
-        PowerStatsSpan span = new PowerStatsSpan(monotonicStartTime);
-        span.addTimeFrame(monotonicStartTime, batteryUsageStats.getStatsStartTimestamp(),
-                batteryUsageStats.getStatsDuration());
-        span.addSection(new BatteryUsageStatsSection(batteryUsageStats));
-        storePowerStatsSpan(span);
+        mHandler.post(() -> {
+            try {
+                PowerStatsSpan span = new PowerStatsSpan(monotonicStartTime);
+                span.addTimeFrame(monotonicStartTime, batteryUsageStats.getStatsStartTimestamp(),
+                        batteryUsageStats.getStatsDuration());
+                span.addSection(new BatteryUsageStatsSection(batteryUsageStats));
+                storePowerStatsSpan(span);
+            } finally {
+                try {
+                    batteryUsageStats.close();
+                } catch (IOException e) {
+                    Slog.e(TAG, "Cannot close BatteryUsageStats", e);
+                }
+            }
+        });
     }
 
     /**
diff --git a/services/core/java/com/android/server/power/stats/format/BinaryStatePowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/BinaryStatePowerStatsLayout.java
index 4a26d83..657701b 100644
--- a/services/core/java/com/android/server/power/stats/format/BinaryStatePowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/BinaryStatePowerStatsLayout.java
@@ -21,7 +21,9 @@
 public class BinaryStatePowerStatsLayout extends EnergyConsumerPowerStatsLayout {
     public BinaryStatePowerStatsLayout() {
         addDeviceSectionUsageDuration();
+        addDeviceSectionPowerEstimate();
         addUidSectionUsageDuration();
+        addUidSectionPowerEstimate();
     }
 
     public BinaryStatePowerStatsLayout(PowerStats.Descriptor descriptor) {
diff --git a/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java b/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
index c7ad564..c8170a1 100644
--- a/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
+++ b/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
@@ -117,6 +117,7 @@
                     PowerStatsSpan.Section section = sections.get(k);
                     populateBatteryUsageStatsBuilder(batteryUsageStatsBuilder,
                             ((AggregatedPowerStatsSection) section).getAggregatedPowerStats());
+                    // TODO(b/371614748): close the builder
                 }
             }
 
@@ -197,6 +198,7 @@
                 && powerState == BatteryConsumer.POWER_STATE_BATTERY;
 
         double[] totalPower = new double[1];
+        long[] durationMs = new long[1];
         MultiStateStats.States.forEachTrackedStateCombination(
                 powerComponentStats.getConfig().getDeviceStateConfig(),
                 states -> {
@@ -209,6 +211,7 @@
                     }
 
                     totalPower[0] += layout.getDevicePowerEstimate(deviceStats);
+                    durationMs[0] += layout.getUsageDuration(deviceStats);
 
                     if (hasBatteryLevelProperties) {
                         gatherBatteryLevelInfo(batteryLevelInfo, deviceStats);
@@ -223,9 +226,13 @@
         if (key != null) {
             deviceScope.addConsumedPower(key, totalPower[0],
                     BatteryConsumer.POWER_MODEL_UNDEFINED);
+            deviceScope.addUsageDurationMillis(key, durationMs[0]);
         }
-        deviceScope.addConsumedPower(powerComponentId, totalPower[0],
-                BatteryConsumer.POWER_MODEL_UNDEFINED);
+        key = deviceScope.getKey(powerComponentId, BatteryConsumer.PROCESS_STATE_UNSPECIFIED);
+        if (key != null) {
+            deviceScope.addConsumedPower(key, totalPower[0], BatteryConsumer.POWER_MODEL_UNDEFINED);
+            deviceScope.addUsageDurationMillis(key, durationMs[0]);
+        }
     }
 
     private void gatherBatteryLevelInfo(BatteryLevelInfo batteryLevelInfo, long[] deviceStats) {
@@ -373,9 +380,15 @@
 
             if (resultScreenState != BatteryConsumer.SCREEN_STATE_UNSPECIFIED
                     || resultPowerState != BatteryConsumer.POWER_STATE_UNSPECIFIED) {
-                builder.addConsumedPower(powerComponentId,
-                        powerByProcState[BatteryConsumer.PROCESS_STATE_UNSPECIFIED],
-                        BatteryConsumer.POWER_MODEL_UNDEFINED);
+                BatteryConsumer.Key key = builder.getKey(powerComponentId,
+                        BatteryConsumer.PROCESS_STATE_UNSPECIFIED);
+                if (key != null) {
+                    builder.addConsumedPower(key,
+                            powerByProcState[BatteryConsumer.PROCESS_STATE_UNSPECIFIED],
+                            BatteryConsumer.POWER_MODEL_UNDEFINED);
+                    builder.addUsageDurationMillis(key,
+                            durationByProcState[BatteryConsumer.PROCESS_STATE_UNSPECIFIED]);
+                }
             }
             powerAllApps += powerByProcState[BatteryConsumer.PROCESS_STATE_UNSPECIFIED];
         }
diff --git a/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java b/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java
index dcb47a7..4c9cbc4 100644
--- a/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java
+++ b/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java
@@ -42,27 +42,28 @@
             return;
         }
         StatsEvent.Builder builder = StatsEvent.newBuilder().setAtomId(atom.atomId);
-        for (StatsBootstrapAtomValue value : atom.values) {
+        for (StatsBootstrapAtomValue atomValue : atom.values) {
+            StatsBootstrapAtomValue.Primitive value = atomValue.value;
             switch (value.getTag()) {
-                case StatsBootstrapAtomValue.boolValue:
+                case StatsBootstrapAtomValue.Primitive.boolValue:
                     builder.writeBoolean(value.getBoolValue());
                     break;
-                case StatsBootstrapAtomValue.intValue:
+                case StatsBootstrapAtomValue.Primitive.intValue:
                     builder.writeInt(value.getIntValue());
                     break;
-                case StatsBootstrapAtomValue.longValue:
+                case StatsBootstrapAtomValue.Primitive.longValue:
                     builder.writeLong(value.getLongValue());
                     break;
-                case StatsBootstrapAtomValue.floatValue:
+                case StatsBootstrapAtomValue.Primitive.floatValue:
                     builder.writeFloat(value.getFloatValue());
                     break;
-                case StatsBootstrapAtomValue.stringValue:
+                case StatsBootstrapAtomValue.Primitive.stringValue:
                     builder.writeString(value.getStringValue());
                     break;
-                case StatsBootstrapAtomValue.bytesValue:
+                case StatsBootstrapAtomValue.Primitive.bytesValue:
                     builder.writeByteArray(value.getBytesValue());
                     break;
-                case StatsBootstrapAtomValue.stringArrayValue:
+                case StatsBootstrapAtomValue.Primitive.stringArrayValue:
                     builder.writeStringArray(value.getStringArrayValue());
                     break;
                 default:
@@ -71,6 +72,25 @@
                     return;
 
             }
+            StatsBootstrapAtomValue.Annotation[] annotations = atomValue.annotations;
+            for (StatsBootstrapAtomValue.Annotation annotation : atomValue.annotations) {
+                if (annotation.id != StatsBootstrapAtomValue.Annotation.Id.IS_UID) {
+                    Slog.e(TAG, "Unexpected annotation ID: " + annotation.id
+                            + ", for atom " + atom.atomId + ": only UIDs are supported!");
+                    return;
+                }
+
+                switch (annotation.value.getTag()) {
+                    case StatsBootstrapAtomValue.Annotation.Primitive.boolValue:
+                        builder.addBooleanAnnotation(
+                                annotation.id, annotation.value.getBoolValue());
+                        break;
+                    default:
+                        Slog.e(TAG, "Unexpected value type " + annotation.value.getTag()
+                                + " when logging UID for atom " + atom.atomId);
+                        return;
+                }
+            }
         }
         StatsLog.write(builder.usePooledBuffer().build());
     }
diff --git a/services/core/java/com/android/server/vcn/VcnContext.java b/services/core/java/com/android/server/vcn/VcnContext.java
index 6a4c9c2..a492a72 100644
--- a/services/core/java/com/android/server/vcn/VcnContext.java
+++ b/services/core/java/com/android/server/vcn/VcnContext.java
@@ -70,10 +70,6 @@
         return mIsInTestMode;
     }
 
-    public boolean isFlagNetworkMetricMonitorEnabled() {
-        return mFeatureFlags.networkMetricMonitor();
-    }
-
     public boolean isFlagIpSecTransformStateEnabled() {
         // TODO: b/328844044: Ideally this code should gate the behavior by checking the
         // android.net.platform.flags.ipsec_transform_state flag but that flag is not accessible
diff --git a/services/core/java/com/android/server/vcn/VcnGatewayConnection.java b/services/core/java/com/android/server/vcn/VcnGatewayConnection.java
index b574782..a81ad22 100644
--- a/services/core/java/com/android/server/vcn/VcnGatewayConnection.java
+++ b/services/core/java/com/android/server/vcn/VcnGatewayConnection.java
@@ -1913,7 +1913,6 @@
                 mIpSecManager.applyTunnelModeTransform(tunnelIface, direction, transform);
 
                 if (direction == IpSecManager.DIRECTION_IN
-                        && mVcnContext.isFlagNetworkMetricMonitorEnabled()
                         && mVcnContext.isFlagIpSecTransformStateEnabled()) {
                     mUnderlyingNetworkController.updateInboundTransform(mUnderlying, transform);
                 }
diff --git a/services/core/java/com/android/server/vcn/routeselection/NetworkMetricMonitor.java b/services/core/java/com/android/server/vcn/routeselection/NetworkMetricMonitor.java
index b9b1060..0d4c373 100644
--- a/services/core/java/com/android/server/vcn/routeselection/NetworkMetricMonitor.java
+++ b/services/core/java/com/android/server/vcn/routeselection/NetworkMetricMonitor.java
@@ -62,12 +62,6 @@
             @Nullable PersistableBundleWrapper carrierConfig,
             @NonNull NetworkMetricMonitorCallback callback)
             throws IllegalAccessException {
-        if (!vcnContext.isFlagNetworkMetricMonitorEnabled()) {
-            // Caller error
-            logWtf("networkMetricMonitor flag disabled");
-            throw new IllegalAccessException("networkMetricMonitor flag disabled");
-        }
-
         mVcnContext = Objects.requireNonNull(vcnContext, "Missing vcnContext");
         mNetwork = Objects.requireNonNull(network, "Missing network");
         mCallback = Objects.requireNonNull(callback, "Missing callback");
diff --git a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java
index 2b0ca08..ad5bc72 100644
--- a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java
+++ b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkController.java
@@ -204,8 +204,7 @@
         List<NetworkCallback> oldCellCallbacks = new ArrayList<>(mCellBringupCallbacks);
         mCellBringupCallbacks.clear();
 
-        if (mVcnContext.isFlagNetworkMetricMonitorEnabled()
-                && mVcnContext.isFlagIpSecTransformStateEnabled()) {
+        if (mVcnContext.isFlagIpSecTransformStateEnabled()) {
             for (UnderlyingNetworkEvaluator evaluator : mUnderlyingNetworkRecords.values()) {
                 evaluator.close();
             }
@@ -431,8 +430,7 @@
                 .getAllSubIdsInGroup(mSubscriptionGroup)
                 .equals(newSnapshot.getAllSubIdsInGroup(mSubscriptionGroup))) {
 
-            if (mVcnContext.isFlagNetworkMetricMonitorEnabled()
-                    && mVcnContext.isFlagIpSecTransformStateEnabled()) {
+            if (mVcnContext.isFlagIpSecTransformStateEnabled()) {
                 reevaluateNetworks();
             }
             return;
@@ -447,8 +445,7 @@
      */
     public void updateInboundTransform(
             @NonNull UnderlyingNetworkRecord currentNetwork, @NonNull IpSecTransform transform) {
-        if (!mVcnContext.isFlagNetworkMetricMonitorEnabled()
-                || !mVcnContext.isFlagIpSecTransformStateEnabled()) {
+        if (!mVcnContext.isFlagIpSecTransformStateEnabled()) {
             logWtf("#updateInboundTransform: unexpected call; flags missing");
             return;
         }
@@ -575,8 +572,7 @@
 
         @Override
         public void onLost(@NonNull Network network) {
-            if (mVcnContext.isFlagNetworkMetricMonitorEnabled()
-                    && mVcnContext.isFlagIpSecTransformStateEnabled()) {
+            if (mVcnContext.isFlagIpSecTransformStateEnabled()) {
                 mUnderlyingNetworkRecords.get(network).close();
             }
 
@@ -652,8 +648,7 @@
     class NetworkEvaluatorCallbackImpl implements NetworkEvaluatorCallback {
         @Override
         public void onEvaluationResultChanged() {
-            if (!mVcnContext.isFlagNetworkMetricMonitorEnabled()
-                    || !mVcnContext.isFlagIpSecTransformStateEnabled()) {
+            if (!mVcnContext.isFlagIpSecTransformStateEnabled()) {
                 logWtf("#onEvaluationResultChanged: unexpected call; flags missing");
                 return;
             }
diff --git a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkEvaluator.java b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkEvaluator.java
index c852fb4..53b0751 100644
--- a/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkEvaluator.java
+++ b/services/core/java/com/android/server/vcn/routeselection/UnderlyingNetworkEvaluator.java
@@ -193,8 +193,7 @@
     }
 
     private static boolean isIpSecPacketLossDetectorEnabled(VcnContext vcnContext) {
-        return vcnContext.isFlagIpSecTransformStateEnabled()
-                && vcnContext.isFlagNetworkMetricMonitorEnabled();
+        return vcnContext.isFlagIpSecTransformStateEnabled();
     }
 
     /** Get the comparator for UnderlyingNetworkEvaluator */
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 460de01..054f931 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -5508,7 +5508,8 @@
                 clearAllDrawn();
                 // Reset the draw state in order to prevent the starting window to be immediately
                 // dismissed when the app still has the surface.
-                if (!isVisible() && !isClientVisible()) {
+                if (!Flags.resetDrawStateOnClientInvisible()
+                        && !isVisible() && !isClientVisible()) {
                     forAllWindows(w -> {
                         if (w.mWinAnimator.mDrawState == HAS_DRAWN) {
                             w.mWinAnimator.resetDrawState();
@@ -6852,7 +6853,7 @@
         } else if (associatedTask.getActivity(
                 r -> r.isVisibleRequested() && !r.firstWindowDrawn) == null) {
             // The last drawn activity may not be the one that owns the starting window.
-            final ActivityRecord r = associatedTask.topActivityContainsStartingWindow();
+            final ActivityRecord r = associatedTask.getActivity(ar -> ar.mStartingData != null);
             if (r != null) {
                 r.removeStartingWindow();
             }
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index 70f9ebb..ccd5996 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -201,6 +201,9 @@
             infoBuilder.setTouchableRegion(window.getFrame());
             infoBuilder.setAppProgressAllowed((window.getAttrs().privateFlags
                     & PRIVATE_FLAG_APP_PROGRESS_GENERATION_ALLOWED) != 0);
+            if (currentTask != null) {
+                infoBuilder.setFocusedTaskId(currentTask.mTaskId);
+            }
             mNavigationMonitor.startMonitor(window, navigationObserver);
 
             ProtoLog.d(WM_DEBUG_BACK_PREVIEW, "startBackNavigation currentTask=%s, "
diff --git a/services/core/java/com/android/server/wm/ScreenRotationAnimation.java b/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
index 0d3fa1b..4bd294b 100644
--- a/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
+++ b/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
@@ -32,7 +32,6 @@
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 import static com.android.server.wm.utils.CoordinateTransforms.computeRotationMatrix;
-import static com.android.window.flags.Flags.deleteCaptureDisplay;
 
 import android.animation.ArgbEvaluator;
 import android.content.Context;
@@ -41,11 +40,9 @@
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
-import android.os.IBinder;
 import android.os.Trace;
 import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
-import android.view.DisplayAddress;
 import android.view.DisplayInfo;
 import android.view.Surface;
 import android.view.Surface.OutOfResourcesException;
@@ -58,7 +55,6 @@
 import com.android.internal.R;
 import com.android.internal.policy.TransitionAnimation;
 import com.android.internal.protolog.ProtoLog;
-import com.android.server.display.DisplayControl;
 import com.android.server.wm.SurfaceAnimator.AnimationType;
 import com.android.server.wm.SurfaceAnimator.OnAnimationFinishedCallback;
 
@@ -171,33 +167,7 @@
 
         try {
             final ScreenCapture.ScreenshotHardwareBuffer screenshotBuffer;
-            if (isSizeChanged && !deleteCaptureDisplay()) {
-                final DisplayAddress address = displayInfo.address;
-                if (!(address instanceof DisplayAddress.Physical)) {
-                    Slog.e(TAG, "Display does not have a physical address: " + displayId);
-                    return;
-                }
-                final DisplayAddress.Physical physicalAddress =
-                        (DisplayAddress.Physical) address;
-                final IBinder displayToken = DisplayControl.getPhysicalDisplayToken(
-                        physicalAddress.getPhysicalDisplayId());
-                if (displayToken == null) {
-                    Slog.e(TAG, "Display token is null.");
-                    return;
-                }
-                // Temporarily not skip screenshot for the rounded corner overlays and screenshot
-                // the whole display to include the rounded corner overlays.
-                setSkipScreenshotForRoundedCornerOverlays(false, t);
-                mRoundedCornerOverlay = displayContent.findRoundedCornerOverlays();
-                final ScreenCapture.DisplayCaptureArgs captureArgs =
-                        new ScreenCapture.DisplayCaptureArgs.Builder(displayToken)
-                                .setSourceCrop(new Rect(0, 0, width, height))
-                                .setAllowProtected(true)
-                                .setCaptureSecureLayers(true)
-                                .setHintForSeamlessTransition(true)
-                                .build();
-                screenshotBuffer = ScreenCapture.captureDisplay(captureArgs);
-            } else if (isSizeChanged) {
+            if (isSizeChanged) {
                 // Temporarily not skip screenshot for the rounded corner overlays and screenshot
                 // the whole display to include the rounded corner overlays.
                 setSkipScreenshotForRoundedCornerOverlays(false, t);
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 76f2437..6141876 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -3304,6 +3304,16 @@
                 android.os.Process.killProcess(mSession.mPid);
             }
         }
+
+        // Because the client is notified to be invisible, it should no longer be considered as
+        // drawn state. This prevent the app from showing incomplete content if the app is
+        // requested to be visible in a short time (e.g. before activity stopped).
+        if (Flags.resetDrawStateOnClientInvisible() && !clientVisible && mActivityRecord != null
+                && mWinAnimator.mDrawState == HAS_DRAWN) {
+            mWinAnimator.resetDrawState();
+            // Make sure the app can report drawn if it becomes visible again.
+            forceReportingResized();
+        }
     }
 
     void onStartFreezingScreen() {
diff --git a/services/core/java/com/android/server/wm/WindowTracingPerfetto.java b/services/core/java/com/android/server/wm/WindowTracingPerfetto.java
index 6e8094a..1b42e13 100644
--- a/services/core/java/com/android/server/wm/WindowTracingPerfetto.java
+++ b/services/core/java/com/android/server/wm/WindowTracingPerfetto.java
@@ -159,22 +159,28 @@
 
     void onStart(WindowTracingDataSource.Config config) {
         if (config.mLogFrequency == WindowTracingLogFrequency.FRAME) {
+            Log.i(TAG, "Started session (frequency=FRAME, log level=" + config.mLogFrequency + ")");
             mCountSessionsOnFrame.incrementAndGet();
         } else if (config.mLogFrequency == WindowTracingLogFrequency.TRANSACTION) {
+            Log.i(TAG, "Started session (frequency=TRANSACTION, log level="
+                    + config.mLogFrequency + ")");
             mCountSessionsOnTransaction.incrementAndGet();
         }
 
-        Log.i(TAG, "Started with logLevel: " + config.mLogLevel
-                + " logFrequency: " + config.mLogFrequency);
+        Log.i(TAG, getStatus());
+
         log(WHERE_START_TRACING);
     }
 
     void onStop(WindowTracingDataSource.Config config) {
         if (config.mLogFrequency == WindowTracingLogFrequency.FRAME) {
+            Log.i(TAG, "Stopped session (frequency=FRAME)");
             mCountSessionsOnFrame.decrementAndGet();
+            Log.i(TAG, "Stopped session (frequency=TRANSACTION)");
         } else if (config.mLogFrequency == WindowTracingLogFrequency.TRANSACTION) {
             mCountSessionsOnTransaction.decrementAndGet();
         }
-        Log.i(TAG, "Stopped");
+
+        Log.i(TAG, getStatus());
     }
 }
diff --git a/services/core/xsd/vts/Android.bp b/services/core/xsd/vts/Android.bp
index 4d3c79e..e1478d6 100644
--- a/services/core/xsd/vts/Android.bp
+++ b/services/core/xsd/vts/Android.bp
@@ -15,6 +15,7 @@
 //
 
 package {
+    default_team: "trendy_team_android_kernel",
     // 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"
diff --git a/services/java/com/android/server/flags.aconfig b/services/java/com/android/server/flags.aconfig
index 0990691..e2ac22d 100644
--- a/services/java/com/android/server/flags.aconfig
+++ b/services/java/com/android/server/flags.aconfig
@@ -10,6 +10,14 @@
 }
 
 flag {
+    name: "remove_java_service_manager_cache"
+    namespace: "system_performance"
+    description: "This flag turns off Java's Service Manager caching mechanism."
+    bug: "333854840"
+    is_fixed_read_only: true
+}
+
+flag {
      name: "remove_text_service"
      namespace: "wear_frameworks"
      description: "Remove TextServiceManagerService on Wear"
diff --git a/services/tests/mockingservicestests/src/com/android/server/alarm/OWNERS b/services/tests/mockingservicestests/src/com/android/server/alarm/OWNERS
index 6f207fb1..6eb986b 100644
--- a/services/tests/mockingservicestests/src/com/android/server/alarm/OWNERS
+++ b/services/tests/mockingservicestests/src/com/android/server/alarm/OWNERS
@@ -1 +1 @@
-include /apex/jobscheduler/OWNERS
+include /apex/jobscheduler/ALARM_OWNERS
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/OWNERS b/services/tests/mockingservicestests/src/com/android/server/job/OWNERS
index 6f207fb1..c8345f7 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/OWNERS
+++ b/services/tests/mockingservicestests/src/com/android/server/job/OWNERS
@@ -1 +1 @@
-include /apex/jobscheduler/OWNERS
+include /apex/jobscheduler/JOB_OWNERS
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
index 7b635d4..d7b60cf 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryExternalStatsWorkerTest.java
@@ -39,7 +39,7 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.os.connectivity.WifiActivityEnergyInfo;
-import android.platform.test.ravenwood.RavenwoodRule;
+import android.platform.test.ravenwood.RavenwoodConfig;
 import android.power.PowerStatsInternal;
 import android.util.IntArray;
 import android.util.SparseArray;
@@ -52,7 +52,6 @@
 import com.android.internal.os.PowerProfile;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.Arrays;
@@ -67,25 +66,27 @@
 @SuppressWarnings("GuardedBy")
 @android.platform.test.annotations.DisabledOnRavenwood
 public class BatteryExternalStatsWorkerTest {
-    @Rule
-    public final RavenwoodRule mRavenwood = new RavenwoodRule();
+    @RavenwoodConfig.Config
+    public final RavenwoodConfig mRavenwood = new RavenwoodConfig.Builder().build();
 
     private BatteryExternalStatsWorker mBatteryExternalStatsWorker;
     private TestPowerStatsInternal mPowerStatsInternal;
+    private Handler mHandler;
 
     @Before
     public void setUp() {
         final Context context = InstrumentationRegistry.getContext();
 
+        mHandler = new Handler(Looper.getMainLooper());
         BatteryStatsImpl batteryStats = new BatteryStatsImpl(
                 new BatteryStatsImpl.BatteryStatsConfig.Builder().build(), Clock.SYSTEM_CLOCK,
                 new MonotonicClock(0, Clock.SYSTEM_CLOCK), null,
-                new Handler(Looper.getMainLooper()), null, null, null,
+                mHandler, null, null, null,
                 new PowerProfile(context, true /* forTest */), buildScalingPolicies(),
                 new PowerStatsUidResolver());
         mPowerStatsInternal = new TestPowerStatsInternal();
         mBatteryExternalStatsWorker =
-                new BatteryExternalStatsWorker(new TestInjector(context), batteryStats);
+                new BatteryExternalStatsWorker(new TestInjector(context), batteryStats, mHandler);
     }
 
     @Test
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryIteratorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryIteratorTest.java
index d36b553..d6f5036 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryIteratorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryIteratorTest.java
@@ -38,7 +38,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Random;
-import java.util.concurrent.Future;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -324,9 +323,8 @@
         private boolean mSyncScheduled;
 
         @Override
-        public Future<?> scheduleCpuSyncDueToWakelockChange(long delayMillis) {
+        public void scheduleCpuSyncDueToWakelockChange(long delayMillis) {
             mSyncScheduled = true;
-            return null;
         }
 
         public boolean isSyncScheduled() {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java
index e40a3e3..b67ec8b 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsHistoryTest.java
@@ -653,6 +653,44 @@
     }
 
     @Test
+    public void getMonotonicHistorySize() {
+        long lastHistorySize = mHistory.getMonotonicHistorySize();
+        mHistory.forceRecordAllHistory();
+
+        mClock.realtime = 1000;
+        mClock.uptime = 1000;
+        mHistory.recordEvent(mClock.realtime, mClock.uptime,
+                BatteryStats.HistoryItem.EVENT_JOB_START, "job", 42);
+        long size = mHistory.getMonotonicHistorySize();
+        assertThat(size).isGreaterThan(lastHistorySize);
+        lastHistorySize = size;
+
+        mHistory.startNextFile(mClock.realtime);
+
+        size = mHistory.getMonotonicHistorySize();
+        assertThat(size).isEqualTo(lastHistorySize);
+
+        mClock.realtime = 2000;
+        mClock.uptime = 2000;
+        mHistory.recordEvent(mClock.realtime, mClock.uptime,
+                BatteryStats.HistoryItem.EVENT_JOB_FINISH, "job", 42);
+
+        size = mHistory.getMonotonicHistorySize();
+        assertThat(size).isGreaterThan(lastHistorySize);
+        lastHistorySize = size;
+
+        mHistory.startNextFile(mClock.realtime);
+
+        mClock.realtime = 3000;
+        mClock.uptime = 3000;
+        mHistory.recordEvent(mClock.realtime, mClock.uptime,
+                HistoryItem.EVENT_ALARM, "alarm", 42);
+
+        size = mHistory.getMonotonicHistorySize();
+        assertThat(size).isGreaterThan(lastHistorySize);
+    }
+
+    @Test
     public void testVarintParceler() {
         long[] values = {
                 0,
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
index 2408cc1..177f30a 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
@@ -159,7 +159,7 @@
         }
         mPowerStatsStore = new PowerStatsStore(systemDir, mHandler);
         mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mPowerAttributor,
-                mPowerProfile, mBatteryStatsImpl.getCpuScalingPolicies(), mPowerStatsStore,
+                mPowerProfile, mBatteryStatsImpl.getCpuScalingPolicies(), mPowerStatsStore, 0,
                 mMockClock);
     }
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
index 3931201..5912a99 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsUserLifecycleTests.java
@@ -134,7 +134,7 @@
 
     private int getNumberOfUidsInBatteryStats() throws Exception {
         ArraySet<Integer> uids = new ArraySet<>();
-        final String dumpsys = executeShellCommand("dumpsys batterystats --checkin");
+        final String dumpsys = executeShellCommand("dumpsys batterystats -c");
         for (String line : dumpsys.split("\n")) {
             final String[] parts = line.trim().split(",");
             if (parts.length < 5 ||
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
index b30224b..e9d95fc 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
@@ -22,6 +22,7 @@
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -119,7 +120,7 @@
                 .isWithin(PRECISION).of(0.4);
 
         assertThat(batteryUsageStats.getStatsStartTimestamp()).isEqualTo(12345);
-        assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(54321);
+        assertThat(batteryUsageStats.getStatsEndTimestamp()).isEqualTo(180 * MINUTE_IN_MS);
     }
 
     @Test
@@ -142,7 +143,7 @@
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), 0, mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
                 provider.getBatteryUsageStats(batteryStats,
@@ -249,7 +250,8 @@
             }
         }
 
-        mStatsRule.setCurrentTime(54321);
+        setTime(180 * MINUTE_IN_MS);
+
         return batteryStats;
     }
 
@@ -266,7 +268,7 @@
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 powerAttributor, mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), 0, mMockClock);
 
         return provider.getBatteryUsageStats(batteryStats, BatteryUsageStatsQuery.DEFAULT);
     }
@@ -296,7 +298,7 @@
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), 0, mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
                 provider.getBatteryUsageStats(batteryStats,
@@ -385,7 +387,7 @@
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), 0, mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
                 provider.getBatteryUsageStats(batteryStats,
@@ -474,7 +476,7 @@
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), powerStatsStore, mMockClock);
+                mStatsRule.getCpuScalingPolicies(), powerStatsStore, 0, mMockClock);
 
         batteryStats.saveBatteryUsageStatsOnReset(provider, powerStatsStore,
                 /* accumulateBatteryUsageStats */ false);
@@ -564,8 +566,19 @@
     }
 
     @Test
-    public void accumulateBatteryUsageStats() {
+    public void accumulateBatteryUsageStats() throws Throwable {
+        accumulateBatteryUsageStats(10000000, 1);
+        // Accumulate every 200 bytes of battery history
+        accumulateBatteryUsageStats(200, 2);
+        accumulateBatteryUsageStats(50, 5);
+        // Accumulate on every invocation of accumulateBatteryUsageStats
+        accumulateBatteryUsageStats(0, 7);
+    }
+
+    private void accumulateBatteryUsageStats(int accumulatedBatteryUsageStatsSpanSize,
+            int expectedNumberOfUpdates) throws Throwable {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+        batteryStats.forceRecordAllHistory();
 
         setTime(5 * MINUTE_IN_MS);
 
@@ -574,69 +587,86 @@
             batteryStats.resetAllStatsAndHistoryLocked(BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
         }
 
-        PowerStatsStore powerStatsStore = new PowerStatsStore(
+        PowerStatsStore powerStatsStore = spy(new PowerStatsStore(
                 new File(mStatsRule.getHistoryDir(), getClass().getSimpleName()),
-                mStatsRule.getHandler());
+                mStatsRule.getHandler()));
         powerStatsStore.reset();
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
-                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), powerStatsStore, mMockClock);
+        int[] count = new int[1];
+        doAnswer(inv -> {
+            count[0]++;
+            return null;
+        }).when(powerStatsStore).storePowerStatsSpan(any(PowerStatsSpan.class));
 
-        batteryStats.saveBatteryUsageStatsOnReset(provider, powerStatsStore,
-                /* accumulateBatteryUsageStats */ true);
+        MultiStatePowerAttributor powerAttributor = new MultiStatePowerAttributor(mContext,
+                powerStatsStore, mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(),
+                () -> 3500, new PowerStatsUidResolver());
+        for (int powerComponentId = 0; powerComponentId < BatteryConsumer.POWER_COMPONENT_COUNT;
+                powerComponentId++) {
+            powerAttributor.setPowerComponentSupported(powerComponentId, true);
+        }
+        powerAttributor.setPowerComponentSupported(BatteryConsumer.POWER_COMPONENT_ANY, true);
+
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                powerAttributor, mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), powerStatsStore,
+                accumulatedBatteryUsageStatsSpanSize, mMockClock);
+
+        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
 
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOnLocked(APP_UID,
                     10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
         }
+
+        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOffLocked(APP_UID,
                     20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
         }
 
-        synchronized (batteryStats) {
-            batteryStats.resetAllStatsAndHistoryLocked(BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
-        }
+        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
 
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOnLocked(APP_UID,
                     30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS);
         }
+
+        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOffLocked(APP_UID,
                     50 * MINUTE_IN_MS, 50 * MINUTE_IN_MS);
         }
         setTime(55 * MINUTE_IN_MS);
-        synchronized (batteryStats) {
-            batteryStats.resetAllStatsAndHistoryLocked(BatteryStatsImpl.RESET_REASON_ADB_COMMAND);
-        }
+
+        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
 
         // This section has not been saved yet, but should be added to the accumulated totals
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOnLocked(APP_UID,
                     80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         }
+
+        provider.accumulateBatteryUsageStatsAsync(batteryStats, mStatsRule.getHandler());
+
         synchronized (batteryStats) {
             batteryStats.noteFlashlightOffLocked(APP_UID,
                     110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
         }
         setTime(115 * MINUTE_IN_MS);
 
-        // Await completion
-        ConditionVariable done = new ConditionVariable();
-        mStatsRule.getHandler().post(done::open);
-        done.block();
+        // Pick up the remainder of battery history that has not yet been accumulated
+        provider.accumulateBatteryUsageStats(batteryStats);
+
+        mStatsRule.waitForBackgroundThread();
 
         BatteryUsageStats stats = provider.getBatteryUsageStats(batteryStats,
                 new BatteryUsageStatsQuery.Builder().accumulated().build());
-
         assertThat(stats.getStatsStartTimestamp()).isEqualTo(5 * MINUTE_IN_MS);
         assertThat(stats.getStatsEndTimestamp()).isEqualTo(115 * MINUTE_IN_MS);
 
-        // Section 1 (saved): 20 - 10 = 10
-        // Section 2 (saved): 50 - 30 = 20
-        // Section 3 (fresh): 110 - 80 = 30
         // Total: 10 + 20 + 30 = 60
         assertThat(stats.getAggregateBatteryConsumer(
                         BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
@@ -657,6 +687,8 @@
         assertThat(uidBatteryConsumer
                 .getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT))
                 .isEqualTo(60 * MINUTE_IN_MS);
+
+        assertThat(count[0]).isEqualTo(expectedNumberOfUpdates);
     }
 
     private void setTime(long timeMs) {
@@ -682,7 +714,7 @@
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), 0, mMockClock);
 
         PowerStatsStore powerStatsStore = mock(PowerStatsStore.class);
         doAnswer(invocation -> {
@@ -702,7 +734,7 @@
             assertThat(uid.getConsumedPower(componentId1))
                     .isWithin(PRECISION).of(8.33333);
             return null;
-        }).when(powerStatsStore).storeBatteryUsageStats(anyLong(), any());
+        }).when(powerStatsStore).storeBatteryUsageStatsAsync(anyLong(), any());
 
         mStatsRule.getBatteryStats().saveBatteryUsageStatsOnReset(provider, powerStatsStore,
                 /* accumulateBatteryUsageStats */ false);
@@ -714,7 +746,7 @@
 
         mStatsRule.waitForBackgroundThread();
 
-        verify(powerStatsStore).storeBatteryUsageStats(anyLong(), any());
+        verify(powerStatsStore).storeBatteryUsageStatsAsync(anyLong(), any());
     }
 
     @Test
@@ -746,7 +778,7 @@
 
         BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
                 mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
-                mStatsRule.getCpuScalingPolicies(), powerStatsStore, mMockClock);
+                mStatsRule.getCpuScalingPolicies(), powerStatsStore, 0, mMockClock);
 
         BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
                 .aggregateSnapshots(0, 3000)
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
index 2c03f9d..1e4454c 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MockBatteryStatsImpl.java
@@ -43,7 +43,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Queue;
-import java.util.concurrent.Future;
 
 /**
  * Mocks a BatteryStatsImpl object.
@@ -288,30 +287,25 @@
         public int flags = 0;
 
         @Override
-        public Future<?> scheduleSync(String reason, int flags) {
-            return null;
+        public void scheduleSync(String reason, int flags) {
         }
 
         @Override
-        public Future<?> scheduleCleanupDueToRemovedUser(int userId) {
-            return null;
+        public void scheduleCleanupDueToRemovedUser(int userId) {
         }
 
         @Override
-        public Future<?> scheduleCpuSyncDueToRemovedUid(int uid) {
-            return null;
+        public void scheduleCpuSyncDueToRemovedUid(int uid) {
         }
 
         @Override
-        public Future<?> scheduleSyncDueToScreenStateChange(int flag, boolean onBattery,
+        public void scheduleSyncDueToScreenStateChange(int flag, boolean onBattery,
                 boolean onBatteryScreenOff, int screenState, int[] perDisplayScreenStates) {
             flags |= flag;
-            return null;
         }
 
         @Override
-        public Future<?> scheduleCpuSyncDueToWakelockChange(long delayMillis) {
-            return null;
+        public void scheduleCpuSyncDueToWakelockChange(long delayMillis) {
         }
 
         @Override
@@ -319,8 +313,7 @@
         }
 
         @Override
-        public Future<?> scheduleSyncDueToBatteryLevelChange(long delayMillis) {
-            return null;
+        public void scheduleSyncDueToBatteryLevelChange(long delayMillis) {
         }
 
         @Override
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java
index 03491bc..0d5d277 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockPowerStatsCollectorTest.java
@@ -65,12 +65,14 @@
     private WakelockPowerStatsLayout mStatsLayout = new WakelockPowerStatsLayout();
 
     @Before
-    public void setup() {
-        mBatteryStats = new MockBatteryStatsImpl(mClock);
+    public void setup() throws Throwable {
+        mBatteryStats = mStatsRule.getBatteryStats();
         mBatteryStats.setPowerStatsCollectorEnabled(POWER_COMPONENT_WAKELOCK, true);
         mBatteryStats.getPowerStatsCollector(POWER_COMPONENT_WAKELOCK)
                 .addConsumer(ps -> mPowerStats = ps);
         mBatteryStats.onSystemReady(mock(Context.class));
+        // onSystemReady schedules the initial power stats collection. Wait for it to finish
+        mStatsRule.waitForBackgroundThread();
     }
 
     @Test
@@ -79,9 +81,6 @@
         PowerStatsCollector powerStatsCollector = mBatteryStats.getPowerStatsCollector(
                 POWER_COMPONENT_WAKELOCK);
 
-        // Establish a baseline
-        powerStatsCollector.collectAndDeliverStats();
-
         mBatteryStats.forceRecordAllHistory();
 
         mStatsRule.advanceSuspendedTime(1000);
diff --git a/services/tests/servicestests/src/com/android/server/OWNERS b/services/tests/servicestests/src/com/android/server/OWNERS
index d49bc43..d8a9400 100644
--- a/services/tests/servicestests/src/com/android/server/OWNERS
+++ b/services/tests/servicestests/src/com/android/server/OWNERS
@@ -1,4 +1,4 @@
-per-file *Alarm* = file:/apex/jobscheduler/OWNERS
+per-file *Alarm* = file:/apex/jobscheduler/ALARM_OWNERS
 per-file *AppOp* = file:/core/java/android/permission/OWNERS
 per-file *BinaryTransparency* = file:/core/java/android/transparency/OWNERS
 per-file *Bluetooth* = file:platform/packages/modules/Bluetooth:master:/framework/OWNERS
diff --git a/services/tests/servicestests/src/com/android/server/am/BatteryStatsServiceTest.java b/services/tests/servicestests/src/com/android/server/am/BatteryStatsServiceTest.java
index 467c15d..ea70287 100644
--- a/services/tests/servicestests/src/com/android/server/am/BatteryStatsServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/BatteryStatsServiceTest.java
@@ -55,6 +55,7 @@
         File systemDir = context.getCacheDir();
         Handler handler = new Handler(mBgThread.getLooper());
         mBatteryStatsService = new BatteryStatsService(context, systemDir);
+        mBatteryStatsService.setRailsStatsCollectionEnabled(false);
     }
 
     @After
diff --git a/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java b/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java
index e652df5..f994660 100644
--- a/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/pm/UserManagerTest.java
@@ -539,6 +539,79 @@
 
     @MediumTest
     @Test
+    public void testRemoveUser_shouldRemovePrivateUser() {
+        UserInfo privateProfileUser =
+                createProfileForUser(
+                        "Private profile",
+                        UserManager.USER_TYPE_PROFILE_PRIVATE,
+                        mUserManager.getMainUser().getIdentifier());
+        assertThat(privateProfileUser).isNotNull();
+        assertThat(hasUser(privateProfileUser.id)).isTrue();
+
+        removeUser(privateProfileUser.id);
+
+        assertThat(hasUser(privateProfileUser.id)).isFalse();
+    }
+
+    @MediumTest
+    @Test
+    @RequiresFlagsEnabled(
+            android.multiuser.Flags.FLAG_IGNORE_RESTRICTIONS_WHEN_DELETING_PRIVATE_PROFILE)
+    public void testRemoveUser_shouldRemovePrivateUser_withDisallowRemoveUserRestriction() {
+        UserHandle mainUser = mUserManager.getMainUser();
+        mUserManager.setUserRestriction(
+                UserManager.DISALLOW_REMOVE_USER, /* value= */ true, mainUser);
+        try {
+            UserInfo privateProfileUser =
+                    createProfileForUser(
+                            "Private profile",
+                            UserManager.USER_TYPE_PROFILE_PRIVATE,
+                            mainUser.getIdentifier());
+            assertThat(privateProfileUser).isNotNull();
+            assertThat(hasUser(privateProfileUser.id)).isTrue();
+            removeUser(privateProfileUser.id);
+
+            assertThat(hasUser(privateProfileUser.id)).isFalse();
+        } finally {
+            mUserManager.setUserRestriction(
+                    UserManager.DISALLOW_REMOVE_USER, /* value= */ false, mainUser);
+        }
+    }
+
+    @MediumTest
+    @Test
+    public void testRemoveUser_withDisallowRemoveUserRestrictionAndMultipleUsersPresent() {
+        UserInfo privateProfileUser =
+                createProfileForUser(
+                        "Private profile",
+                        UserManager.USER_TYPE_PROFILE_PRIVATE,
+                        mUserManager.getMainUser().getIdentifier());
+        assertThat(privateProfileUser).isNotNull();
+        assertThat(hasUser(privateProfileUser.id)).isTrue();
+        UserInfo testUser = createUser("TestUser", /* flags= */ 0);
+        assertThat(testUser).isNotNull();
+        assertThat(hasUser(testUser.id)).isTrue();
+        UserHandle mainUser = mUserManager.getMainUser();
+        mUserManager.setUserRestriction(
+                UserManager.DISALLOW_REMOVE_USER, /* value= */ true, mainUser);
+        try {
+            assertThat(
+                    mUserManager.removeUserWhenPossible(
+                            testUser.getUserHandle(), /* overrideDevicePolicy= */ false))
+                    .isEqualTo(UserManager.REMOVE_RESULT_ERROR_USER_RESTRICTION);
+
+            // Non private profile users should be prevented from being removed.
+            assertThat(mUserManager.removeUser(testUser.id)).isEqualTo(false);
+
+            assertThat(hasUser(testUser.id)).isTrue();
+        } finally {
+            mUserManager.setUserRestriction(
+                    UserManager.DISALLOW_REMOVE_USER, /* value= */ false, mainUser);
+        }
+    }
+
+    @MediumTest
+    @Test
     public void testRemoveUserShouldNotRemoveTargetUser_DuringUserSwitch() {
         final int startUser = ActivityManager.getCurrentUser();
         final UserInfo testUser = createUser("TestUser", /* flags= */ 0);
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
index 48bc9d7..b5724b5c 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
@@ -21,8 +21,10 @@
 import static android.os.UserManager.USER_TYPE_PROFILE_CLONE;
 import static android.os.UserManager.USER_TYPE_PROFILE_MANAGED;
 import static android.os.UserManager.USER_TYPE_PROFILE_PRIVATE;
+import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
 import static android.service.notification.NotificationListenerService.META_DATA_DEFAULT_AUTOBIND;
 
+import static com.android.server.notification.Flags.FLAG_NOTIFICATION_NLS_REBIND;
 import static com.android.server.notification.ManagedServices.APPROVAL_BY_COMPONENT;
 import static com.android.server.notification.ManagedServices.APPROVAL_BY_PACKAGE;
 import static com.android.server.notification.NotificationManagerService.privateSpaceFlagsEnabled;
@@ -63,11 +65,14 @@
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.IInterface;
+import android.os.Looper;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
+import android.testing.TestableLooper;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.ArraySet;
@@ -82,7 +87,9 @@
 
 import com.google.android.collect.Lists;
 
+import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
@@ -103,7 +110,10 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 
+
 public class ManagedServicesTest extends UiServiceTestCase {
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
 
     @Mock
     private IPackageManager mIpm;
@@ -115,6 +125,7 @@
     private ManagedServices.UserProfiles mUserProfiles;
     @Mock private DevicePolicyManager mDpm;
     Object mLock = new Object();
+    private TestableLooper mTestableLooper;
 
     UserInfo mZero = new UserInfo(0, "zero", 0);
     UserInfo mTen = new UserInfo(10, "ten", 0);
@@ -142,6 +153,7 @@
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
+        mTestableLooper = new TestableLooper(Looper.getMainLooper());
 
         mContext.setMockPackageManager(mPm);
         mContext.addMockSystemService(Context.USER_SERVICE, mUm);
@@ -199,6 +211,11 @@
                 mIpm, APPROVAL_BY_COMPONENT);
     }
 
+    @After
+    public void tearDown() throws Exception {
+        mTestableLooper.destroy();
+    }
+
     @Test
     public void testBackupAndRestore_migration() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
@@ -888,7 +905,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a", 0, true);
 
         service.reregisterService(cn, 0);
@@ -919,7 +936,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a", 0, false);
 
         service.reregisterService(cn, 0);
@@ -950,7 +967,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a/a", 0, true);
 
         service.reregisterService(cn, 0);
@@ -981,7 +998,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a/a", 0, false);
 
         service.reregisterService(cn, 0);
@@ -1053,6 +1070,78 @@
     }
 
     @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void registerService_bindingDied_rebindIsClearedOnUserSwitch() throws Exception {
+        Context context = mock(Context.class);
+        PackageManager pm = mock(PackageManager.class);
+        ApplicationInfo ai = new ApplicationInfo();
+        ai.targetSdkVersion = Build.VERSION_CODES.CUR_DEVELOPMENT;
+
+        when(context.getPackageName()).thenReturn(mPkg);
+        when(context.getUserId()).thenReturn(mUser.getIdentifier());
+        when(context.getPackageManager()).thenReturn(pm);
+        when(pm.getApplicationInfo(anyString(), anyInt())).thenReturn(ai);
+
+        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles, mIpm,
+                APPROVAL_BY_PACKAGE);
+        service = spy(service);
+        ComponentName cn = ComponentName.unflattenFromString("a/a");
+
+        // Trigger onBindingDied for component when registering
+        //  => will schedule a rebind in 10 seconds
+        when(context.bindServiceAsUser(any(), any(), anyInt(), any())).thenAnswer(invocation -> {
+            Object[] args = invocation.getArguments();
+            ServiceConnection sc = (ServiceConnection) args[1];
+            sc.onBindingDied(cn);
+            return true;
+        });
+        service.registerService(cn, 0);
+        assertThat(service.isBound(cn, 0)).isFalse();
+
+        // Switch to user 10
+        service.onUserSwitched(10);
+
+        // Check that the scheduled rebind for user 0 was cleared
+        mTestableLooper.moveTimeForward(ManagedServices.ON_BINDING_DIED_REBIND_DELAY_MS);
+        mTestableLooper.processAllMessages();
+        verify(service, never()).reregisterService(any(), anyInt());
+    }
+
+    @Test
+    public void registerService_bindingDied_rebindIsExecutedAfterTimeout() throws Exception {
+        Context context = mock(Context.class);
+        PackageManager pm = mock(PackageManager.class);
+        ApplicationInfo ai = new ApplicationInfo();
+        ai.targetSdkVersion = Build.VERSION_CODES.CUR_DEVELOPMENT;
+
+        when(context.getPackageName()).thenReturn(mPkg);
+        when(context.getUserId()).thenReturn(mUser.getIdentifier());
+        when(context.getPackageManager()).thenReturn(pm);
+        when(pm.getApplicationInfo(anyString(), anyInt())).thenReturn(ai);
+
+        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles, mIpm,
+                APPROVAL_BY_PACKAGE);
+        service = spy(service);
+        ComponentName cn = ComponentName.unflattenFromString("a/a");
+
+        // Trigger onBindingDied for component when registering
+        //  => will schedule a rebind in 10 seconds
+        when(context.bindServiceAsUser(any(), any(), anyInt(), any())).thenAnswer(invocation -> {
+            Object[] args = invocation.getArguments();
+            ServiceConnection sc = (ServiceConnection) args[1];
+            sc.onBindingDied(cn);
+            return true;
+        });
+        service.registerService(cn, 0);
+        assertThat(service.isBound(cn, 0)).isFalse();
+
+        // Check that the scheduled rebind is run
+        mTestableLooper.moveTimeForward(ManagedServices.ON_BINDING_DIED_REBIND_DELAY_MS);
+        mTestableLooper.processAllMessages();
+        verify(service, times(1)).reregisterService(eq(cn), eq(0));
+    }
+
+    @Test
     public void testPackageUninstall_packageNoLongerInApprovedList() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1211,6 +1300,65 @@
     }
 
     @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void testUpgradeAppNoIntentFilterNoRebind() throws Exception {
+        Context context = spy(getContext());
+        doReturn(true).when(context).bindServiceAsUser(any(), any(), anyInt(), any());
+
+        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles,
+                mIpm, APPROVAL_BY_COMPONENT);
+
+        List<String> packages = new ArrayList<>();
+        packages.add("package");
+        addExpectedServices(service, packages, 0);
+
+        final ComponentName unapprovedComponent = ComponentName.unflattenFromString("package/C1");
+        final ComponentName approvedComponent = ComponentName.unflattenFromString("package/C2");
+
+        // Both components are approved initially
+        mExpectedPrimaryComponentNames.clear();
+        mExpectedPrimaryPackages.clear();
+        mExpectedPrimaryComponentNames.put(0, "package/C1:package/C2");
+        mExpectedSecondaryComponentNames.clear();
+        mExpectedSecondaryPackages.clear();
+
+        loadXml(service);
+
+        //Component package/C1 loses serviceInterface intent filter
+        ManagedServices.Config config = service.getConfig();
+        when(mPm.queryIntentServicesAsUser(any(), anyInt(), anyInt()))
+                .thenAnswer(new Answer<List<ResolveInfo>>() {
+                    @Override
+                    public List<ResolveInfo> answer(InvocationOnMock invocationOnMock)
+                            throws Throwable {
+                        Object[] args = invocationOnMock.getArguments();
+                        Intent invocationIntent = (Intent) args[0];
+                        if (invocationIntent != null) {
+                            if (invocationIntent.getAction().equals(config.serviceInterface)
+                                    && packages.contains(invocationIntent.getPackage())) {
+                                List<ResolveInfo> dummyServices = new ArrayList<>();
+                                ResolveInfo resolveInfo = new ResolveInfo();
+                                ServiceInfo serviceInfo = new ServiceInfo();
+                                serviceInfo.packageName = invocationIntent.getPackage();
+                                serviceInfo.name = approvedComponent.getClassName();
+                                serviceInfo.permission = service.getConfig().bindPermission;
+                                resolveInfo.serviceInfo = serviceInfo;
+                                dummyServices.add(resolveInfo);
+                                return dummyServices;
+                            }
+                        }
+                        return new ArrayList<>();
+                    }
+                });
+
+        // Trigger package update
+        service.onPackagesChanged(false, new String[]{"package"}, new int[]{0});
+
+        assertFalse(service.isComponentEnabledForCurrentProfiles(unapprovedComponent));
+        assertTrue(service.isComponentEnabledForCurrentProfiles(approvedComponent));
+    }
+
+    @Test
     public void testSetPackageOrComponentEnabled() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1223,6 +1371,21 @@
                             "user10package1/K", "user10.3/Component", "user10package2/L",
                             "user10.4/Component"}));
 
+            // mock permissions for services
+            PackageManager pm = mock(PackageManager.class);
+            when(getContext().getPackageManager()).thenReturn(pm);
+            List<ComponentName> enabledComponents = List.of(
+                    ComponentName.unflattenFromString("package/Comp"),
+                    ComponentName.unflattenFromString("package/C2"),
+                    ComponentName.unflattenFromString("again/M4"),
+                    ComponentName.unflattenFromString("user10package/B"),
+                    ComponentName.unflattenFromString("user10/Component"),
+                    ComponentName.unflattenFromString("user10package1/K"),
+                    ComponentName.unflattenFromString("user10.3/Component"),
+                    ComponentName.unflattenFromString("user10package2/L"),
+                    ComponentName.unflattenFromString("user10.4/Component"));
+            mockServiceInfoWithMetaData(enabledComponents, service, pm, new ArrayMap<>());
+
             for (int userId : expectedEnabled.keySet()) {
                 ArrayList<String> expectedForUser = expectedEnabled.get(userId);
                 for (int i = 0; i < expectedForUser.size(); i++) {
@@ -1284,6 +1447,90 @@
     }
 
     @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void testSetPackageOrComponentEnabled_pkgInstalledAfterEnabling() throws Exception {
+        ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
+                mIpm, APPROVAL_BY_COMPONENT);
+
+        final int userId = 0;
+        final String validComponent = "again/M4";
+        ArrayList<String> expectedEnabled = Lists.newArrayList("package/Comp", "package/C2",
+                validComponent);
+
+        PackageManager pm = mock(PackageManager.class);
+        when(getContext().getPackageManager()).thenReturn(pm);
+        service = spy(service);
+
+        // Component again/M4 is a valid service and the package is available
+        doReturn(true).when(service)
+                .isValidService(ComponentName.unflattenFromString(validComponent), userId);
+        when(pm.isPackageAvailable("again")).thenReturn(true);
+
+        // "package" is not available and its services are not valid
+        doReturn(false).when(service)
+                .isValidService(ComponentName.unflattenFromString("package/Comp"), userId);
+        doReturn(false).when(service)
+                .isValidService(ComponentName.unflattenFromString("package/C2"), userId);
+        when(pm.isPackageAvailable("package")).thenReturn(false);
+
+        // Enable all components
+        for (String component: expectedEnabled) {
+            service.setPackageOrComponentEnabled(component, userId, true, true);
+        }
+
+        // Verify everything added is approved
+        for (String component: expectedEnabled) {
+            assertTrue("Not allowed: user: " + userId + " entry: " + component
+                    + " for approval level " + APPROVAL_BY_COMPONENT,
+                    service.isPackageOrComponentAllowed(component, userId));
+        }
+
+        // Add missing package "package"
+        service.onPackagesChanged(false, new String[]{"package"}, new int[]{0});
+
+        // Check that component of "package" are not enabled
+        assertFalse(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString("package/Comp")));
+        assertFalse(service.isPackageOrComponentAllowed("package/Comp", userId));
+
+        assertFalse(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString("package/C2")));
+        assertFalse(service.isPackageOrComponentAllowed("package/C2", userId));
+
+        // Check that the valid components are still enabled
+        assertTrue(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString(validComponent)));
+        assertTrue(service.isPackageOrComponentAllowed(validComponent, userId));
+    }
+
+    @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void testSetPackageOrComponentEnabled_invalidComponent() throws Exception {
+        ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
+                mIpm, APPROVAL_BY_COMPONENT);
+
+        final int userId = 0;
+        final String invalidComponent = "package/Comp";
+
+        PackageManager pm = mock(PackageManager.class);
+        when(getContext().getPackageManager()).thenReturn(pm);
+        service = spy(service);
+
+        // Component is an invalid service and the package is available
+        doReturn(false).when(service)
+                .isValidService(ComponentName.unflattenFromString(invalidComponent), userId);
+        when(pm.isPackageAvailable("package")).thenReturn(true);
+        service.setPackageOrComponentEnabled(invalidComponent, userId, true, true);
+
+        // Verify that the component was not enabled
+        assertFalse("Not allowed: user: " + userId + " entry: " + invalidComponent
+                    + " for approval level " + APPROVAL_BY_COMPONENT,
+                service.isPackageOrComponentAllowed(invalidComponent, userId));
+        assertFalse(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString(invalidComponent)));
+    }
+
+    @Test
     public void testGetAllowedPackages_byUser() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1944,7 +2191,7 @@
         metaDataAutobindAllow.putBoolean(META_DATA_DEFAULT_AUTOBIND, true);
         metaDatas.put(cn_allowed, metaDataAutobindAllow);
 
-        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
 
         service.addApprovedList(cn_allowed.flattenToString(), 0, true);
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
@@ -1989,7 +2236,7 @@
         metaDataAutobindDisallow.putBoolean(META_DATA_DEFAULT_AUTOBIND, false);
         metaDatas.put(cn_disallowed, metaDataAutobindDisallow);
 
-        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
 
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
 
@@ -2028,7 +2275,7 @@
         metaDataAutobindDisallow.putBoolean(META_DATA_DEFAULT_AUTOBIND, false);
         metaDatas.put(cn_disallowed, metaDataAutobindDisallow);
 
-        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
 
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
 
@@ -2099,8 +2346,8 @@
     }
 
     private void mockServiceInfoWithMetaData(List<ComponentName> componentNames,
-            ManagedServices service, ArrayMap<ComponentName, Bundle> metaDatas)
-            throws RemoteException {
+            ManagedServices service, PackageManager packageManager,
+            ArrayMap<ComponentName, Bundle> metaDatas) throws RemoteException {
         when(mIpm.getServiceInfo(any(), anyLong(), anyInt())).thenAnswer(
                 (Answer<ServiceInfo>) invocation -> {
                     ComponentName invocationCn = invocation.getArgument(0);
@@ -2115,6 +2362,39 @@
                     return null;
                 }
         );
+
+        // add components to queryIntentServicesAsUser response
+        final List<String> packages = new ArrayList<>();
+        for (ComponentName cn: componentNames) {
+            packages.add(cn.getPackageName());
+        }
+        ManagedServices.Config config = service.getConfig();
+        when(packageManager.queryIntentServicesAsUser(any(), anyInt(), anyInt())).
+                thenAnswer(new Answer<List<ResolveInfo>>() {
+                @Override
+                public List<ResolveInfo> answer(InvocationOnMock invocationOnMock)
+                    throws Throwable {
+                    Object[] args = invocationOnMock.getArguments();
+                    Intent invocationIntent = (Intent) args[0];
+                    if (invocationIntent != null) {
+                        if (invocationIntent.getAction().equals(config.serviceInterface)
+                            && packages.contains(invocationIntent.getPackage())) {
+                            List<ResolveInfo> dummyServices = new ArrayList<>();
+                            for (ComponentName cn: componentNames) {
+                                ResolveInfo resolveInfo = new ResolveInfo();
+                                ServiceInfo serviceInfo = new ServiceInfo();
+                                serviceInfo.packageName = invocationIntent.getPackage();
+                                serviceInfo.name = cn.getClassName();
+                                serviceInfo.permission = service.getConfig().bindPermission;
+                                resolveInfo.serviceInfo = serviceInfo;
+                                dummyServices.add(resolveInfo);
+                            }
+                            return dummyServices;
+                        }
+                    }
+                    return new ArrayList<>();
+                }
+            });
     }
 
     private void resetComponentsAndPackages() {
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
index 0f7de7d..2c645e0 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
@@ -28,6 +28,7 @@
 import static junit.framework.Assert.assertTrue;
 
 import static org.junit.Assert.assertNull;
+
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.any;
@@ -197,6 +198,8 @@
     public void testWriteXml_userTurnedOffNAS() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
+
         mAssistants.loadDefaultsFromConfig(true);
 
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
@@ -432,6 +435,10 @@
     public void testSetPackageOrComponentEnabled_onlyOnePackage() throws Exception {
         ComponentName component1 = ComponentName.unflattenFromString("package/Component1");
         ComponentName component2 = ComponentName.unflattenFromString("package/Component2");
+
+        doReturn(true).when(mAssistants).isValidService(eq(component1), eq(mZero.id));
+        doReturn(true).when(mAssistants).isValidService(eq(component2), eq(mZero.id));
+
         mAssistants.setPackageOrComponentEnabled(component1.flattenToString(), mZero.id, true,
                 true, true);
         verify(mNm, never()).setNotificationAssistantAccessGrantedForUserInternal(
@@ -577,6 +584,7 @@
     public void testSetAdjustmentTypeSupportedState() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
@@ -600,6 +608,7 @@
     public void testSetAdjustmentTypeSupportedState_readWriteXml_entries() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
@@ -623,6 +632,7 @@
     public void testSetAdjustmentTypeSupportedState_readWriteXml_empty() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index 577b02a..c30b4bb 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -3214,23 +3214,32 @@
         assertFalse(activity.mDisplayContent.mClosingApps.contains(activity));
     }
 
+    @SetupWindows(addWindows = W_ACTIVITY)
     @Test
     public void testSetVisibility_visibleToInvisible() {
-        final ActivityRecord activity = new ActivityBuilder(mAtm)
-                .setCreateTask(true).build();
+        final TestTransitionPlayer player = registerTestTransitionPlayer();
+        final ActivityRecord activity = mAppWindow.mActivityRecord;
+        makeWindowVisibleAndDrawn(mAppWindow);
         // By default, activity is visible.
         assertTrue(activity.isVisible());
         assertTrue(activity.isVisibleRequested());
-        assertFalse(activity.mDisplayContent.mClosingApps.contains(activity));
+        assertTrue(mAppWindow.isDrawn());
+        assertFalse(mAppWindow.setReportResizeHints());
 
         // Request the activity to be invisible. Since the visibility changes, app transition
         // animation should be applied on this activity.
-        mDisplayContent.prepareAppTransition(0);
+        activity.mTransitionController.requestCloseTransitionIfNeeded(activity);
         activity.setVisibility(false);
         assertTrue(activity.isVisible());
         assertFalse(activity.isVisibleRequested());
-        assertFalse(activity.mDisplayContent.mOpeningApps.contains(activity));
-        assertTrue(activity.mDisplayContent.mClosingApps.contains(activity));
+
+        player.start();
+        mSetFlagsRule.enableFlags(Flags.FLAG_RESET_DRAW_STATE_ON_CLIENT_INVISIBLE);
+        // ActivityRecord#commitVisibility(false) -> WindowState#sendAppVisibilityToClients().
+        player.finish();
+        assertFalse(activity.isVisible());
+        assertFalse("Reset draw state after committing invisible", mAppWindow.isDrawn());
+        assertTrue("Set pending redraw hint", mAppWindow.setReportResizeHints());
     }
 
     @Test
diff --git a/services/usage/OWNERS b/services/usage/OWNERS
index f825f55..678c7ac 100644
--- a/services/usage/OWNERS
+++ b/services/usage/OWNERS
@@ -3,7 +3,6 @@
 
 mwachens@google.com
 varunshah@google.com
-yamasani@google.com
 guanxin@google.com
 
 per-file *StorageStats* = file:/core/java/android/os/storage/OWNERS
diff --git a/telephony/java/android/telephony/CarrierConfigManager.java b/telephony/java/android/telephony/CarrierConfigManager.java
index 02999c8..6535b9b 100644
--- a/telephony/java/android/telephony/CarrierConfigManager.java
+++ b/telephony/java/android/telephony/CarrierConfigManager.java
@@ -9856,6 +9856,16 @@
     public static final String KEY_REMOVE_SATELLITE_PLMN_IN_MANUAL_NETWORK_SCAN_BOOL =
             "remove_satellite_plmn_in_manual_network_scan_bool";
 
+    /**
+     * This value is used to set the max datagram size, if the value is not available then the
+     * default one will be used.
+     * If key is {@code true}, retrieve the max datagram value and use this value always,
+     * {@code false} the default value from the modem will be used.
+     *
+     * @hide
+     */
+    public static final String KEY_SATELLITE_SOS_MAX_DATAGRAM_SIZE =
+            "satellite_sos_max_datagram_size";
 
     /** @hide */
     @IntDef({
diff --git a/telephony/java/android/telephony/satellite/SatelliteCapabilities.java b/telephony/java/android/telephony/satellite/SatelliteCapabilities.java
index 0d8f101..f34522a 100644
--- a/telephony/java/android/telephony/satellite/SatelliteCapabilities.java
+++ b/telephony/java/android/telephony/satellite/SatelliteCapabilities.java
@@ -200,6 +200,15 @@
     }
 
     /**
+     * Setting the maximum number of bytes per datagram that can be sent over satellite.
+     *
+     * @hide
+     */
+    public void setMaxBytesPerOutgoingDatagram(int maxBytesPerOutgoingDatagram) {
+        mMaxBytesPerOutgoingDatagram = maxBytesPerOutgoingDatagram;
+    }
+
+    /**
      * Antenna Position received from satellite modem which gives information about antenna
      * direction to be used with satellite communication and suggested device hold positions.
      * @return Map key: {@link SatelliteManager.DeviceHoldPosition} value: AntennaPosition
diff --git a/tests/JobSchedulerPerfTests/OWNERS b/tests/JobSchedulerPerfTests/OWNERS
index 6f207fb1..c8345f7 100644
--- a/tests/JobSchedulerPerfTests/OWNERS
+++ b/tests/JobSchedulerPerfTests/OWNERS
@@ -1 +1 @@
-include /apex/jobscheduler/OWNERS
+include /apex/jobscheduler/JOB_OWNERS
diff --git a/tests/JobSchedulerTestApp/OWNERS b/tests/JobSchedulerTestApp/OWNERS
index 6f207fb1..c8345f7 100644
--- a/tests/JobSchedulerTestApp/OWNERS
+++ b/tests/JobSchedulerTestApp/OWNERS
@@ -1 +1 @@
-include /apex/jobscheduler/OWNERS
+include /apex/jobscheduler/JOB_OWNERS
diff --git a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
index e29e462..e045f10 100644
--- a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
+++ b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
@@ -224,7 +224,6 @@
         doReturn(mFeatureFlags).when(mVcnContext).getFeatureFlags();
         doReturn(true).when(mVcnContext).isFlagSafeModeTimeoutConfigEnabled();
         doReturn(true).when(mVcnContext).isFlagIpSecTransformStateEnabled();
-        doReturn(true).when(mVcnContext).isFlagNetworkMetricMonitorEnabled();
 
         doReturn(mUnderlyingNetworkController)
                 .when(mDeps)
diff --git a/tests/vcn/java/com/android/server/vcn/routeselection/NetworkEvaluationTestBase.java b/tests/vcn/java/com/android/server/vcn/routeselection/NetworkEvaluationTestBase.java
index 421e1ad..bc7ff47 100644
--- a/tests/vcn/java/com/android/server/vcn/routeselection/NetworkEvaluationTestBase.java
+++ b/tests/vcn/java/com/android/server/vcn/routeselection/NetworkEvaluationTestBase.java
@@ -127,7 +127,6 @@
                                 false /* isInTestMode */));
         doNothing().when(mVcnContext).ensureRunningOnLooperThread();
 
-        doReturn(true).when(mVcnContext).isFlagNetworkMetricMonitorEnabled();
         doReturn(true).when(mVcnContext).isFlagIpSecTransformStateEnabled();
 
         setupSystemService(
diff --git a/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java b/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java
index 588624b..6f31d8d 100644
--- a/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java
+++ b/tests/vcn/java/com/android/server/vcn/routeselection/UnderlyingNetworkControllerTest.java
@@ -226,7 +226,6 @@
     private void resetVcnContext(VcnContext vcnContext) {
         reset(vcnContext);
         doNothing().when(vcnContext).ensureRunningOnLooperThread();
-        doReturn(true).when(vcnContext).isFlagNetworkMetricMonitorEnabled();
         doReturn(true).when(vcnContext).isFlagIpSecTransformStateEnabled();
     }