Merge "Update sign-in titles" into main
diff --git a/core/api/current.txt b/core/api/current.txt
index c189a24c..c86e4cd 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -282,7 +282,7 @@
     field public static final String REQUEST_IGNORE_BATTERY_OPTIMIZATIONS = "android.permission.REQUEST_IGNORE_BATTERY_OPTIMIZATIONS";
     field public static final String REQUEST_INSTALL_PACKAGES = "android.permission.REQUEST_INSTALL_PACKAGES";
     field public static final String REQUEST_OBSERVE_COMPANION_DEVICE_PRESENCE = "android.permission.REQUEST_OBSERVE_COMPANION_DEVICE_PRESENCE";
-    field @FlaggedApi("android.companion.flags.device_presence") public static final String REQUEST_OBSERVE_DEVICE_UUID_PRESENCE = "android.permission.REQUEST_OBSERVE_DEVICE_UUID_PRESENCE";
+    field @FlaggedApi("android.companion.device_presence") public static final String REQUEST_OBSERVE_DEVICE_UUID_PRESENCE = "android.permission.REQUEST_OBSERVE_DEVICE_UUID_PRESENCE";
     field public static final String REQUEST_PASSWORD_COMPLEXITY = "android.permission.REQUEST_PASSWORD_COMPLEXITY";
     field @Deprecated public static final String RESTART_PACKAGES = "android.permission.RESTART_PACKAGES";
     field public static final String RUN_USER_INITIATED_JOBS = "android.permission.RUN_USER_INITIATED_JOBS";
@@ -39388,10 +39388,8 @@
   }
 
   public final class FileIntegrityManager {
-    method @FlaggedApi("android.security.fsverity_api") @Nullable public byte[] getFsVerityDigest(@NonNull java.io.File) throws java.io.IOException;
     method public boolean isApkVeritySupported();
     method @Deprecated @RequiresPermission(anyOf={android.Manifest.permission.INSTALL_PACKAGES, android.Manifest.permission.REQUEST_INSTALL_PACKAGES}) public boolean isAppSourceCertificateTrusted(@NonNull java.security.cert.X509Certificate) throws java.security.cert.CertificateEncodingException;
-    method @FlaggedApi("android.security.fsverity_api") public void setupFsVerity(@NonNull java.io.File) throws java.io.IOException;
   }
 
   public final class KeyChain {
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index b767c52..bed8c41 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -355,6 +355,7 @@
     field public static final String SEND_SHOW_SUSPENDED_APP_DETAILS = "android.permission.SEND_SHOW_SUSPENDED_APP_DETAILS";
     field public static final String SEND_SMS_NO_CONFIRMATION = "android.permission.SEND_SMS_NO_CONFIRMATION";
     field public static final String SERIAL_PORT = "android.permission.SERIAL_PORT";
+    field @FlaggedApi("android.security.fsverity_api") public static final String SETUP_FSVERITY = "android.permission.SETUP_FSVERITY";
     field public static final String SET_ACTIVITY_WATCHER = "android.permission.SET_ACTIVITY_WATCHER";
     field public static final String SET_CLIP_SOURCE = "android.permission.SET_CLIP_SOURCE";
     field public static final String SET_DEFAULT_ACCOUNT_FOR_CONTACTS = "android.permission.SET_DEFAULT_ACCOUNT_FOR_CONTACTS";
@@ -1144,6 +1145,7 @@
   public static final class StatusBarManager.DisableInfo implements android.os.Parcelable {
     method public boolean areAllComponentsEnabled();
     method public int describeContents();
+    method public boolean isBackDisabled();
     method public boolean isNavigateToHomeDisabled();
     method public boolean isNotificationPeekingDisabled();
     method public boolean isRecentsDisabled();
@@ -12106,6 +12108,11 @@
 
 package android.security {
 
+  public final class FileIntegrityManager {
+    method @FlaggedApi("android.security.fsverity_api") @Nullable public byte[] getFsVerityDigest(@NonNull java.io.File) throws java.io.IOException;
+    method @FlaggedApi("android.security.fsverity_api") public void setupFsVerity(@NonNull java.io.File) throws java.io.IOException;
+  }
+
   public final class KeyChain {
     method @Nullable @WorkerThread public static String getWifiKeyGrantAsUser(@NonNull android.content.Context, @NonNull android.os.UserHandle, @NonNull String);
     method @WorkerThread public static boolean hasWifiKeyGrantAsUser(@NonNull android.content.Context, @NonNull android.os.UserHandle, @NonNull String);
diff --git a/core/java/android/app/StatusBarManager.java b/core/java/android/app/StatusBarManager.java
index 301fef8..e6e46dd 100644
--- a/core/java/android/app/StatusBarManager.java
+++ b/core/java/android/app/StatusBarManager.java
@@ -1448,6 +1448,7 @@
          *
          * @hide
          */
+        @SystemApi
         public boolean isBackDisabled() {
             return mBack;
         }
@@ -1861,38 +1862,38 @@
         };
 
         @DataClass.Generated(
-                time = 1708625947132L,
+                time = 1707345957771L,
                 codegenVersion = "1.0.23",
                 sourceFile = "frameworks/base/core/java/android/app/StatusBarManager.java",
-                inputSignatures = "private  boolean mStatusBarExpansion\nprivate  boolean "
-                        + "mNavigateHome\nprivate  boolean mNotificationPeeking\nprivate  "
-                        + "boolean mRecents\nprivate  boolean mBack\nprivate  boolean mSearch\n"
-                        + "private  boolean mSystemIcons\nprivate  boolean mClock\nprivate  "
-                        + "boolean mNotificationIcons\nprivate  boolean mRotationSuggestion\n"
+                inputSignatures = "private  boolean mStatusBarExpansion\nprivate  "
+                        + "boolean mNavigateHome\nprivate  boolean mNotificationPeeking\nprivate  "
+                        + "boolean mRecents\nprivate  boolean mBack\nprivate  boolean "
+                        + "mSearch\nprivate  boolean mSystemIcons\nprivate  boolean mClock\nprivate"
+                        + "  boolean mNotificationIcons\nprivate  boolean mRotationSuggestion\n"
                         + "private  boolean mNotificationTicker\npublic "
                         + "@android.annotation.SystemApi boolean isStatusBarExpansionDisabled()\n"
                         + "public  void setStatusBarExpansionDisabled(boolean)\npublic "
-                        + "@android.annotation.SystemApi boolean isNavigateToHomeDisabled()\npublic"
-                        + "  void setNavigationHomeDisabled(boolean)\npublic "
-                        + "@android.annotation.SystemApi boolean isNotificationPeekingDisabled()"
-                        + "\npublic  void setNotificationPeekingDisabled(boolean)\npublic "
+                        + "@android.annotation.SystemApi boolean isNavigateToHomeDisabled()\n"
+                        + "public  void setNavigationHomeDisabled(boolean)\npublic "
+                        + "@android.annotation.SystemApi boolean isNotificationPeekingDisabled()\n"
+                        + "public  void setNotificationPeekingDisabled(boolean)\npublic "
                         + "@android.annotation.SystemApi boolean isRecentsDisabled()\npublic  "
-                        + "void setRecentsDisabled(boolean)\npublic  boolean isBackDisabled()"
-                        + "\npublic  void setBackDisabled(boolean)\npublic "
+                        + "void setRecentsDisabled(boolean)\npublic @android.annotation.SystemApi "
+                        + "boolean isBackDisabled()\npublic  void setBackDisabled(boolean)\npublic "
                         + "@android.annotation.SystemApi boolean isSearchDisabled()\npublic  "
                         + "void setSearchDisabled(boolean)\npublic  boolean "
-                        + "areSystemIconsDisabled()\npublic  void setSystemIconsDisabled(boolean)\n"
-                        + "public  boolean isClockDisabled()\npublic  "
-                        + "void setClockDisabled(boolean)\npublic  boolean "
-                        + "areNotificationIconsDisabled()\npublic  void "
-                        + "setNotificationIconsDisabled(boolean)\npublic  boolean "
+                        + "areSystemIconsDisabled()\npublic  void setSystemIconsDisabled(boolean)"
+                        + "\npublic  boolean isClockDisabled()\npublic  "
+                        + "void setClockDisabled(boolean)\npublic  "
+                        + "boolean areNotificationIconsDisabled()\npublic  "
+                        + "void setNotificationIconsDisabled(boolean)\npublic  boolean "
                         + "isNotificationTickerDisabled()\npublic  void "
                         + "setNotificationTickerDisabled(boolean)\npublic "
                         + "@android.annotation.TestApi boolean isRotationSuggestionDisabled()\n"
                         + "public  void setRotationSuggestionDisabled(boolean)\npublic "
-                        + "@android.annotation.SystemApi boolean areAllComponentsEnabled()\npublic"
-                        + "  void setEnableAll()\npublic  boolean areAllComponentsDisabled()\n"
-                        + "public  void setDisableAll()\npublic @android.annotation.NonNull "
+                        + "@android.annotation.SystemApi boolean areAllComponentsEnabled()\n"
+                        + "public  void setEnableAll()\npublic  boolean areAllComponentsDisabled()"
+                        + "\npublic  void setDisableAll()\npublic @android.annotation.NonNull "
                         + "@java.lang.Override java.lang.String toString()\npublic  "
                         + "android.util.Pair<java.lang.Integer,java.lang.Integer> toFlags()\n"
                         + "class DisableInfo extends java.lang.Object implements "
diff --git a/core/java/android/security/FileIntegrityManager.java b/core/java/android/security/FileIntegrityManager.java
index 025aac9..478435b 100644
--- a/core/java/android/security/FileIntegrityManager.java
+++ b/core/java/android/security/FileIntegrityManager.java
@@ -20,6 +20,8 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
+import android.annotation.SuppressLint;
+import android.annotation.SystemApi;
 import android.annotation.SystemService;
 import android.content.Context;
 import android.os.IInstalld.IFsveritySetupAuthToken;
@@ -99,8 +101,11 @@
      * @throws IOException If the operation failed.
      *
      * @see <a href="https://www.kernel.org/doc/html/next/filesystems/fsverity.html">Kernel doc</a>
+     * @hide
      */
     @FlaggedApi(Flags.FLAG_FSVERITY_API)
+    @SuppressLint("StreamFiles")
+    @SystemApi
     public void setupFsVerity(@NonNull File file) throws IOException {
         if (!file.isAbsolute()) {
             // fs-verity is to be enabled by installd, which enforces the validation to the
@@ -138,8 +143,11 @@
      * @param file The file to measure the fs-verity digest.
      * @return The fs-verity digest in byte[], null if none.
      * @see <a href="https://www.kernel.org/doc/html/next/filesystems/fsverity.html">Kernel doc</a>
+     * @hide
      */
     @FlaggedApi(Flags.FLAG_FSVERITY_API)
+    @SuppressLint("StreamFiles")
+    @SystemApi
     public @Nullable byte[] getFsVerityDigest(@NonNull File file) throws IOException {
         return VerityUtils.getFsverityDigest(file.getPath());
     }
diff --git a/core/java/android/tracing/perfetto/DataSource.java b/core/java/android/tracing/perfetto/DataSource.java
index c33fa7d..bbf378e 100644
--- a/core/java/android/tracing/perfetto/DataSource.java
+++ b/core/java/android/tracing/perfetto/DataSource.java
@@ -18,8 +18,6 @@
 
 import android.util.proto.ProtoInputStream;
 
-import com.android.internal.annotations.VisibleForTesting;
-
 import dalvik.annotation.optimization.CriticalNative;
 
 /**
@@ -73,7 +71,8 @@
      * @param fun The tracing lambda that will be called with the tracing contexts of each active
      *            tracing instance.
      */
-    public final void trace(TraceFunction<TlsStateType, IncrementalStateType> fun) {
+    public final void trace(
+            TraceFunction<DataSourceInstanceType, TlsStateType, IncrementalStateType> fun) {
         boolean startedIterator = nativePerfettoDsTraceIterateBegin(mNativeObj);
 
         if (!startedIterator) {
@@ -82,8 +81,10 @@
 
         try {
             do {
-                TracingContext<TlsStateType, IncrementalStateType> ctx =
-                        new TracingContext<>(mNativeObj);
+                int instanceIndex = nativeGetPerfettoDsInstanceIndex(mNativeObj);
+
+                TracingContext<DataSourceInstanceType, TlsStateType, IncrementalStateType> ctx =
+                        new TracingContext<>(this, instanceIndex);
                 fun.trace(ctx);
 
                 ctx.flush();
@@ -104,9 +105,7 @@
      * Override this method to create a custom TlsState object for your DataSource. A new instance
      * will be created per trace instance per thread.
      *
-     * NOTE: Should only be called from native side.
      */
-    @VisibleForTesting
     public TlsStateType createTlsState(CreateTlsStateArgs<DataSourceInstanceType> args) {
         return null;
     }
@@ -114,9 +113,8 @@
     /**
      * Override this method to create and use a custom IncrementalState object for your DataSource.
      *
-     * NOTE: Should only be called from native side.
      */
-    protected IncrementalStateType createIncrementalState(
+    public IncrementalStateType createIncrementalState(
             CreateIncrementalStateArgs<DataSourceInstanceType> args) {
         return null;
     }
@@ -185,4 +183,6 @@
     private static native boolean nativePerfettoDsTraceIterateNext(long dataSourcePtr);
     @CriticalNative
     private static native void nativePerfettoDsTraceIterateBreak(long dataSourcePtr);
+    @CriticalNative
+    private static native int nativeGetPerfettoDsInstanceIndex(long dataSourcePtr);
 }
diff --git a/core/java/android/tracing/perfetto/TraceFunction.java b/core/java/android/tracing/perfetto/TraceFunction.java
index 13e663d..d8854f9 100644
--- a/core/java/android/tracing/perfetto/TraceFunction.java
+++ b/core/java/android/tracing/perfetto/TraceFunction.java
@@ -19,12 +19,14 @@
 /**
  * The interface for the trace function called from native on a trace call with a context.
  *
+ * @param <DataSourceInstanceType> The type of DataSource this tracing context is for.
  * @param <TlsStateType> The type of the custom TLS state, if any is used.
  * @param <IncrementalStateType> The type of the custom incremental state, if any is used.
  *
  * @hide
  */
-public interface TraceFunction<TlsStateType, IncrementalStateType> {
+public interface TraceFunction<DataSourceInstanceType extends DataSourceInstance,
+        TlsStateType, IncrementalStateType> {
 
     /**
      * This function will be called synchronously (i.e., always before trace() returns) only if
@@ -34,5 +36,5 @@
      *
      * @param ctx the tracing context to trace for in the trace function.
      */
-    void trace(TracingContext<TlsStateType, IncrementalStateType> ctx);
+    void trace(TracingContext<DataSourceInstanceType, TlsStateType, IncrementalStateType> ctx);
 }
diff --git a/core/java/android/tracing/perfetto/TracingContext.java b/core/java/android/tracing/perfetto/TracingContext.java
index 060f964..6b7df54 100644
--- a/core/java/android/tracing/perfetto/TracingContext.java
+++ b/core/java/android/tracing/perfetto/TracingContext.java
@@ -24,18 +24,25 @@
 /**
  * Argument passed to the lambda function passed to Trace().
  *
+ * @param <DataSourceInstanceType> The type of the datasource this tracing context is for.
  * @param <TlsStateType> The type of the custom TLS state, if any is used.
  * @param <IncrementalStateType> The type of the custom incremental state, if any is used.
  *
  * @hide
  */
-public class TracingContext<TlsStateType, IncrementalStateType> {
+public class TracingContext<DataSourceInstanceType extends DataSourceInstance, TlsStateType,
+        IncrementalStateType> {
 
-    private final long mNativeDsPtr;
+    private final DataSource<DataSourceInstanceType, TlsStateType, IncrementalStateType>
+            mDataSource;
+    private final int mInstanceIndex;
     private final List<ProtoOutputStream> mTracePackets = new ArrayList<>();
 
-    TracingContext(long nativeDsPtr) {
-        this.mNativeDsPtr = nativeDsPtr;
+    TracingContext(DataSource<DataSourceInstanceType, TlsStateType, IncrementalStateType>
+            dataSource,
+            int instanceIndex) {
+        this.mDataSource = dataSource;
+        this.mInstanceIndex = instanceIndex;
     }
 
     /**
@@ -61,18 +68,26 @@
      * Stop timeout expires.
      */
     public void flush() {
-        nativeFlush(mNativeDsPtr, getAndClearAllPendingTracePackets());
+        nativeFlush(mDataSource.mNativeObj, getAndClearAllPendingTracePackets());
     }
 
     /**
      * Can optionally be used to store custom per-sequence
      * session data, which is not reset when incremental state is cleared
      * (e.g. configuration options).
-     *
+     *h
      * @return The TlsState instance for the tracing thread and instance.
      */
     public TlsStateType getCustomTlsState() {
-        return (TlsStateType) nativeGetCustomTls(mNativeDsPtr);
+        TlsStateType tlsState = (TlsStateType) nativeGetCustomTls(mDataSource.mNativeObj);
+        if (tlsState == null) {
+            final CreateTlsStateArgs<DataSourceInstanceType> args =
+                    new CreateTlsStateArgs<>(mDataSource, mInstanceIndex);
+            tlsState = mDataSource.createTlsState(args);
+            nativeSetCustomTls(mDataSource.mNativeObj, tlsState);
+        }
+
+        return tlsState;
     }
 
     /**
@@ -82,7 +97,16 @@
      * @return The current IncrementalState object instance.
      */
     public IncrementalStateType getIncrementalState() {
-        return (IncrementalStateType) nativeGetIncrementalState(mNativeDsPtr);
+        IncrementalStateType incrementalState =
+                (IncrementalStateType) nativeGetIncrementalState(mDataSource.mNativeObj);
+        if (incrementalState == null) {
+            final CreateIncrementalStateArgs<DataSourceInstanceType> args =
+                    new CreateIncrementalStateArgs<>(mDataSource, mInstanceIndex);
+            incrementalState = mDataSource.createIncrementalState(args);
+            nativeSetIncrementalState(mDataSource.mNativeObj, incrementalState);
+        }
+
+        return incrementalState;
     }
 
     private byte[][] getAndClearAllPendingTracePackets() {
@@ -97,6 +121,10 @@
     }
 
     private static native void nativeFlush(long dataSourcePtr, byte[][] packetData);
+
     private static native Object nativeGetCustomTls(long nativeDsPtr);
+    private static native void nativeSetCustomTls(long nativeDsPtr, Object tlsState);
+
     private static native Object nativeGetIncrementalState(long nativeDsPtr);
+    private static native void nativeSetIncrementalState(long nativeDsPtr, Object incrementalState);
 }
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 87dfa029..09f98d7 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -239,6 +239,7 @@
 import android.view.autofill.AutofillManager;
 import android.view.contentcapture.ContentCaptureManager;
 import android.view.contentcapture.ContentCaptureSession;
+import android.view.flags.Flags;
 import android.view.inputmethod.ImeTracker;
 import android.view.inputmethod.InputMethodManager;
 import android.widget.Scroller;
@@ -1153,7 +1154,9 @@
     private static boolean sToolkitFrameRateTypingReadOnlyFlagValue;
     private static final boolean sToolkitFrameRateViewEnablingReadOnlyFlagValue;
     private static boolean sToolkitFrameRateVelocityMappingReadOnlyFlagValue =
-            toolkitFrameRateVelocityMappingReadOnly();;
+            toolkitFrameRateVelocityMappingReadOnly();
+    private static boolean sToolkitEnableInvalidateCheckThreadFlagValue =
+            Flags.enableInvalidateCheckThread();
 
     static {
         sToolkitSetFrameRateReadOnlyFlagValue = toolkitSetFrameRateReadOnly();
@@ -2375,8 +2378,9 @@
 
     @Override
     public void onDescendantInvalidated(@NonNull View child, @NonNull View descendant) {
-        // TODO: Re-enable after camera is fixed or consider targetSdk checking this
-        // checkThread();
+        if (sToolkitEnableInvalidateCheckThreadFlagValue) {
+            checkThread();
+        }
         if ((descendant.mPrivateFlags & PFLAG_DRAW_ANIMATION) != 0) {
             mIsAnimating = true;
         }
diff --git a/core/java/android/view/flags/view_flags.aconfig b/core/java/android/view/flags/view_flags.aconfig
index c482f8b..486c2ab 100644
--- a/core/java/android/view/flags/view_flags.aconfig
+++ b/core/java/android/view/flags/view_flags.aconfig
@@ -50,3 +50,11 @@
     description: "Enable default arrow icon when hovering on buttons or clickable widgets."
     bug: "299269803"
 }
+
+flag {
+    name: "enable_invalidate_check_thread"
+    namespace: "toolkit"
+    description: "Enable checkThread call in ViewRootImpl#onDescendentInvalidated"
+    bug: "333752000"
+    is_fixed_read_only: true
+}
diff --git a/core/java/android/window/flags/lse_desktop_experience.aconfig b/core/java/android/window/flags/lse_desktop_experience.aconfig
index 9524a6e..65e5f1a 100644
--- a/core/java/android/window/flags/lse_desktop_experience.aconfig
+++ b/core/java/android/window/flags/lse_desktop_experience.aconfig
@@ -49,3 +49,10 @@
     description: "Shows running apps in Desktop Mode Taskbar"
     bug: "332504528"
 }
+
+flag {
+    name: "enable_desktop_windowing_wallpaper_activity"
+    namespace: "lse_desktop_experience"
+    description: "Enables desktop wallpaper activity to show wallpaper in the desktop mode"
+    bug: "309014605"
+}
diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index fca4e91..b6558cb 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -440,6 +440,7 @@
     }
 
     private int internStacktraceString(TracingContext<
+            ProtoLogDataSource.Instance,
             ProtoLogDataSource.TlsState,
             ProtoLogDataSource.IncrementalState> ctx,
             String stacktrace) {
@@ -449,7 +450,8 @@
     }
 
     private int internStringArg(
-            TracingContext<ProtoLogDataSource.TlsState, ProtoLogDataSource.IncrementalState> ctx,
+            TracingContext<ProtoLogDataSource.Instance, ProtoLogDataSource.TlsState,
+                    ProtoLogDataSource.IncrementalState> ctx,
             String string
     ) {
         final ProtoLogDataSource.IncrementalState incrementalState = ctx.getIncrementalState();
@@ -458,7 +460,8 @@
     }
 
     private int internString(
-            TracingContext<ProtoLogDataSource.TlsState, ProtoLogDataSource.IncrementalState> ctx,
+            TracingContext<ProtoLogDataSource.Instance, ProtoLogDataSource.TlsState,
+                    ProtoLogDataSource.IncrementalState> ctx,
             Map<String, Integer> internMap,
             long fieldId,
             String string
diff --git a/core/jni/android_tracing_PerfettoDataSource.cpp b/core/jni/android_tracing_PerfettoDataSource.cpp
index 5c7b470..da690b0 100644
--- a/core/jni/android_tracing_PerfettoDataSource.cpp
+++ b/core/jni/android_tracing_PerfettoDataSource.cpp
@@ -93,49 +93,6 @@
     return instance;
 }
 
-jobject PerfettoDataSource::createTlsStateGlobalRef(JNIEnv* env, PerfettoDsInstanceIndex inst_id) {
-    ScopedLocalRef<jobject> args(env,
-                                 env->NewObject(gCreateTlsStateArgsClassInfo.clazz,
-                                                gCreateTlsStateArgsClassInfo.init, mJavaDataSource,
-                                                inst_id));
-
-    ScopedLocalRef<jobject> tslState(env,
-                                     env->CallObjectMethod(mJavaDataSource,
-                                                           gPerfettoDataSourceClassInfo
-                                                                   .createTlsState,
-                                                           args.get()));
-
-    if (env->ExceptionCheck()) {
-        LOGE_EX(env);
-        env->ExceptionClear();
-        LOG_ALWAYS_FATAL("Failed to create new Java Perfetto incremental state");
-    }
-
-    return env->NewGlobalRef(tslState.get());
-}
-
-jobject PerfettoDataSource::createIncrementalStateGlobalRef(JNIEnv* env,
-                                                            PerfettoDsInstanceIndex inst_id) {
-    ScopedLocalRef<jobject> args(env,
-                                 env->NewObject(gCreateIncrementalStateArgsClassInfo.clazz,
-                                                gCreateIncrementalStateArgsClassInfo.init,
-                                                mJavaDataSource, inst_id));
-
-    ScopedLocalRef<jobject> incrementalState(env,
-                                             env->CallObjectMethod(mJavaDataSource,
-                                                                   gPerfettoDataSourceClassInfo
-                                                                           .createIncrementalState,
-                                                                   args.get()));
-
-    if (env->ExceptionCheck()) {
-        LOGE_EX(env);
-        env->ExceptionClear();
-        LOG_ALWAYS_FATAL("Failed to create Java Perfetto incremental state");
-    }
-
-    return env->NewGlobalRef(incrementalState.get());
-}
-
 bool PerfettoDataSource::TraceIterateBegin() {
     if (gInIteration) {
         return false;
@@ -177,6 +134,15 @@
     gInIteration = false;
 }
 
+PerfettoDsInstanceIndex PerfettoDataSource::GetInstanceIndex() {
+    if (!gInIteration) {
+        LOG_ALWAYS_FATAL("Tried calling GetInstanceIndex outside of a tracer iteration.");
+        return -1;
+    }
+
+    return gIterator.impl.inst_id;
+}
+
 jobject PerfettoDataSource::GetCustomTls() {
     if (!gInIteration) {
         LOG_ALWAYS_FATAL("Tried getting CustomTls outside of a tracer iteration.");
@@ -189,6 +155,18 @@
     return tls_state->jobj;
 }
 
+void PerfettoDataSource::SetCustomTls(jobject tlsState) {
+    if (!gInIteration) {
+        LOG_ALWAYS_FATAL("Tried getting CustomTls outside of a tracer iteration.");
+        return;
+    }
+
+    TlsState* tls_state =
+            reinterpret_cast<TlsState*>(PerfettoDsGetCustomTls(&dataSource, &gIterator));
+
+    tls_state->jobj = tlsState;
+}
+
 jobject PerfettoDataSource::GetIncrementalState() {
     if (!gInIteration) {
         LOG_ALWAYS_FATAL("Tried getting IncrementalState outside of a tracer iteration.");
@@ -201,6 +179,18 @@
     return incr_state->jobj;
 }
 
+void PerfettoDataSource::SetIncrementalState(jobject incrementalState) {
+    if (!gInIteration) {
+        LOG_ALWAYS_FATAL("Tried getting IncrementalState outside of a tracer iteration.");
+        return;
+    }
+
+    IncrementalState* incr_state = reinterpret_cast<IncrementalState*>(
+            PerfettoDsGetIncrementalState(&dataSource, &gIterator));
+
+    incr_state->jobj = incrementalState;
+}
+
 void PerfettoDataSource::WritePackets(JNIEnv* env, jobjectArray packets) {
     if (!gInIteration) {
         LOG_ALWAYS_FATAL("Tried writing packets outside of a tracer iteration.");
@@ -211,7 +201,7 @@
     for (int i = 0; i < packets_count; i++) {
         jbyteArray packet_proto_buffer = (jbyteArray)env->GetObjectArrayElement(packets, i);
 
-        jbyte* raw_proto_buffer = env->GetByteArrayElements(packet_proto_buffer, 0);
+        jbyte* raw_proto_buffer = env->GetByteArrayElements(packet_proto_buffer, nullptr);
         int buffer_size = env->GetArrayLength(packet_proto_buffer);
 
         struct PerfettoDsRootTracePacket trace_packet;
@@ -219,6 +209,8 @@
         PerfettoPbMsgAppendBytes(&trace_packet.msg.msg, (const uint8_t*)raw_proto_buffer,
                                  buffer_size);
         PerfettoDsTracerPacketEnd(&gIterator, &trace_packet);
+
+        env->ReleaseByteArrayElements(packet_proto_buffer, raw_proto_buffer, 0 /* default mode */);
     }
 }
 
@@ -264,7 +256,7 @@
 }
 
 void nativeRegisterDataSource(JNIEnv* env, jclass clazz, jlong datasource_ptr,
-                              int buffer_exhausted_policy) {
+                              jint buffer_exhausted_policy) {
     sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(datasource_ptr);
 
     struct PerfettoDsParams params = PerfettoDsParamsDefault();
@@ -291,13 +283,8 @@
 
     params.on_create_tls_cb = [](struct PerfettoDsImpl* ds_impl, PerfettoDsInstanceIndex inst_id,
                                  struct PerfettoDsTracerImpl* tracer, void* user_arg) -> void* {
-        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
-
-        auto* datasource = reinterpret_cast<PerfettoDataSource*>(user_arg);
-
-        jobject java_tls_state = datasource->createTlsStateGlobalRef(env, inst_id);
-
-        auto* tls_state = new TlsState(java_tls_state);
+        // Populated later and only if required by the java side
+        auto* tls_state = new TlsState(NULL);
         return static_cast<void*>(tls_state);
     };
 
@@ -306,18 +293,16 @@
 
         TlsState* tls_state = reinterpret_cast<TlsState*>(ptr);
 
-        env->DeleteGlobalRef(tls_state->jobj);
+        if (tls_state->jobj != NULL) {
+            env->DeleteGlobalRef(tls_state->jobj);
+        }
         delete tls_state;
     };
 
     params.on_create_incr_cb = [](struct PerfettoDsImpl* ds_impl, PerfettoDsInstanceIndex inst_id,
                                   struct PerfettoDsTracerImpl* tracer, void* user_arg) -> void* {
-        JNIEnv* env = GetOrAttachJNIEnvironment(gVm, JNI_VERSION_1_6);
-
-        auto* datasource = reinterpret_cast<PerfettoDataSource*>(user_arg);
-        jobject java_incr_state = datasource->createIncrementalStateGlobalRef(env, inst_id);
-
-        auto* incr_state = new IncrementalState(java_incr_state);
+        // Populated later and only if required by the java side
+        auto* incr_state = new IncrementalState(NULL);
         return static_cast<void*>(incr_state);
     };
 
@@ -326,7 +311,9 @@
 
         IncrementalState* incr_state = reinterpret_cast<IncrementalState*>(ptr);
 
-        env->DeleteGlobalRef(incr_state->jobj);
+        if (incr_state->jobj != NULL) {
+            env->DeleteGlobalRef(incr_state->jobj);
+        }
         delete incr_state;
     };
 
@@ -401,16 +388,34 @@
     return datasource->TraceIterateBreak();
 }
 
+jint nativeGetPerfettoDsInstanceIndex(jlong dataSourcePtr) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
+    return (jint)datasource->GetInstanceIndex();
+}
+
 jobject nativeGetCustomTls(JNIEnv* /* env */, jclass /* clazz */, jlong dataSourcePtr) {
     sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
     return datasource->GetCustomTls();
 }
 
+void nativeSetCustomTls(JNIEnv* env, jclass /* clazz */, jlong dataSourcePtr, jobject tlsState) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
+    tlsState = env->NewGlobalRef(tlsState);
+    return datasource->SetCustomTls(tlsState);
+}
+
 jobject nativeGetIncrementalState(JNIEnv* /* env */, jclass /* clazz */, jlong dataSourcePtr) {
     sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
     return datasource->GetIncrementalState();
 }
 
+void nativeSetIncrementalState(JNIEnv* env, jclass /* clazz */, jlong dataSourcePtr,
+                               jobject incrementalState) {
+    sp<PerfettoDataSource> datasource = reinterpret_cast<PerfettoDataSource*>(dataSourcePtr);
+    incrementalState = env->NewGlobalRef(incrementalState);
+    return datasource->SetIncrementalState(incrementalState);
+}
+
 const JNINativeMethod gMethods[] = {
         /* name, signature, funcPtr */
         {"nativeCreate", "(Landroid/tracing/perfetto/DataSource;Ljava/lang/String;)J",
@@ -425,13 +430,16 @@
 
         {"nativePerfettoDsTraceIterateBegin", "(J)Z", (void*)nativePerfettoDsTraceIterateBegin},
         {"nativePerfettoDsTraceIterateNext", "(J)Z", (void*)nativePerfettoDsTraceIterateNext},
-        {"nativePerfettoDsTraceIterateBreak", "(J)V", (void*)nativePerfettoDsTraceIterateBreak}};
+        {"nativePerfettoDsTraceIterateBreak", "(J)V", (void*)nativePerfettoDsTraceIterateBreak},
+        {"nativeGetPerfettoDsInstanceIndex", "(J)I", (void*)nativeGetPerfettoDsInstanceIndex}};
 
 const JNINativeMethod gMethodsTracingContext[] = {
         /* name, signature, funcPtr */
         {"nativeFlush", "(J[[B)V", (void*)nativeFlush},
         {"nativeGetCustomTls", "(J)Ljava/lang/Object;", (void*)nativeGetCustomTls},
         {"nativeGetIncrementalState", "(J)Ljava/lang/Object;", (void*)nativeGetIncrementalState},
+        {"nativeSetCustomTls", "(JLjava/lang/Object;)V", (void*)nativeSetCustomTls},
+        {"nativeSetIncrementalState", "(JLjava/lang/Object;)V", (void*)nativeSetIncrementalState},
 };
 
 int register_android_tracing_PerfettoDataSource(JNIEnv* env) {
diff --git a/core/jni/android_tracing_PerfettoDataSource.h b/core/jni/android_tracing_PerfettoDataSource.h
index 209de29..fe15184 100644
--- a/core/jni/android_tracing_PerfettoDataSource.h
+++ b/core/jni/android_tracing_PerfettoDataSource.h
@@ -44,16 +44,16 @@
     jobject newInstance(JNIEnv* env, void* ds_config, size_t ds_config_size,
                         PerfettoDsInstanceIndex inst_id);
 
-    jobject createTlsStateGlobalRef(JNIEnv* env, PerfettoDsInstanceIndex inst_id);
-    jobject createIncrementalStateGlobalRef(JNIEnv* env, PerfettoDsInstanceIndex inst_id);
-
     bool TraceIterateBegin();
     bool TraceIterateNext();
     void TraceIterateBreak();
+    PerfettoDsInstanceIndex GetInstanceIndex();
     void WritePackets(JNIEnv* env, jobjectArray packets);
 
     jobject GetCustomTls();
+    void SetCustomTls(jobject);
     jobject GetIncrementalState();
+    void SetIncrementalState(jobject);
 
     void flushAll();
 
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index ab714ad..f55f3c7 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -5890,7 +5890,7 @@
     <!-- Allows an application to subscribe to notifications about the nearby devices' presence
          status change base on the UUIDs.
          <p>Not for use by third-party applications.</p>
-         @FlaggedApi("android.companion.flags.device_presence")
+         @FlaggedApi("android.companion.device_presence")
     -->
     <permission android:name="android.permission.REQUEST_OBSERVE_DEVICE_UUID_PRESENCE"
                 android:protectionLevel="signature|privileged" />
@@ -8182,6 +8182,15 @@
     <permission android:name="android.permission.SCREEN_TIMEOUT_OVERRIDE"
                 android:protectionLevel="signature" />
 
+    <!-- @SystemApi
+        @FlaggedApi("android.security.fsverity_api")
+        Allows app to setup fs-verity through FileIntegrityManager.
+        <p>Protection level: signature|privileged
+        @hide
+    -->
+    <permission android:name="android.permission.SETUP_FSVERITY"
+                android:protectionLevel="signature|privileged"/>
+
     <!-- Attribution for Geofencing service. -->
     <attribution android:tag="GeofencingService" android:label="@string/geofencing_service"/>
     <!-- Attribution for Country Detector. -->
diff --git a/core/res/res/layout/notification_template_header.xml b/core/res/res/layout/notification_template_header.xml
index be1c939..6f06d80 100644
--- a/core/res/res/layout/notification_template_header.xml
+++ b/core/res/res/layout/notification_template_header.xml
@@ -80,6 +80,7 @@
         android:layout_height="match_parent"
         android:layout_alignParentStart="true"
         android:importantForAccessibility="no"
+        android:focusable="false"
         />
 
     <include layout="@layout/notification_expand_button"
diff --git a/core/res/res/layout/notification_template_material_base.xml b/core/res/res/layout/notification_template_material_base.xml
index 710a70a..64227d8 100644
--- a/core/res/res/layout/notification_template_material_base.xml
+++ b/core/res/res/layout/notification_template_material_base.xml
@@ -55,6 +55,7 @@
         android:layout_height="match_parent"
         android:layout_gravity="start"
         android:importantForAccessibility="no"
+        android:focusable="false"
         />
 
     <LinearLayout
diff --git a/core/res/res/layout/notification_template_material_media.xml b/core/res/res/layout/notification_template_material_media.xml
index df32d30..8a94c48 100644
--- a/core/res/res/layout/notification_template_material_media.xml
+++ b/core/res/res/layout/notification_template_material_media.xml
@@ -54,6 +54,7 @@
         android:layout_height="match_parent"
         android:layout_gravity="start"
         android:importantForAccessibility="no"
+        android:focusable="false"
         />
 
     <LinearLayout
diff --git a/core/res/res/layout/notification_template_material_messaging.xml b/core/res/res/layout/notification_template_material_messaging.xml
index 3e82bd1..a83d923 100644
--- a/core/res/res/layout/notification_template_material_messaging.xml
+++ b/core/res/res/layout/notification_template_material_messaging.xml
@@ -67,6 +67,7 @@
                 android:layout_height="match_parent"
                 android:layout_gravity="start"
                 android:importantForAccessibility="no"
+                android:focusable="false"
                 />
 
             <!--
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index 0f04321..46a3e7f 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -861,9 +861,9 @@
         // container update.
         updateDivider(wct, taskContainer);
 
-        // If the last direct activity of the host task is dismissed and the overlay container is
-        // the only taskFragment, the overlay container should also be dismissed.
-        dismissOverlayContainerIfNeeded(wct, taskContainer);
+        // If the last direct activity of the host task is dismissed and there's an always-on-top
+        // overlay container in the task, the overlay container should also be dismissed.
+        dismissAlwaysOnTopOverlayIfNeeded(wct, taskContainer);
 
         if (!shouldUpdateContainer) {
             return;
@@ -1990,7 +1990,7 @@
             @NonNull TaskFragmentContainer container) {
         final TaskContainer taskContainer = container.getTaskContainer();
 
-        if (dismissOverlayContainerIfNeeded(wct, taskContainer)) {
+        if (dismissAlwaysOnTopOverlayIfNeeded(wct, taskContainer)) {
             return;
         }
 
@@ -2014,22 +2014,27 @@
         }
     }
 
-    /** Dismisses the overlay container in the {@code taskContainer} if needed. */
+    /**
+     * Dismisses {@link TaskFragmentContainer#isAlwaysOnTopOverlay()} in the {@code taskContainer}
+     * if needed.
+     */
     @GuardedBy("mLock")
-    private boolean dismissOverlayContainerIfNeeded(@NonNull WindowContainerTransaction wct,
-            @NonNull TaskContainer taskContainer) {
-        final TaskFragmentContainer overlayContainer = taskContainer.getOverlayContainer();
-        if (overlayContainer == null) {
+    private boolean dismissAlwaysOnTopOverlayIfNeeded(@NonNull WindowContainerTransaction wct,
+                                                      @NonNull TaskContainer taskContainer) {
+        // Dismiss always-on-top overlay container if it's the only container in the task and
+        // there's no direct activity in the parent task.
+        final List<TaskFragmentContainer> containers = taskContainer.getTaskFragmentContainers();
+        if (containers.size() != 1 || taskContainer.hasDirectActivity()) {
             return false;
         }
-        // Dismiss the overlay container if it's the only container in the task and there's no
-        // direct activity in the parent task.
-        if (taskContainer.getTaskFragmentContainers().size() == 1
-                && !taskContainer.hasDirectActivity()) {
-            mPresenter.cleanupContainer(wct, overlayContainer, false /* shouldFinishDependant */);
-            return true;
+
+        final TaskFragmentContainer container = containers.getLast();
+        if (!container.isAlwaysOnTopOverlay()) {
+            return false;
         }
-        return false;
+
+        mPresenter.cleanupContainer(wct, container, false /* shouldFinishDependant */);
+        return true;
     }
 
     /**
@@ -2620,25 +2625,42 @@
     /**
      * Gets all overlay containers from all tasks in this process, or an empty list if there's
      * no overlay container.
-     * <p>
-     * Note that we only support one overlay container for each task, but an app could have multiple
-     * tasks.
      */
     @VisibleForTesting
     @GuardedBy("mLock")
     @NonNull
-    List<TaskFragmentContainer> getAllOverlayTaskFragmentContainers() {
+    List<TaskFragmentContainer> getAllNonFinishingOverlayContainers() {
         final List<TaskFragmentContainer> overlayContainers = new ArrayList<>();
         for (int i = 0; i < mTaskContainers.size(); i++) {
             final TaskContainer taskContainer = mTaskContainers.valueAt(i);
-            final TaskFragmentContainer overlayContainer = taskContainer.getOverlayContainer();
-            if (overlayContainer != null) {
-                overlayContainers.add(overlayContainer);
-            }
+            final List<TaskFragmentContainer> overlayContainersPerTask = taskContainer
+                    .getTaskFragmentContainers()
+                    .stream()
+                    .filter(c -> c.isOverlay() && !c.isFinished())
+                    .toList();
+            overlayContainers.addAll(overlayContainersPerTask);
         }
         return overlayContainers;
     }
 
+    /**
+     * Creates an overlay container or updates a visible overlay container if its
+     * {@link TaskFragmentContainer#getTaskId()}, {@link TaskFragmentContainer#getOverlayTag()}
+     * and {@link TaskFragmentContainer#getAssociatedActivityToken()} matches.
+     * <p>
+     * This method will also dismiss any existing overlay container if:
+     * <ul>
+     *   <li>it's visible but not meet the criteria to update overlay</li>
+     *   <li>{@link TaskFragmentContainer#getOverlayTag()} matches but not meet the criteria to
+     *   update overlay</li>
+     * </ul>
+     *
+     * @param wct the {@link WindowContainerTransaction}
+     * @param options the {@link ActivityOptions} to launch the overlay
+     * @param intent the intent of activity to launch
+     * @param launchActivity the activity to launch the overlay container
+     * @return the overlay container
+     */
     @VisibleForTesting
     // Suppress GuardedBy warning because lint ask to mark this method as
     // @GuardedBy(container.mController.mLock), which is mLock itself
@@ -2649,7 +2671,7 @@
             @NonNull WindowContainerTransaction wct, @NonNull Bundle options,
             @NonNull Intent intent, @NonNull Activity launchActivity) {
         final List<TaskFragmentContainer> overlayContainers =
-                getAllOverlayTaskFragmentContainers();
+                getAllNonFinishingOverlayContainers();
         final String overlayTag = Objects.requireNonNull(options.getString(KEY_OVERLAY_TAG));
         final boolean associateLaunchingActivity = options
                 .getBoolean(KEY_OVERLAY_ASSOCIATE_WITH_LAUNCHING_ACTIVITY, true);
@@ -2673,62 +2695,86 @@
         final int taskId = getTaskId(launchActivity);
         if (!overlayContainers.isEmpty()) {
             for (final TaskFragmentContainer overlayContainer : overlayContainers) {
-                if (!overlayTag.equals(overlayContainer.getOverlayTag())
-                        && taskId == overlayContainer.getTaskId()) {
-                    // If there's an overlay container with different tag shown in the same
-                    // task, dismiss the existing overlay container.
-                    mPresenter.cleanupContainer(wct, overlayContainer,
-                            false /* shouldFinishDependant */);
-                }
-                if (overlayTag.equals(overlayContainer.getOverlayTag())
-                        && taskId != overlayContainer.getTaskId()) {
-                    Log.w(TAG, "The overlay container with tag:"
-                            + overlayContainer.getOverlayTag() + " is dismissed because"
-                            + " there's an existing overlay container with the same tag but"
-                            + " different task ID:" + overlayContainer.getTaskId() + ". "
-                            + "The new associated activity is " + launchActivity);
+                final boolean isTopNonFinishingOverlay = overlayContainer.equals(
+                        overlayContainer.getTaskContainer().getTopNonFinishingTaskFragmentContainer(
+                                true /* includePin */, true /* includeOverlay */));
+                if (taskId != overlayContainer.getTaskId()) {
                     // If there's an overlay container with same tag in a different task,
                     // dismiss the overlay container since the tag must be unique per process.
-                    mPresenter.cleanupContainer(wct, overlayContainer,
-                            false /* shouldFinishDependant */);
-                }
-                if (overlayTag.equals(overlayContainer.getOverlayTag())
-                        && taskId == overlayContainer.getTaskId()) {
-                    if (associateLaunchingActivity && !launchActivity.getActivityToken()
-                            .equals(overlayContainer.getAssociatedActivityToken())) {
+                    if (overlayTag.equals(overlayContainer.getOverlayTag())) {
                         Log.w(TAG, "The overlay container with tag:"
                                 + overlayContainer.getOverlayTag() + " is dismissed because"
                                 + " there's an existing overlay container with the same tag but"
-                                + " different associated launching activity. The new associated"
-                                + " activity is " + launchActivity);
-                        // The associated activity must be the same, or it will be dismissed.
+                                + " different task ID:" + overlayContainer.getTaskId() + ". "
+                                + "The new associated activity is " + launchActivity);
                         mPresenter.cleanupContainer(wct, overlayContainer,
                                 false /* shouldFinishDependant */);
-                    } else if (!associateLaunchingActivity
-                            && overlayContainer.isAssociatedWithActivity()) {
+                    }
+                    continue;
+                }
+                if (!overlayTag.equals(overlayContainer.getOverlayTag())) {
+                    // If there's an overlay container with different tag on top in the same
+                    // task, dismiss the existing overlay container.
+                    if (isTopNonFinishingOverlay) {
+                        mPresenter.cleanupContainer(wct, overlayContainer,
+                                false /* shouldFinishDependant */);
+                    }
+                    continue;
+                }
+                // The overlay container has the same tag and task ID with the new launching
+                // overlay container.
+                if (!isTopNonFinishingOverlay) {
+                    // Dismiss the invisible overlay container regardless of activity
+                    // association if it collides the tag of new launched overlay container .
+                    Log.w(TAG, "The invisible overlay container with tag:"
+                            + overlayContainer.getOverlayTag() + " is dismissed because"
+                            + " there's a launching overlay container with the same tag."
+                            + " The new associated activity is " + launchActivity);
+                    mPresenter.cleanupContainer(wct, overlayContainer,
+                            false /* shouldFinishDependant */);
+                    continue;
+                }
+                // Requesting an always-on-top overlay.
+                if (!associateLaunchingActivity) {
+                    if (overlayContainer.isAssociatedWithActivity()) {
+                        // Dismiss the overlay container since it has associated with an activity.
                         Log.w(TAG, "The overlay container with tag:"
                                 + overlayContainer.getOverlayTag() + " is dismissed because"
                                 + " there's an existing overlay container with the same tag but"
                                 + " different associated launching activity. The overlay container"
                                 + " doesn't associate with any activity.");
-                        // Dismiss the overlay container since it has been associated with an
-                        // activity.
                         mPresenter.cleanupContainer(wct, overlayContainer,
                                 false /* shouldFinishDependant */);
+                        continue;
                     } else {
-                        // Just update the overlay container if
-                        // - should associate with an activity and associated activity matches
-                        // - should not associate with an activity and the overlay container
-                        //     don't have an associated activity
+                        // The existing overlay container doesn't associate an activity as well.
+                        // Just update the overlay and return.
+                        // Note that going to this condition means the tag, task ID matches a
+                        // visible always-on-top overlay, and won't dismiss any overlay any more.
                         mPresenter.applyActivityStackAttributes(wct, overlayContainer, attrs,
                                 getMinDimensions(intent));
-                        // We can just return the updated overlay container and don't need to
-                        // check other condition since we only have one OverlayCreateParams, and
-                        // if the tag and task are matched, it's impossible to match another task
-                        // or tag since tags and tasks are all unique.
                         return overlayContainer;
                     }
                 }
+                if (launchActivity.getActivityToken()
+                        != overlayContainer.getAssociatedActivityToken()) {
+                    Log.w(TAG, "The overlay container with tag:"
+                            + overlayContainer.getOverlayTag() + " is dismissed because"
+                            + " there's an existing overlay container with the same tag but"
+                            + " different associated launching activity. The new associated"
+                            + " activity is " + launchActivity);
+                    // The associated activity must be the same, or it will be dismissed.
+                    mPresenter.cleanupContainer(wct, overlayContainer,
+                            false /* shouldFinishDependant */);
+                    continue;
+                }
+                // Reaching here means the launching activity launch an overlay container with the
+                // same task ID, tag, while there's a previously launching visible overlay
+                // container. We'll regard it as updating the existing overlay container.
+                mPresenter.applyActivityStackAttributes(wct, overlayContainer, attrs,
+                        getMinDimensions(intent));
+                return overlayContainer;
+
             }
         }
         // Launch the overlay container to the task with taskId.
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
index b56f671..6231ea0 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
@@ -467,6 +467,11 @@
             reorderTaskFragmentToFront(wct,
                     pinnedContainer.getSecondaryContainer().getTaskFragmentToken());
         }
+        final TaskFragmentContainer alwaysOnTopOverlayContainer = container.getTaskContainer()
+                .getAlwaysOnTopOverlayContainer();
+        if (alwaysOnTopOverlayContainer != null) {
+            reorderTaskFragmentToFront(wct, alwaysOnTopOverlayContainer.getTaskFragmentToken());
+        }
     }
 
     @Override
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
index 3dd96c4..fdf0910 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
@@ -70,9 +70,11 @@
     @Nullable
     private SplitPinContainer mSplitPinContainer;
 
-    /** The overlay container in this Task. */
+    /**
+     * The {@link TaskFragmentContainer#isAlwaysOnTopOverlay()} in the Task.
+     */
     @Nullable
-    private TaskFragmentContainer mOverlayContainer;
+    private TaskFragmentContainer mAlwaysOnTopOverlayContainer;
 
     @NonNull
     private final Configuration mConfiguration;
@@ -316,10 +318,12 @@
         return null;
     }
 
-    /** Returns the overlay container in the task, or {@code null} if it doesn't exist. */
+    /**
+     * Returns the always-on-top overlay container in the task, or {@code null} if it doesn't exist.
+     */
     @Nullable
-    TaskFragmentContainer getOverlayContainer() {
-        return mOverlayContainer;
+    TaskFragmentContainer getAlwaysOnTopOverlayContainer() {
+        return mAlwaysOnTopOverlayContainer;
     }
 
     int indexOf(@NonNull TaskFragmentContainer child) {
@@ -531,7 +535,21 @@
         updateSplitPinContainerIfNecessary();
         // Update overlay container after split pin container since the overlay should be on top of
         // pin container.
-        updateOverlayContainerIfNecessary();
+        updateAlwaysOnTopOverlayIfNecessary();
+    }
+
+    private void updateAlwaysOnTopOverlayIfNecessary() {
+        final List<TaskFragmentContainer> alwaysOnTopOverlays = mContainers
+                .stream().filter(TaskFragmentContainer::isAlwaysOnTopOverlay).toList();
+        if (alwaysOnTopOverlays.size() > 1) {
+            throw new IllegalStateException("There must be at most one always-on-top overlay "
+                    + "container per Task");
+        }
+        mAlwaysOnTopOverlayContainer = alwaysOnTopOverlays.isEmpty()
+                ? null : alwaysOnTopOverlays.getFirst();
+        if (mAlwaysOnTopOverlayContainer != null) {
+            moveContainerToLastIfNecessary(mAlwaysOnTopOverlayContainer);
+        }
     }
 
     private void updateSplitPinContainerIfNecessary() {
@@ -559,18 +577,6 @@
         }
     }
 
-    private void updateOverlayContainerIfNecessary() {
-        final List<TaskFragmentContainer> overlayContainers = mContainers.stream()
-                .filter(TaskFragmentContainer::isOverlay).toList();
-        if (overlayContainers.size() > 1) {
-            throw new IllegalStateException("There must be at most one overlay container per Task");
-        }
-        mOverlayContainer = overlayContainers.isEmpty() ? null : overlayContainers.get(0);
-        if (mOverlayContainer != null) {
-            moveContainerToLastIfNecessary(mOverlayContainer);
-        }
-    }
-
     /** Moves the {@code container} to the last to align taskFragments' z-order. */
     private void moveContainerToLastIfNecessary(@NonNull TaskFragmentContainer container) {
         final int index = mContainers.indexOf(container);
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
index 5dbb016..094ebcb 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
@@ -1023,6 +1023,14 @@
         return mAssociatedActivityToken != null;
     }
 
+    /**
+     * Returns {@code true} if the overlay container should be always on top, which should be
+     * a non-fill-parent overlay without activity association.
+     */
+    boolean isAlwaysOnTopOverlay() {
+        return isOverlay() && !isAssociatedWithActivity();
+    }
+
     @Override
     public String toString() {
         return toString(true /* includeContainersToFinishOnExit */);
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
index dcdbe59..b1b1984 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
@@ -178,37 +178,59 @@
     }
 
     @Test
-    public void testGetOverlayContainers() {
-        assertThat(mSplitController.getAllOverlayTaskFragmentContainers()).isEmpty();
+    public void testGetAllNonFinishingOverlayContainers() {
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers()).isEmpty();
 
         final TaskFragmentContainer overlayContainer1 =
                 createTestOverlayContainer(TASK_ID, "test1");
 
-        assertThat(mSplitController.getAllOverlayTaskFragmentContainers())
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
                 .containsExactly(overlayContainer1);
 
-        assertThrows(
-                "The exception must throw if there are two overlay containers in the same task.",
-                IllegalStateException.class,
-                () -> createTestOverlayContainer(TASK_ID, "test2"));
+        final TaskFragmentContainer overlayContainer2 =
+                createTestOverlayContainer(TASK_ID, "test2");
+
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
+                .containsExactly(overlayContainer1, overlayContainer2);
 
         final TaskFragmentContainer overlayContainer3 =
                 createTestOverlayContainer(TASK_ID + 1, "test3");
 
-        assertThat(mSplitController.getAllOverlayTaskFragmentContainers())
-                .containsExactly(overlayContainer1, overlayContainer3);
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
+                .containsExactly(overlayContainer1, overlayContainer2, overlayContainer3);
+
+        final TaskFragmentContainer finishingOverlayContainer =
+                createTestOverlayContainer(TASK_ID, "test4");
+        spyOn(finishingOverlayContainer);
+        doReturn(true).when(finishingOverlayContainer).isFinished();
+
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
+                .containsExactly(overlayContainer1, overlayContainer2, overlayContainer3);
     }
 
     @Test
-    public void testCreateOrUpdateOverlayTaskFragmentIfNeeded_anotherTagInTask_dismissOverlay() {
+    public void testCreateOrUpdateOverlayTaskFragmentIfNeeded_anotherTagInTask() {
+        createExistingOverlayContainers(false /* visible */);
+        createMockTaskFragmentContainer(mActivity);
+
+        final TaskFragmentContainer overlayContainer =
+                createOrUpdateOverlayTaskFragmentIfNeeded("test3");
+
+        assertWithMessage("overlayContainer1 is still there since it's not visible.")
+                .that(mSplitController.getAllNonFinishingOverlayContainers())
+                .containsExactly(mOverlayContainer1, mOverlayContainer2, overlayContainer);
+    }
+
+    @Test
+    public void testCreateOrUpdateOverlay_visibleOverlaySameTagInTask_dismissOverlay() {
         createExistingOverlayContainers();
 
         final TaskFragmentContainer overlayContainer =
                 createOrUpdateOverlayTaskFragmentIfNeeded("test3");
 
-        assertWithMessage("overlayContainer1 must be dismissed since the new overlay container"
-                + " is launched to the same task")
-                .that(mSplitController.getAllOverlayTaskFragmentContainers())
+        assertWithMessage("overlayContainer1 must be dismissed since it's visible"
+                + " in the same task.")
+                .that(mSplitController.getAllNonFinishingOverlayContainers())
                 .containsExactly(mOverlayContainer2, overlayContainer);
     }
 
@@ -223,7 +245,7 @@
         assertWithMessage("overlayContainer1 must be dismissed since the new overlay container"
                 + " is launched with the same tag as an existing overlay container in a different "
                 + "task")
-                .that(mSplitController.getAllOverlayTaskFragmentContainers())
+                .that(mSplitController.getAllNonFinishingOverlayContainers())
                 .containsExactly(mOverlayContainer2, overlayContainer);
     }
 
@@ -240,7 +262,7 @@
 
         assertWithMessage("overlayContainer1 must be updated since the new overlay container"
                 + " is launched with the same tag and task")
-                .that(mSplitController.getAllOverlayTaskFragmentContainers())
+                .that(mSplitController.getAllNonFinishingOverlayContainers())
                 .containsExactly(mOverlayContainer1, mOverlayContainer2);
 
         assertThat(overlayContainer).isEqualTo(mOverlayContainer1);
@@ -260,11 +282,26 @@
 
         assertWithMessage("overlayContainer1 must be dismissed since the new overlay container"
                 + " is associated with different launching activity")
-                .that(mSplitController.getAllOverlayTaskFragmentContainers())
+                .that(mSplitController.getAllNonFinishingOverlayContainers())
                 .containsExactly(mOverlayContainer2, overlayContainer);
     }
 
     @Test
+    public void testCreateOrUpdateOverlayTaskFragmentIfNeeded_dismissOverlay() {
+        createExistingOverlayContainers(false /* visible */);
+        createMockTaskFragmentContainer(mActivity);
+
+        final TaskFragmentContainer overlayContainer =
+                createOrUpdateOverlayTaskFragmentIfNeeded("test2");
+
+        // OverlayContainer2 is dismissed since new container is launched with the
+        // same tag in different task.
+        assertWithMessage("overlayContainer1 must be dismissed")
+                .that(mSplitController.getAllNonFinishingOverlayContainers())
+                .containsExactly(mOverlayContainer1, overlayContainer);
+    }
+
+    @Test
     public void testCreateOrUpdateOverlayTaskFragmentIfNeeded_dismissMultipleOverlays() {
         createExistingOverlayContainers();
 
@@ -275,15 +312,19 @@
         // different tag. OverlayContainer2 is dismissed since new container is launched with the
         // same tag in different task.
         assertWithMessage("overlayContainer1 and overlayContainer2 must be dismissed")
-                .that(mSplitController.getAllOverlayTaskFragmentContainers())
+                .that(mSplitController.getAllNonFinishingOverlayContainers())
                 .containsExactly(overlayContainer);
     }
 
     private void createExistingOverlayContainers() {
-        mOverlayContainer1 = createTestOverlayContainer(TASK_ID, "test1");
-        mOverlayContainer2 = createTestOverlayContainer(TASK_ID + 1, "test2");
+        createExistingOverlayContainers(true /* visible */);
+    }
+
+    private void createExistingOverlayContainers(boolean visible) {
+        mOverlayContainer1 = createTestOverlayContainer(TASK_ID, "test1", visible);
+        mOverlayContainer2 = createTestOverlayContainer(TASK_ID + 1, "test2", visible);
         List<TaskFragmentContainer> overlayContainers = mSplitController
-                .getAllOverlayTaskFragmentContainers();
+                .getAllNonFinishingOverlayContainers();
         assertThat(overlayContainers).containsExactly(mOverlayContainer1, mOverlayContainer2);
     }
 
@@ -314,7 +355,7 @@
                 createOrUpdateOverlayTaskFragmentIfNeeded("test");
         setupTaskFragmentInfo(overlayContainer, mActivity, true /* isVisible */);
 
-        assertThat(mSplitController.getAllOverlayTaskFragmentContainers())
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
                 .containsExactly(overlayContainer);
         assertThat(overlayContainer.getTaskId()).isEqualTo(TASK_ID);
         assertThat(overlayContainer.areLastRequestedBoundsEqual(bounds)).isTrue();
@@ -322,14 +363,16 @@
     }
 
     @Test
-    public void testGetTopNonFishingTaskFragmentContainerWithOverlay() {
-        final TaskFragmentContainer overlayContainer =
-                createTestOverlayContainer(TASK_ID, "test1");
-
-        // Add a SplitPinContainer, the overlay should be on top
+    public void testGetTopNonFishingTaskFragmentContainerWithoutAssociatedOverlay() {
         final Activity primaryActivity = createMockActivity();
         final Activity secondaryActivity = createMockActivity();
-
+        final Rect bounds = new Rect(0, 0, 100, 100);
+        mSplitController.setActivityStackAttributesCalculator(params ->
+                new ActivityStackAttributes.Builder().setRelativeBounds(bounds).build());
+        final TaskFragmentContainer overlayContainer =
+                createTestOverlayContainer(TASK_ID, "test1", true /* isVisible */,
+                        false /* shouldAssociateWithActivity */);
+        overlayContainer.setIsolatedNavigationEnabled(true);
         final TaskFragmentContainer primaryContainer =
                 createMockTaskFragmentContainer(primaryActivity);
         final TaskFragmentContainer secondaryContainer =
@@ -371,10 +414,10 @@
 
     @Test
     public void testGetTopNonFinishingActivityWithOverlay() {
-        TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test1");
-
         final Activity activity = createMockActivity();
         final TaskFragmentContainer container = createMockTaskFragmentContainer(activity);
+        final TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID,
+                "test1");
         final TaskContainer task = container.getTaskContainer();
 
         assertThat(task.getTopNonFinishingActivity(true /* includeOverlay */))
@@ -391,8 +434,9 @@
     }
 
     @Test
-    public void testUpdateOverlayContainer_dismissOverlayIfNeeded() {
-        TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test");
+    public void testUpdateOverlayContainer_dismissNonAssociatedOverlayIfNeeded() {
+        TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test",
+                true /* isVisible */, false /* associatedLaunchingActivity */);
 
         mSplitController.updateOverlayContainer(mTransaction, overlayContainer);
 
@@ -461,11 +505,10 @@
 
     @Test
     public void testOnTaskFragmentParentInfoChanged_positionOnlyChange_earlyReturn() {
-        final TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test");
+        final TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test",
+                true /* isVisible */, false /* associatedLaunchingActivity */);
         final TaskContainer taskContainer = overlayContainer.getTaskContainer();
 
-        assertThat(taskContainer.getOverlayContainer()).isEqualTo(overlayContainer);
-
         spyOn(taskContainer);
         final TaskContainer.TaskProperties taskProperties = taskContainer.getTaskProperties();
         final TaskFragmentParentInfo parentInfo = new TaskFragmentParentInfo(
@@ -481,16 +524,15 @@
 
         assertWithMessage("The overlay container must still be dismissed even if "
                 + "#updateContainer is not called")
-                .that(taskContainer.getOverlayContainer()).isNull();
+                .that(taskContainer.getTaskFragmentContainers()).isEmpty();
     }
 
     @Test
-    public void testOnTaskFragmentParentInfoChanged_invisibleTask_callDismissOverlayContainer() {
-        final TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test");
+    public void testOnTaskFragmentParentInfoChanged_invisibleTask_callDismissNonAssocOverlay() {
+        final TaskFragmentContainer overlayContainer = createTestOverlayContainer(TASK_ID, "test",
+                true /* isVisible */, false /* associatedLaunchingActivity */);
         final TaskContainer taskContainer = overlayContainer.getTaskContainer();
 
-        assertThat(taskContainer.getOverlayContainer()).isEqualTo(overlayContainer);
-
         spyOn(taskContainer);
         final TaskContainer.TaskProperties taskProperties = taskContainer.getTaskProperties();
         final TaskFragmentParentInfo parentInfo = new TaskFragmentParentInfo(
@@ -505,7 +547,7 @@
 
         assertWithMessage("The overlay container must still be dismissed even if "
                 + "#updateContainer is not called")
-                .that(taskContainer.getOverlayContainer()).isNull();
+                .that(taskContainer.getTaskFragmentContainers()).isEmpty();
     }
 
     @Test
@@ -529,8 +571,7 @@
 
     @Test
     public void testApplyActivityStackAttributesForOverlayContainerAssociatedWithActivity() {
-        final TaskFragmentContainer container = createTestOverlayContainer(TASK_ID,
-                TEST_TAG, true /* associatedWithLaunchingActivity */);
+        final TaskFragmentContainer container = createTestOverlayContainer(TASK_ID, TEST_TAG);
         final IBinder token = container.getTaskFragmentToken();
         final ActivityStackAttributes attributes = new ActivityStackAttributes.Builder()
                 .setRelativeBounds(new Rect(0, 0, 200, 200))
@@ -555,7 +596,7 @@
     @Test
     public void testApplyActivityStackAttributesForOverlayContainerWithoutAssociatedActivity() {
         final TaskFragmentContainer container = createTestOverlayContainer(TASK_ID, TEST_TAG,
-                false /* associatedWithLaunchingActivity */);
+                true, /* isVisible */ false /* associatedWithLaunchingActivity */);
         final IBinder token = container.getTaskFragmentToken();
         final ActivityStackAttributes attributes = new ActivityStackAttributes.Builder()
                 .setRelativeBounds(new Rect(0, 0, 200, 200))
@@ -610,8 +651,7 @@
         mSplitPresenter.applyActivityStackAttributes(mTransaction, container, attributes,
                 new Size(relativeBounds.width() + 1, relativeBounds.height()));
 
-        verify(mSplitPresenter).resizeTaskFragmentIfRegistered(mTransaction, container,
-                new Rect());
+        verify(mSplitPresenter).resizeTaskFragmentIfRegistered(mTransaction, container, new Rect());
         verify(mSplitPresenter).updateTaskFragmentWindowingModeIfRegistered(mTransaction, container,
                 WINDOWING_MODE_UNDEFINED);
         verify(mSplitPresenter).updateAnimationParams(mTransaction, token,
@@ -635,14 +675,14 @@
                 mActivity.getActivityToken());
 
         verify(mSplitPresenter, never()).cleanupContainer(any(), any(), anyBoolean());
-        assertThat(mSplitController.getAllOverlayTaskFragmentContainers())
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
                 .contains(overlayWithoutAssociation);
 
         TaskFragmentContainer overlayWithAssociation =
                 createOrUpdateOverlayTaskFragmentIfNeeded("test");
         overlayWithAssociation.setInfo(mTransaction, createMockTaskFragmentInfo(
                 overlayWithAssociation, mActivity, true /* isVisible */));
-        assertThat(mSplitController.getAllOverlayTaskFragmentContainers())
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
                 .contains(overlayWithAssociation);
         clearInvocations(mSplitPresenter);
 
@@ -655,7 +695,7 @@
 
         verify(mSplitPresenter).cleanupContainer(mTransaction, overlayWithAssociation, false);
 
-        assertThat(mSplitController.getAllOverlayTaskFragmentContainers())
+        assertThat(mSplitController.getAllNonFinishingOverlayContainers())
                 .doesNotContain(overlayWithAssociation);
     }
 
@@ -692,7 +732,14 @@
 
     @NonNull
     private TaskFragmentContainer createTestOverlayContainer(int taskId, @NonNull String tag) {
-        return createTestOverlayContainer(taskId, tag,
+        return createTestOverlayContainer(taskId, tag, false /* isVisible */,
+                true /* associateLaunchingActivity */);
+    }
+
+    @NonNull
+    private TaskFragmentContainer createTestOverlayContainer(int taskId, @NonNull String tag,
+            boolean isVisible) {
+        return createTestOverlayContainer(taskId, tag, isVisible,
                 true /* associateLaunchingActivity */);
     }
 
@@ -700,13 +747,13 @@
     //  once we have use cases.
     @NonNull
     private TaskFragmentContainer createTestOverlayContainer(int taskId, @NonNull String tag,
-                boolean associateLaunchingActivity) {
+                boolean isVisible, boolean associateLaunchingActivity) {
         Activity activity = createMockActivity();
         TaskFragmentContainer overlayContainer = mSplitController.newContainer(
                 null /* pendingAppearedActivity */, mIntent, activity, taskId,
                 null /* pairedPrimaryContainer */, tag, Bundle.EMPTY,
                 associateLaunchingActivity);
-        setupTaskFragmentInfo(overlayContainer, activity, false /* isVisible */);
+        setupTaskFragmentInfo(overlayContainer, activity, isVisible);
         return overlayContainer;
     }
 
diff --git a/libs/WindowManager/Shell/AndroidManifest.xml b/libs/WindowManager/Shell/AndroidManifest.xml
index 36d3313..7a98683 100644
--- a/libs/WindowManager/Shell/AndroidManifest.xml
+++ b/libs/WindowManager/Shell/AndroidManifest.xml
@@ -23,4 +23,12 @@
     <uses-permission android:name="android.permission.ROTATE_SURFACE_FLINGER" />
     <uses-permission android:name="android.permission.WAKEUP_SURFACE_FLINGER" />
     <uses-permission android:name="android.permission.READ_FRAME_BUFFER" />
+
+    <application>
+        <activity
+            android:name=".desktopmode.DesktopWallpaperActivity"
+            android:excludeFromRecents="true"
+            android:launchMode="singleInstance"
+            android:theme="@style/DesktopWallpaperTheme" />
+    </application>
 </manifest>
diff --git a/libs/WindowManager/Shell/res/values/styles.xml b/libs/WindowManager/Shell/res/values/styles.xml
index 08c2a02..13c0e66 100644
--- a/libs/WindowManager/Shell/res/values/styles.xml
+++ b/libs/WindowManager/Shell/res/values/styles.xml
@@ -23,6 +23,14 @@
         <item name="android:windowAnimationStyle">@style/Animation.ForcedResizable</item>
     </style>
 
+    <!-- Theme used for the activity that shows below the desktop mode windows to show wallpaper -->
+    <style name="DesktopWallpaperTheme" parent="@android:style/Theme.Wallpaper.NoTitleBar">
+        <item name="android:statusBarColor">@android:color/transparent</item>
+        <item name="android:navigationBarColor">@android:color/transparent</item>
+        <item name="android:windowDrawsSystemBarBackgrounds">true</item>
+        <item name="android:windowAnimationStyle">@null</item>
+    </style>
+
     <style name="Animation.ForcedResizable" parent="@android:style/Animation">
         <item name="android:activityOpenEnterAnimation">@anim/forced_resizable_enter</item>
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index ff00a7b..1408ead 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -59,6 +59,7 @@
 import com.android.wm.shell.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
+import com.android.wm.shell.desktopmode.DesktopTasksTransitionObserver;
 import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.EnterDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.ExitDesktopTaskTransitionHandler;
@@ -569,6 +570,18 @@
 
     @WMSingleton
     @Provides
+    static Optional<DesktopTasksTransitionObserver> provideDesktopTasksTransitionObserver(
+            Optional<DesktopModeTaskRepository> desktopModeTaskRepository,
+            Transitions transitions,
+            ShellInit shellInit
+    ) {
+        return desktopModeTaskRepository.flatMap(repository ->
+                Optional.of(new DesktopTasksTransitionObserver(repository, transitions, shellInit))
+        );
+    }
+
+    @WMSingleton
+    @Provides
     static DesktopModeLoggerTransitionObserver provideDesktopModeLoggerTransitionObserver(
             ShellInit shellInit,
             Transitions transitions,
@@ -623,7 +636,8 @@
     @Provides
     static Object provideIndependentShellComponentsToCreate(
             DragAndDropController dragAndDropController,
-            DefaultMixedHandler defaultMixedHandler) {
+            DefaultMixedHandler defaultMixedHandler,
+            Optional<DesktopTasksTransitionObserver> desktopTasksTransitionObserverOptional) {
         return new Object();
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt
index e1e41ee..f1a475a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt
@@ -36,7 +36,14 @@
                     true
                 }
             }
-
+            "moveToNextDisplay" -> {
+                if (!runMoveToNextDisplay(args, pw)) {
+                    pw.println("Task not found. Please enter a valid taskId.")
+                    false
+                } else {
+                    true
+                }
+            }
             else -> {
                 pw.println("Invalid command: ${args[0]}")
                 false
@@ -61,8 +68,28 @@
         return controller.moveToDesktop(taskId, WindowContainerTransaction())
     }
 
+    private fun runMoveToNextDisplay(args: Array<String>, pw: PrintWriter): Boolean {
+        if (args.size < 2) {
+            // First argument is the action name.
+            pw.println("Error: task id should be provided as arguments")
+            return false
+        }
+
+        val taskId = try {
+            args[1].toInt()
+        } catch (e: NumberFormatException) {
+            pw.println("Error: task id should be an integer")
+            return false
+        }
+
+        controller.moveToNextDisplay(taskId)
+        return true
+    }
+
     override fun printShellCommandHelp(pw: PrintWriter, prefix: String) {
         pw.println("$prefix moveToDesktop <taskId> ")
         pw.println("$prefix  Move a task with given id to desktop mode.")
+        pw.println("$prefix moveToNextDisplay <taskId> ")
+        pw.println("$prefix  Move a task with given id to next display.")
     }
 }
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
index 120d681..50cea01 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -22,6 +22,7 @@
 import android.util.ArraySet
 import android.util.SparseArray
 import android.view.Display.INVALID_DISPLAY
+import android.window.WindowContainerToken
 import androidx.core.util.forEach
 import androidx.core.util.keyIterator
 import androidx.core.util.valueIterator
@@ -49,6 +50,8 @@
         var stashed: Boolean = false
     )
 
+    // Token of the current wallpaper activity, used to remove it when the last task is removed
+    var wallpaperActivityToken: WindowContainerToken? = null
     // Tasks currently in freeform mode, ordered from top to bottom (top is at index 0).
     private val freeformTasksInZOrder = mutableListOf<Int>()
     private val activeTasksListeners = ArraySet<ActiveTasksListener>()
@@ -200,6 +203,15 @@
     }
 
     /**
+     *  Check if a task with the given [taskId] is the only active task on its display
+     */
+    fun isOnlyActiveTask(taskId: Int): Boolean {
+        return displayData.valueIterator().asSequence().any { data ->
+            data.activeTasks.singleOrNull() == taskId
+        }
+    }
+
+    /**
      * Get a set of the active tasks for given [displayId]
      */
     fun getActiveTasks(displayId: Int): ArraySet<Int> {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 487bbfb..068661a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -40,6 +40,7 @@
 import android.view.WindowManager.TRANSIT_CHANGE
 import android.view.WindowManager.TRANSIT_NONE
 import android.view.WindowManager.TRANSIT_OPEN
+import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.RemoteTransition
 import android.window.TransitionInfo
@@ -381,7 +382,6 @@
         )
         val wct = WindowContainerTransaction()
         exitSplitIfApplicable(wct, taskInfo)
-        moveHomeTaskToFront(wct)
         bringDesktopAppsToFront(taskInfo.displayId, wct)
         addMoveToDesktopChanges(wct, taskInfo)
         wct.setBounds(taskInfo.token, freeformBounds)
@@ -401,6 +401,22 @@
         shellTaskOrganizer.applyTransaction(wct)
     }
 
+    /**
+     * Perform clean up of the desktop wallpaper activity if the closed window task is
+     * the last active task.
+     *
+     * @param wct transaction to modify if the last active task is closed
+     * @param taskId task id of the window that's being closed
+     */
+    fun onDesktopWindowClose(
+        wct: WindowContainerTransaction,
+        taskId: Int
+    ) {
+        if (desktopModeTaskRepository.isOnlyActiveTask(taskId)) {
+            removeWallpaperActivity(wct)
+        }
+    }
+
     /** Move a task with given `taskId` to fullscreen */
     fun moveToFullscreen(taskId: Int) {
         shellTaskOrganizer.getRunningTaskInfo(taskId)?.let { task ->
@@ -676,9 +692,15 @@
         KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: bringDesktopAppsToFront")
         val activeTasks = desktopModeTaskRepository.getActiveTasks(displayId)
 
-        // First move home to front and then other tasks on top of it
-        moveHomeTaskToFront(wct)
+        if (Flags.enableDesktopWindowingWallpaperActivity()) {
+            // Add translucent wallpaper activity to show the wallpaper underneath
+            addWallpaperActivity(wct)
+        } else {
+            // Move home to front
+            moveHomeTaskToFront(wct)
+        }
 
+        // Then move other tasks on top of it
         val allTasksInZOrder = desktopModeTaskRepository.getFreeformTasksInZOrder()
         activeTasks
             // Sort descending as the top task is at index 0. It should be ordered to top last
@@ -694,6 +716,26 @@
             ?.let { homeTask -> wct.reorder(homeTask.getToken(), true /* onTop */) }
     }
 
+    private fun addWallpaperActivity(wct: WindowContainerTransaction) {
+        KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: addWallpaper")
+        val intent = Intent(context, DesktopWallpaperActivity::class.java)
+        val options = ActivityOptions.makeBasic().apply {
+            isPendingIntentBackgroundActivityLaunchAllowedByPermission = true
+            pendingIntentBackgroundActivityStartMode =
+                ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED
+        }
+        val pendingIntent = PendingIntent.getActivity(context, /* requestCode = */ 0, intent,
+            PendingIntent.FLAG_IMMUTABLE)
+        wct.sendPendingIntent(pendingIntent, intent, options.toBundle())
+    }
+
+    private fun removeWallpaperActivity(wct: WindowContainerTransaction) {
+        desktopModeTaskRepository.wallpaperActivityToken?.let { token ->
+            KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: removeWallpaper")
+            wct.removeTask(token)
+        }
+    }
+
     fun releaseVisualIndicator() {
         val t = SurfaceControl.Transaction()
         visualIndicator?.releaseVisualIndicator(t)
@@ -741,6 +783,9 @@
                     reason = "recents animation is running"
                     false
                 }
+                // Handle back navigation for the last window if wallpaper available
+                shouldRemoveWallpaper(request) ->
+                    true
                 // Only handle open or to front transitions
                 request.type != TRANSIT_OPEN && request.type != TRANSIT_TO_FRONT -> {
                     reason = "transition type not handled (${request.type})"
@@ -777,6 +822,7 @@
 
         val result = triggerTask?.let { task ->
             when {
+                request.type == TRANSIT_TO_BACK -> handleBackNavigation(task)
                 // If display has tasks stashed, handle as stashed launch
                 task.isStashed -> handleStashedTaskLaunch(task)
                 // Check if the task has a top transparent activity
@@ -824,6 +870,14 @@
         return Flags.enableDesktopWindowingModalsPolicy() && isSingleTopActivityTranslucent(task)
     }
 
+    private fun shouldRemoveWallpaper(request: TransitionRequestInfo): Boolean {
+        return Flags.enableDesktopWindowingWallpaperActivity() &&
+                request.type == TRANSIT_TO_BACK &&
+                request.triggerTask?.let { task ->
+                    desktopModeTaskRepository.isOnlyActiveTask(task.taskId)
+                } ?: false
+    }
+
     private fun handleFreeformTaskLaunch(task: RunningTaskInfo): WindowContainerTransaction? {
         KtProtoLog.v(WM_SHELL_DESKTOP_MODE, "DesktopTasksController: handleFreeformTaskLaunch")
         val activeTasks = desktopModeTaskRepository.getActiveTasks(task.displayId)
@@ -881,6 +935,19 @@
         }
     }
 
+    /** Handle back navigation by removing wallpaper activity if it's the last active task */
+    private fun handleBackNavigation(task: RunningTaskInfo): WindowContainerTransaction? {
+        if (desktopModeTaskRepository.isOnlyActiveTask(task.taskId) &&
+            desktopModeTaskRepository.wallpaperActivityToken != null) {
+            // Remove wallpaper activity when the last active task is removed
+            return WindowContainerTransaction().also { wct ->
+                removeWallpaperActivity(wct)
+            }
+        } else {
+            return null
+        }
+    }
+
     private fun addMoveToDesktopChanges(
         wct: WindowContainerTransaction,
         taskInfo: RunningTaskInfo
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
new file mode 100644
index 0000000..20df264
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.desktopmode
+
+import android.os.IBinder
+import android.view.SurfaceControl
+import android.view.WindowManager
+import android.window.TransitionInfo
+import com.android.window.flags.Flags.enableDesktopWindowingWallpaperActivity
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.transition.Transitions
+import com.android.wm.shell.util.KtProtoLog
+
+/**
+ * A [Transitions.TransitionObserver] that observes shell transitions and updates
+ * the [DesktopModeTaskRepository] state TODO: b/332682201
+ * This observes transitions related to desktop mode
+ * and other transitions that originate both within and outside shell.
+ */
+class DesktopTasksTransitionObserver(
+    private val desktopModeTaskRepository: DesktopModeTaskRepository,
+    private val transitions: Transitions,
+    shellInit: ShellInit
+) : Transitions.TransitionObserver {
+
+    init {
+        if (Transitions.ENABLE_SHELL_TRANSITIONS && DesktopModeStatus.isEnabled()) {
+            shellInit.addInitCallback(::onInit, this)
+        }
+    }
+
+    fun onInit() {
+        KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTasksTransitionObserver: onInit")
+        transitions.registerObserver(this)
+    }
+
+    override fun onTransitionReady(
+        transition: IBinder,
+        info: TransitionInfo,
+        startTransaction: SurfaceControl.Transaction,
+        finishTransaction: SurfaceControl.Transaction
+    ) {
+        // TODO: b/332682201 Update repository state
+        updateWallpaperToken(info)
+    }
+
+    override fun onTransitionStarting(transition: IBinder) {
+        // TODO: b/332682201 Update repository state
+    }
+
+    override fun onTransitionMerged(merged: IBinder, playing: IBinder) {
+        // TODO: b/332682201 Update repository state
+    }
+
+    override fun onTransitionFinished(transition: IBinder, aborted: Boolean) {
+        // TODO: b/332682201 Update repository state
+    }
+
+    private fun updateWallpaperToken(info: TransitionInfo) {
+        if (!enableDesktopWindowingWallpaperActivity()) {
+            return
+        }
+        info.changes.forEach { change ->
+            change.taskInfo?.let { taskInfo ->
+                if (DesktopWallpaperActivity.isWallpaperTask(taskInfo)) {
+                    when (change.mode) {
+                        WindowManager.TRANSIT_OPEN ->
+                            desktopModeTaskRepository.wallpaperActivityToken = taskInfo.token
+                        WindowManager.TRANSIT_CLOSE ->
+                            desktopModeTaskRepository.wallpaperActivityToken = null
+                        else -> {}
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopWallpaperActivity.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopWallpaperActivity.kt
new file mode 100644
index 0000000..c4a4474
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopWallpaperActivity.kt
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.desktopmode
+
+import android.app.Activity
+import android.app.ActivityManager
+import android.content.ComponentName
+import android.os.Bundle
+import android.view.WindowManager
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.util.KtProtoLog
+
+/**
+ * A transparent activity used in the desktop mode to show the wallpaper under the freeform windows.
+ * This activity will be running in `FULLSCREEN` windowing mode, which ensures it hides Launcher.
+ * When entering desktop, we would ensure that it's added behind desktop apps and removed when
+ * leaving the desktop mode.
+ *
+ * Note! This activity should NOT interact directly with any other code in the Shell without calling
+ * onto the shell main thread. Activities are always started on the main thread.
+ */
+class DesktopWallpaperActivity : Activity() {
+
+    override fun onCreate(savedInstanceState: Bundle?) {
+        KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopWallpaperActivity: onCreate")
+        super.onCreate(savedInstanceState)
+        window.addFlags(WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE)
+    }
+
+    companion object {
+        private const val SYSTEM_UI_PACKAGE_NAME = "com.android.systemui"
+        private val wallpaperActivityComponent =
+            ComponentName(SYSTEM_UI_PACKAGE_NAME, DesktopWallpaperActivity::class.java.name)
+
+        @JvmStatic
+        fun isWallpaperTask(taskInfo: ActivityManager.RunningTaskInfo) =
+            taskInfo.baseIntent.component?.let(::isWallpaperComponent) ?: false
+
+        @JvmStatic
+        fun isWallpaperComponent(component: ComponentName) =
+            component == wallpaperActivityComponent
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java
index c9185ae..b1a1e59 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/HomeTransitionObserver.java
@@ -20,6 +20,7 @@
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.window.TransitionInfo.FLAG_BACK_GESTURE_ANIMATED;
 
+import static com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP;
 import static com.android.wm.shell.transition.Transitions.TransitionObserver;
 
 import android.annotation.NonNull;
@@ -60,7 +61,8 @@
             @NonNull SurfaceControl.Transaction finishTransaction) {
         for (TransitionInfo.Change change : info.getChanges()) {
             final ActivityManager.RunningTaskInfo taskInfo = change.getTaskInfo();
-            if (taskInfo == null
+            if (info.getType() == TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP
+                    || taskInfo == null
                     || taskInfo.displayId != DEFAULT_DISPLAY
                     || taskInfo.taskId == -1
                     || !taskInfo.isRunning) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index 7649a782..777ab9c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -84,6 +84,7 @@
 import com.android.wm.shell.desktopmode.DesktopModeVisualIndicator;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition;
+import com.android.wm.shell.desktopmode.DesktopWallpaperActivity;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
 import com.android.wm.shell.splitscreen.SplitScreen;
 import com.android.wm.shell.splitscreen.SplitScreen.StageType;
@@ -407,7 +408,9 @@
                     mSplitScreenController.moveTaskToFullscreen(getOtherSplitTask(mTaskId).taskId,
                             SplitScreenController.EXIT_REASON_DESKTOP_MODE);
                 } else {
-                    mTaskOperations.closeTask(mTaskToken);
+                    WindowContainerTransaction wct = new WindowContainerTransaction();
+                    mDesktopTasksController.onDesktopWindowClose(wct, mTaskId);
+                    mTaskOperations.closeTask(mTaskToken, wct);
                 }
             } else if (id == R.id.back_button) {
                 mTaskOperations.injectBackKey();
@@ -1005,6 +1008,7 @@
             return false;
         }
         return DesktopModeStatus.isEnabled()
+                && !DesktopWallpaperActivity.isWallpaperTask(taskInfo)
                 && taskInfo.getWindowingMode() != WINDOWING_MODE_PINNED
                 && taskInfo.getActivityType() == ACTIVITY_TYPE_STANDARD
                 && !taskInfo.configuration.windowConfiguration.isAlwaysOnTop()
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java
index 1763e4f..53d4e27 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java
@@ -72,7 +72,10 @@
     }
 
     void closeTask(WindowContainerToken taskToken) {
-        WindowContainerTransaction wct = new WindowContainerTransaction();
+        closeTask(taskToken, new WindowContainerTransaction());
+    }
+
+    void closeTask(WindowContainerToken taskToken, WindowContainerTransaction wct) {
         wct.removeTask(taskToken);
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             mTransitionStarter.startRemoveTransition(wct);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java
index 3672ae3..24f4d92 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java
@@ -23,8 +23,10 @@
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 
+import android.annotation.NonNull;
 import android.app.ActivityManager;
 import android.app.WindowConfiguration;
+import android.content.Intent;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.IBinder;
@@ -38,6 +40,7 @@
 
     private WindowContainerToken mToken = createMockWCToken();
     private int mParentTaskId = INVALID_TASK_ID;
+    private Intent mBaseIntent = new Intent();
     private @WindowConfiguration.ActivityType int mActivityType = ACTIVITY_TYPE_STANDARD;
     private @WindowConfiguration.WindowingMode int mWindowingMode = WINDOWING_MODE_UNDEFINED;
     private int mDisplayId = Display.DEFAULT_DISPLAY;
@@ -68,6 +71,15 @@
         return this;
     }
 
+    /**
+     * Set {@link ActivityManager.RunningTaskInfo#baseIntent} for the task info, by default
+     * an empty intent is assigned
+     */
+    public TestRunningTaskInfoBuilder setBaseIntent(@NonNull Intent intent) {
+        mBaseIntent = intent;
+        return this;
+    }
+
     public TestRunningTaskInfoBuilder setActivityType(
             @WindowConfiguration.ActivityType int activityType) {
         mActivityType = activityType;
@@ -109,6 +121,7 @@
     public ActivityManager.RunningTaskInfo build() {
         final ActivityManager.RunningTaskInfo info = new ActivityManager.RunningTaskInfo();
         info.taskId = sNextTaskId++;
+        info.baseIntent = mBaseIntent;
         info.parentTaskId = mParentTaskId;
         info.displayId = mDisplayId;
         info.configuration.windowConfiguration.setBounds(mBounds);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
index 0c45d52..b2b54ac 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
@@ -119,6 +119,57 @@
     }
 
     @Test
+    fun isOnlyActiveTask_noActiveTasks() {
+        // Not an active task
+        assertThat(repo.isOnlyActiveTask(1)).isFalse()
+    }
+
+    @Test
+    fun isOnlyActiveTask_singleActiveTask() {
+        repo.addActiveTask(DEFAULT_DISPLAY, 1)
+        // The only active task
+        assertThat(repo.isActiveTask(1)).isTrue()
+        assertThat(repo.isOnlyActiveTask(1)).isTrue()
+        // Not an active task
+        assertThat(repo.isActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveTask(99)).isFalse()
+    }
+
+    @Test
+    fun isOnlyActiveTask_multipleActiveTasks() {
+        repo.addActiveTask(DEFAULT_DISPLAY, 1)
+        repo.addActiveTask(DEFAULT_DISPLAY, 2)
+        // Not the only task
+        assertThat(repo.isActiveTask(1)).isTrue()
+        assertThat(repo.isOnlyActiveTask(1)).isFalse()
+        // Not the only task
+        assertThat(repo.isActiveTask(2)).isTrue()
+        assertThat(repo.isOnlyActiveTask(2)).isFalse()
+        // Not an active task
+        assertThat(repo.isActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveTask(99)).isFalse()
+    }
+
+    @Test
+    fun isOnlyActiveTask_multipleDisplays() {
+        repo.addActiveTask(DEFAULT_DISPLAY, 1)
+        repo.addActiveTask(DEFAULT_DISPLAY, 2)
+        repo.addActiveTask(SECOND_DISPLAY, 3)
+        // Not the only task on DEFAULT_DISPLAY
+        assertThat(repo.isActiveTask(1)).isTrue()
+        assertThat(repo.isOnlyActiveTask(1)).isFalse()
+        // Not the only task on DEFAULT_DISPLAY
+        assertThat(repo.isActiveTask(2)).isTrue()
+        assertThat(repo.isOnlyActiveTask(2)).isFalse()
+        // The only active task on SECOND_DISPLAY
+        assertThat(repo.isActiveTask(3)).isTrue()
+        assertThat(repo.isOnlyActiveTask(3)).isTrue()
+        // Not an active task
+        assertThat(repo.isActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveTask(99)).isFalse()
+    }
+
+    @Test
     fun addListener_notifiesVisibleFreeformTask() {
         repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
         val listener = TestVisibilityListener()
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index 4d0f11b..64f6041 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -23,10 +23,13 @@
 import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
 import android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW
 import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
+import android.content.Intent
 import android.graphics.Point
 import android.graphics.PointF
 import android.graphics.Rect
 import android.os.Binder
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
 import android.testing.AndroidTestingRunner
 import android.view.Display.DEFAULT_DISPLAY
@@ -34,11 +37,15 @@
 import android.view.WindowManager
 import android.view.WindowManager.TRANSIT_CHANGE
 import android.view.WindowManager.TRANSIT_OPEN
+import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.DisplayAreaInfo
 import android.window.RemoteTransition
 import android.window.TransitionRequestInfo
+import android.window.WindowContainerToken
 import android.window.WindowContainerTransaction
+import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_PENDING_INTENT
+import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REMOVE_TASK
 import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER
 import androidx.test.filters.SmallTest
 import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
@@ -223,7 +230,8 @@
     }
 
     @Test
-    fun showDesktopApps_allAppsInvisible_bringsToFront() {
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_allAppsInvisible_bringsToFront_desktopWallpaperDisabled() {
         val homeTask = setUpHomeTask()
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
@@ -242,7 +250,27 @@
     }
 
     @Test
-    fun showDesktopApps_appsAlreadyVisible_bringsToFront() {
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_allAppsInvisible_bringsToFront_desktopWallpaperEnabled() {
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskHidden(task1)
+        markTaskHidden(task2)
+
+        controller.showDesktopApps(DEFAULT_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+        val wct =
+            getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+        assertThat(wct.hierarchyOps).hasSize(3)
+        // Expect order to be from bottom: wallpaper intent, task1, task2
+        wct.assertPendingIntentAt(index = 0, desktopWallpaperIntent)
+        wct.assertReorderAt(index = 1, task1)
+        wct.assertReorderAt(index = 2, task2)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_appsAlreadyVisible_bringsToFront_desktopWallpaperDisabled() {
         val homeTask = setUpHomeTask()
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
@@ -261,7 +289,27 @@
     }
 
     @Test
-    fun showDesktopApps_someAppsInvisible_reordersAll() {
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_appsAlreadyVisible_bringsToFront_desktopWallpaperEnabled() {
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskVisible(task1)
+        markTaskVisible(task2)
+
+        controller.showDesktopApps(DEFAULT_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+        val wct =
+            getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+        assertThat(wct.hierarchyOps).hasSize(3)
+        // Expect order to be from bottom: wallpaper intent, task1, task2
+        wct.assertPendingIntentAt(index = 0, desktopWallpaperIntent)
+        wct.assertReorderAt(index = 1, task1)
+        wct.assertReorderAt(index = 2, task2)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_someAppsInvisible_reordersAll_desktopWallpaperDisabled() {
         val homeTask = setUpHomeTask()
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
@@ -280,7 +328,27 @@
     }
 
     @Test
-    fun showDesktopApps_noActiveTasks_reorderHomeToTop() {
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_someAppsInvisible_reordersAll_desktopWallpaperEnabled() {
+        val task1 = setUpFreeformTask()
+        val task2 = setUpFreeformTask()
+        markTaskHidden(task1)
+        markTaskVisible(task2)
+
+        controller.showDesktopApps(DEFAULT_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+        val wct =
+            getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+        assertThat(wct.hierarchyOps).hasSize(3)
+        // Expect order to be from bottom: wallpaper intent, task1, task2
+        wct.assertPendingIntentAt(index = 0, desktopWallpaperIntent)
+        wct.assertReorderAt(index = 1, task1)
+        wct.assertReorderAt(index = 2, task2)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_noActiveTasks_reorderHomeToTop_desktopWallpaperDisabled() {
         val homeTask = setUpHomeTask()
 
         controller.showDesktopApps(DEFAULT_DISPLAY, RemoteTransition(TestRemoteTransition()))
@@ -292,7 +360,18 @@
     }
 
     @Test
-    fun showDesktopApps_twoDisplays_bringsToFrontOnlyOneDisplay() {
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_noActiveTasks_addDesktopWallpaper_desktopWallpaperEnabled() {
+        controller.showDesktopApps(DEFAULT_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+        val wct =
+            getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+        wct.assertPendingIntentAt(index = 0, desktopWallpaperIntent)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_twoDisplays_bringsToFrontOnlyOneDisplay_desktopWallpaperDisabled() {
         val homeTaskDefaultDisplay = setUpHomeTask(DEFAULT_DISPLAY)
         val taskDefaultDisplay = setUpFreeformTask(DEFAULT_DISPLAY)
         setUpHomeTask(SECOND_DISPLAY)
@@ -311,6 +390,25 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun showDesktopApps_twoDisplays_bringsToFrontOnlyOneDisplay_desktopWallpaperEnabled() {
+        val taskDefaultDisplay = setUpFreeformTask(DEFAULT_DISPLAY)
+        setUpHomeTask(SECOND_DISPLAY)
+        val taskSecondDisplay = setUpFreeformTask(SECOND_DISPLAY)
+        markTaskHidden(taskDefaultDisplay)
+        markTaskHidden(taskSecondDisplay)
+
+        controller.showDesktopApps(DEFAULT_DISPLAY, RemoteTransition(TestRemoteTransition()))
+
+        val wct =
+            getLatestWct(type = TRANSIT_TO_FRONT, handlerClass = OneShotRemoteHandler::class.java)
+        assertThat(wct.hierarchyOps).hasSize(2)
+        // Expect order to be from bottom: wallpaper intent, task
+        wct.assertPendingIntentAt(index = 0, desktopWallpaperIntent)
+        wct.assertReorderAt(index = 1, taskDefaultDisplay)
+    }
+
+    @Test
     fun getVisibleTaskCount_noTasks_returnsZero() {
         assertThat(controller.getVisibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
     }
@@ -419,7 +517,8 @@
     }
 
     @Test
-    fun moveToDesktop_otherFreeformTasksBroughtToFront() {
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun moveToDesktop_otherFreeformTasksBroughtToFront_desktopWallpaperDisabled() {
         val homeTask = setUpHomeTask()
         val freeformTask = setUpFreeformTask()
         val fullscreenTask = setUpFullscreenTask()
@@ -437,6 +536,26 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun moveToDesktop_otherFreeformTasksBroughtToFront_desktopWallpaperEnabled() {
+        val freeformTask = setUpFreeformTask()
+        val fullscreenTask = setUpFullscreenTask()
+        markTaskHidden(freeformTask)
+
+        controller.moveToDesktop(fullscreenTask)
+
+        with(getLatestMoveToDesktopWct()) {
+            // Operations should include wallpaper intent, freeform task, fullscreen task
+            assertThat(hierarchyOps).hasSize(3)
+            assertPendingIntentAt(index = 0, desktopWallpaperIntent)
+            assertReorderAt(index = 1, freeformTask)
+            assertReorderAt(index = 2, fullscreenTask)
+            assertThat(changes[fullscreenTask.token.asBinder()]?.windowingMode)
+                .isEqualTo(WINDOWING_MODE_FREEFORM)
+        }
+    }
+
+    @Test
     fun moveToDesktop_onlyFreeformTasksFromCurrentDisplayBroughtToFront() {
         setUpHomeTask(displayId = DEFAULT_DISPLAY)
         val freeformTaskDefault = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
@@ -603,6 +722,48 @@
     }
 
     @Test
+    fun onDesktopWindowClose_noActiveTasks() {
+        val wct = WindowContainerTransaction()
+        controller.onDesktopWindowClose(wct, 1 /* taskId */)
+        // Doesn't modify transaction
+        assertThat(wct.hierarchyOps).isEmpty()
+    }
+
+    @Test
+    fun onDesktopWindowClose_singleActiveTask_noWallpaperActivityToken() {
+        val task = setUpFreeformTask()
+        val wct = WindowContainerTransaction()
+        controller.onDesktopWindowClose(wct, task.taskId)
+        // Doesn't modify transaction
+        assertThat(wct.hierarchyOps).isEmpty()
+    }
+
+    @Test
+    fun onDesktopWindowClose_singleActiveTask_hasWallpaperActivityToken() {
+        val task = setUpFreeformTask()
+        val wallpaperToken = MockToken().token()
+        desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+
+        val wct = WindowContainerTransaction()
+        controller.onDesktopWindowClose(wct, task.taskId)
+        // Adds remove wallpaper operation
+        wct.assertRemoveAt(index = 0, wallpaperToken)
+    }
+
+    @Test
+    fun onDesktopWindowClose_multipleActiveTasks() {
+        val task1 = setUpFreeformTask()
+        setUpFreeformTask()
+        val wallpaperToken = MockToken().token()
+        desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+
+        val wct = WindowContainerTransaction()
+        controller.onDesktopWindowClose(wct, task1.taskId)
+        // Doesn't modify transaction
+        assertThat(wct.hierarchyOps).isEmpty()
+    }
+
+    @Test
     fun handleRequest_fullscreenTask_freeformVisible_returnSwitchToFreeformWCT() {
         assumeTrue(ENABLE_SHELL_TRANSITIONS)
 
@@ -646,6 +807,7 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_fullscreenTask_desktopStashed_returnWCTWithAllAppsBroughtToFront() {
         assumeTrue(ENABLE_SHELL_TRANSITIONS)
         whenever(DesktopModeStatus.isStashingEnabled()).thenReturn(true)
@@ -718,6 +880,7 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_freeformTask_desktopStashed_returnWCTWithAllAppsBroughtToFront() {
         assumeTrue(ENABLE_SHELL_TRANSITIONS)
         whenever(DesktopModeStatus.isStashingEnabled()).thenReturn(true)
@@ -804,6 +967,55 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun handleRequest_backTransition_singleActiveTask_noToken() {
+        val task = setUpFreeformTask()
+        val result =
+            controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
+        // Doesn't handle request
+        assertThat(result).isNull()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun handleRequest_backTransition_singleActiveTask_hasToken_desktopWallpaperDisabled() {
+        desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+
+        val task = setUpFreeformTask()
+        val result =
+            controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
+        // Doesn't handle request
+        assertThat(result).isNull()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun handleRequest_backTransition_singleActiveTask_hasToken_desktopWallpaperEnabled() {
+        val wallpaperToken = MockToken().token()
+        desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+
+        val task = setUpFreeformTask()
+        val result =
+            controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
+        assertThat(result).isNotNull()
+        // Creates remove wallpaper transaction
+        result!!.assertRemoveAt(index = 0, wallpaperToken)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    fun handleRequest_backTransition_multipleActiveTasks() {
+        desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+
+        val task1 = setUpFreeformTask()
+        setUpFreeformTask()
+        val result =
+            controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
+        // Doesn't handle request
+        assertThat(result).isNull()
+    }
+
+    @Test
     fun stashDesktopApps_stateUpdates() {
         whenever(DesktopModeStatus.isStashingEnabled()).thenReturn(true)
 
@@ -1006,6 +1218,9 @@
         assertThat(desktopModeTaskRepository.removeBoundsBeforeMaximize(task.taskId)).isNull()
     }
 
+    private val desktopWallpaperIntent: Intent
+        get() = Intent(context, DesktopWallpaperActivity::class.java)
+
     private fun setUpFreeformTask(
             displayId: Int = DEFAULT_DISPLAY,
             bounds: Rect? = null
@@ -1131,10 +1346,14 @@
     }
 }
 
-private fun WindowContainerTransaction.assertReorderAt(index: Int, task: RunningTaskInfo) {
+private fun WindowContainerTransaction.assertIndexInBounds(index: Int) {
     assertWithMessage("WCT does not have a hierarchy operation at index $index")
         .that(hierarchyOps.size)
         .isGreaterThan(index)
+}
+
+private fun WindowContainerTransaction.assertReorderAt(index: Int, task: RunningTaskInfo) {
+    assertIndexInBounds(index)
     val op = hierarchyOps[index]
     assertThat(op.type).isEqualTo(HIERARCHY_OP_TYPE_REORDER)
     assertThat(op.container).isEqualTo(task.token.asBinder())
@@ -1145,3 +1364,17 @@
         assertReorderAt(i, tasks[i])
     }
 }
+
+private fun WindowContainerTransaction.assertRemoveAt(index: Int, token: WindowContainerToken) {
+    assertIndexInBounds(index)
+    val op = hierarchyOps[index]
+    assertThat(op.type).isEqualTo(HIERARCHY_OP_TYPE_REMOVE_TASK)
+    assertThat(op.container).isEqualTo(token.asBinder())
+}
+
+private fun WindowContainerTransaction.assertPendingIntentAt(index: Int, intent: Intent) {
+    assertIndexInBounds(index)
+    val op = hierarchyOps[index]
+    assertThat(op.type).isEqualTo(HIERARCHY_OP_TYPE_PENDING_INTENT)
+    assertThat(op.pendingIntent?.intent?.component).isEqualTo(intent.component)
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java
index e7d37ad..0db10ef 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/HomeTransitionObserverTest.java
@@ -24,6 +24,7 @@
 import static android.window.TransitionInfo.FLAG_BACK_GESTURE_ANIMATED;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+import static com.android.wm.shell.transition.Transitions.TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP;
 
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.Mockito.mock;
@@ -167,6 +168,25 @@
     }
 
     @Test
+    public void testStartDragToDesktopDoesNotTriggerCallback() throws RemoteException {
+        TransitionInfo info = mock(TransitionInfo.class);
+        TransitionInfo.Change change = mock(TransitionInfo.Change.class);
+        ActivityManager.RunningTaskInfo taskInfo = mock(ActivityManager.RunningTaskInfo.class);
+        when(change.getTaskInfo()).thenReturn(taskInfo);
+        when(info.getChanges()).thenReturn(new ArrayList<>(List.of(change)));
+        when(info.getType()).thenReturn(TRANSIT_DESKTOP_MODE_START_DRAG_TO_DESKTOP);
+
+        setupTransitionInfo(taskInfo, change, ACTIVITY_TYPE_HOME, TRANSIT_OPEN, true);
+
+        mHomeTransitionObserver.onTransitionReady(mock(IBinder.class),
+                info,
+                mock(SurfaceControl.Transaction.class),
+                mock(SurfaceControl.Transaction.class));
+
+        verify(mListener, times(0)).onHomeVisibilityChanged(anyBoolean());
+    }
+
+    @Test
     public void testHomeActivityWithBackGestureNotifiesHomeIsVisible() throws RemoteException {
         TransitionInfo info = mock(TransitionInfo.class);
         TransitionInfo.Change change = mock(TransitionInfo.Change.class);
diff --git a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
index c836df3..c3edb90 100644
--- a/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
+++ b/media/tests/MediaRouter/src/com/android/mediaroutertest/MediaRouter2ManagerTest.java
@@ -361,6 +361,7 @@
      * Tests if MR2.SessionCallback.onSessionCreated is called
      * when a route is selected from MR2Manager.
      */
+    @Ignore // Ignored due to flakiness. No plans to fix though, in favor of removal (b/334970551).
     @Test
     public void testRouterOnSessionCreated() throws Exception {
         Map<String, MediaRoute2Info> routes = waitAndGetRoutesWithManager(FEATURES_ALL);
@@ -908,6 +909,7 @@
      * Tests if getSelectableRoutes and getDeselectableRoutes filter routes based on
      * selected routes
      */
+    @Ignore // Ignored due to flakiness. No plans to fix though, in favor of removal (b/334970551).
     @Test
     public void testGetSelectableRoutes_notReturnsSelectedRoutes() throws Exception {
         Map<String, MediaRoute2Info> routes = waitAndGetRoutesWithManager(FEATURES_ALL);
diff --git a/packages/SettingsLib/DataStore/Android.bp b/packages/SettingsLib/DataStore/Android.bp
index 9fafcab..86c8f0da 100644
--- a/packages/SettingsLib/DataStore/Android.bp
+++ b/packages/SettingsLib/DataStore/Android.bp
@@ -2,12 +2,17 @@
     default_applicable_licenses: ["frameworks_base_license"],
 }
 
+filegroup {
+    name: "SettingsLibDataStore-srcs",
+    srcs: ["src/**/*"],
+}
+
 android_library {
     name: "SettingsLibDataStore",
     defaults: [
         "SettingsLintDefaults",
     ],
-    srcs: ["src/**/*"],
+    srcs: [":SettingsLibDataStore-srcs"],
     static_libs: [
         "androidx.annotation_annotation",
         "androidx.collection_collection-ktx",
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
index 9d3fb66..7644bc9 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
@@ -19,6 +19,7 @@
 import android.app.backup.BackupDataInputStream
 import android.content.Context
 import android.util.Log
+import androidx.annotation.VisibleForTesting
 import java.io.File
 import java.io.InputStream
 import java.io.OutputStream
@@ -33,11 +34,9 @@
  */
 internal class BackupRestoreFileArchiver(
     private val context: Context,
-    private val fileStorages: List<BackupRestoreFileStorage>,
+    @get:VisibleForTesting internal val fileStorages: List<BackupRestoreFileStorage>,
+    override val name: String,
 ) : BackupRestoreStorage() {
-    override val name: String
-        get() = "file_archiver"
-
     override fun createBackupRestoreEntities(): List<BackupRestoreEntity> =
         fileStorages.map { it.toBackupRestoreEntity() }
 
@@ -88,7 +87,8 @@
     }
 }
 
-private fun BackupRestoreFileStorage.toBackupRestoreEntity() =
+@VisibleForTesting
+internal fun BackupRestoreFileStorage.toBackupRestoreEntity() =
     object : BackupRestoreEntity {
         override val key: String
             get() = storageFilePath
@@ -107,7 +107,7 @@
                 Log.i(LOG_TAG, "[$name] $key not exist")
                 return EntityBackupResult.DELETE
             }
-            val codec = codec() ?: defaultCodec()
+            val codec = defaultCodec()
             // MUST close to flush the data
             wrapBackupOutputStream(codec, outputStream).use { stream ->
                 val bytesCopied = file.inputStream().use { it.copyTo(stream) }
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
index c4c00cb..935f9cc 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
@@ -22,6 +22,7 @@
 import android.app.backup.BackupHelper
 import android.os.ParcelFileDescriptor
 import android.util.Log
+import androidx.annotation.VisibleForTesting
 import androidx.collection.MutableScatterMap
 import com.google.common.io.ByteStreams
 import java.io.ByteArrayOutputStream
@@ -60,10 +61,11 @@
      *
      * Map key is the entity key, map value is the checksum of backup data.
      */
-    protected val entityStates = MutableScatterMap<String, Long>()
+    @VisibleForTesting(otherwise = VisibleForTesting.PROTECTED)
+    val entityStates = MutableScatterMap<String, Long>()
 
     /** Entities created by [createBackupRestoreEntities]. This field is for restore only. */
-    private var entities: List<BackupRestoreEntity>? = null
+    @VisibleForTesting internal var entities: List<BackupRestoreEntity>? = null
 
     /** Entities to back up and restore. */
     abstract fun createBackupRestoreEntities(): List<BackupRestoreEntity>
@@ -76,7 +78,7 @@
         data: BackupDataOutput,
         newState: ParcelFileDescriptor,
     ) {
-        oldState.readEntityStates(entityStates)
+        readEntityStates(oldState, entityStates)
         val backupContext = BackupContext(data)
         if (!enableBackup(backupContext)) {
             Log.i(LOG_TAG, "[$name] Backup disabled")
@@ -94,7 +96,10 @@
             val codec = entity.codec() ?: defaultCodec()
             val result =
                 try {
-                    entity.backup(backupContext, wrapBackupOutputStream(codec, checkedOutputStream))
+                    // MUST close to flush all data
+                    wrapBackupOutputStream(codec, checkedOutputStream).use {
+                        entity.backup(backupContext, it)
+                    }
                 } catch (exception: Exception) {
                     Log.e(LOG_TAG, "[$name] Fail to backup entity $key", exception)
                     continue
@@ -191,9 +196,13 @@
     /** Callbacks when restore finished. */
     open fun onRestoreFinished() {}
 
-    private fun ParcelFileDescriptor?.readEntityStates(state: MutableScatterMap<String, Long>) {
+    @VisibleForTesting
+    internal fun readEntityStates(
+        parcelFileDescriptor: ParcelFileDescriptor?,
+        state: MutableScatterMap<String, Long>,
+    ) {
         state.clear()
-        if (this == null) return
+        val fileDescriptor = parcelFileDescriptor?.fileDescriptor ?: return
         // do not close the streams
         val fileInputStream = FileInputStream(fileDescriptor)
         val dataInputStream = DataInputStream(fileInputStream)
@@ -233,6 +242,7 @@
                 dataOutputStream.writeUTF(key)
                 dataOutputStream.writeLong(value)
             }
+            dataOutputStream.flush()
         } catch (exception: Exception) {
             Log.e(LOG_TAG, "[$name] Fail to write state file", exception)
         }
@@ -241,7 +251,7 @@
     }
 
     companion object {
-        private const val STATE_VERSION: Byte = 0
+        internal const val STATE_VERSION: Byte = 0
 
         /** Checksum for entity backup data. */
         fun createChecksum(): Checksum = CRC32()
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
index cfdcaff..8242347 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
@@ -21,23 +21,32 @@
 import android.app.backup.BackupManager
 import android.content.Context
 import android.util.Log
+import androidx.annotation.VisibleForTesting
 import com.google.common.util.concurrent.MoreExecutors
 import java.util.concurrent.ConcurrentHashMap
 
 /** Manager of [BackupRestoreStorage]. */
 class BackupRestoreStorageManager private constructor(private val application: Application) {
-    private val storageWrappers = ConcurrentHashMap<String, StorageWrapper>()
+    @VisibleForTesting internal val storageWrappers = ConcurrentHashMap<String, StorageWrapper>()
 
     private val executor = MoreExecutors.directExecutor()
 
     /**
      * Adds all the registered [BackupRestoreStorage] as the helpers of given [BackupAgentHelper].
      *
-     * All [BackupRestoreFileStorage]s will be wrapped as a single [BackupRestoreFileArchiver].
+     * All [BackupRestoreFileStorage]s will be wrapped as a single [BackupRestoreFileArchiver],
+     * specify [fileArchiverName] to avoid key prefix conflict if needed.
      *
+     * @param backupAgentHelper backup agent helper to add helpers
+     * @param fileArchiverName key prefix of the [BackupRestoreFileArchiver], the value must not be
+     *   changed in future
      * @see BackupAgentHelper.addHelper
      */
-    fun addBackupAgentHelpers(backupAgentHelper: BackupAgentHelper) {
+    @JvmOverloads
+    fun addBackupAgentHelpers(
+        backupAgentHelper: BackupAgentHelper,
+        fileArchiverName: String = "file_archiver",
+    ) {
         val fileStorages = mutableListOf<BackupRestoreFileStorage>()
         for ((keyPrefix, storageWrapper) in storageWrappers) {
             val storage = storageWrapper.storage
@@ -48,7 +57,7 @@
             }
         }
         // Always add file archiver even fileStorages is empty to handle forward compatibility
-        val fileArchiver = BackupRestoreFileArchiver(application, fileStorages)
+        val fileArchiver = BackupRestoreFileArchiver(application, fileStorages, fileArchiverName)
         backupAgentHelper.addHelper(fileArchiver.name, fileArchiver)
     }
 
@@ -106,7 +115,8 @@
     /** Returns storage with given name, exception is raised if not found. */
     fun getOrThrow(name: String): BackupRestoreStorage = storageWrappers[name]!!.storage
 
-    private inner class StorageWrapper(val storage: BackupRestoreStorage) :
+    @VisibleForTesting
+    internal inner class StorageWrapper(val storage: BackupRestoreStorage) :
         Observer, KeyedObserver<Any?> {
         init {
             when (storage) {
@@ -139,7 +149,7 @@
                 LOG_TAG,
                 "Notify BackupManager dataChanged: storage=$name key=$key reason=$reason"
             )
-            BackupManager.dataChanged(application.packageName)
+            BackupManager(application).dataChanged()
         }
 
         fun removeObserver() {
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
index 0c1b417..9f9c0d8 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
@@ -20,9 +20,11 @@
 import android.content.SharedPreferences
 import android.os.Build
 import android.util.Log
-import androidx.core.content.ContextCompat
+import androidx.annotation.VisibleForTesting
 import java.io.File
 
+private fun defaultVerbose() = Build.TYPE == "eng"
+
 /**
  * [SharedPreferences] based storage.
  *
@@ -43,24 +45,35 @@
  * @param verbose Verbose logging on key/value pairs during backup/restore. Enable for dev only!
  * @param filter Filter of key/value pairs for backup and restore.
  */
-class SharedPreferencesStorage
+open class SharedPreferencesStorage
 @JvmOverloads
 constructor(
     context: Context,
     override val name: String,
-    mode: Int,
-    private val verbose: Boolean = (Build.TYPE == "eng"),
+    @get:VisibleForTesting internal val sharedPreferences: SharedPreferences,
+    private val codec: BackupCodec? = null,
+    private val verbose: Boolean = defaultVerbose(),
     private val filter: (String, Any?) -> Boolean = { _, _ -> true },
 ) :
     BackupRestoreFileStorage(context, context.getSharedPreferencesFilePath(name)),
     KeyedObservable<String> by KeyedDataObservable() {
 
-    private val sharedPreferences = context.getSharedPreferences(name, mode)
+    @JvmOverloads
+    constructor(
+        context: Context,
+        name: String,
+        mode: Int,
+        codec: BackupCodec? = null,
+        verbose: Boolean = defaultVerbose(),
+        filter: (String, Any?) -> Boolean = { _, _ -> true },
+    ) : this(context, name, context.getSharedPreferences(name, mode), codec, verbose, filter)
 
     /** Name of the intermediate SharedPreferences. */
-    private val intermediateName: String
+    @VisibleForTesting
+    internal val intermediateName: String
         get() = "_br_$name"
 
+    @Suppress("DEPRECATION")
     private val intermediateSharedPreferences: SharedPreferences
         get() {
             // use MODE_MULTI_PROCESS to ensure a reload
@@ -82,12 +95,15 @@
         sharedPreferences.registerOnSharedPreferenceChangeListener(sharedPreferencesListener)
     }
 
+    override fun defaultCodec() = codec ?: super.defaultCodec()
+
     override val backupFile: File
         // use a different file to avoid multi-thread file write
         get() = context.getSharedPreferencesFile(intermediateName)
 
     override fun prepareBackup(file: File) {
-        val editor = intermediateSharedPreferences.merge(sharedPreferences.all, "Backup")
+        val editor =
+            mergeSharedPreferences(intermediateSharedPreferences, sharedPreferences.all, "Backup")
         // commit to ensure data is write to disk synchronously
         if (!editor.commit()) {
             Log.w(LOG_TAG, "[$name] fail to commit")
@@ -104,8 +120,8 @@
         // observers consistently once restore finished.
         sharedPreferences.unregisterOnSharedPreferenceChangeListener(sharedPreferencesListener)
         val restored = intermediateSharedPreferences
-        val editor = sharedPreferences.merge(restored.all, "Restore")
-        editor.apply() // apply to avoid blocking
+        val editor = mergeSharedPreferences(sharedPreferences, restored.all, "Restore")
+        editor.commit() // commit to avoid race condition
         sharedPreferences.registerOnSharedPreferenceChangeListener(sharedPreferencesListener)
         // clear the intermediate SharedPreferences
         restored.delete(intermediateName)
@@ -115,7 +131,7 @@
         if (deleteSharedPreferences(name)) {
             Log.i(LOG_TAG, "SharedPreferences $name deleted")
         } else {
-            edit().clear().apply()
+            edit().clear().commit() // commit to avoid potential race condition
         }
     }
 
@@ -126,11 +142,13 @@
             false
         }
 
-    private fun SharedPreferences.merge(
+    @VisibleForTesting
+    internal open fun mergeSharedPreferences(
+        sharedPreferences: SharedPreferences,
         entries: Map<String, Any?>,
-        operation: String
+        operation: String,
     ): SharedPreferences.Editor {
-        val editor = edit()
+        val editor = sharedPreferences.edit()
         for ((key, value) in entries) {
             if (!filter.invoke(key, value)) {
                 if (verbose) Log.v(LOG_TAG, "[$name] $operation skips $key=$value")
@@ -184,7 +202,7 @@
     companion object {
         private fun Context.getSharedPreferencesFilePath(name: String): String {
             val file = getSharedPreferencesFile(name)
-            return file.relativeTo(ContextCompat.getDataDir(this)!!).toString()
+            return file.relativeTo(dataDirCompat).toString()
         }
 
         /** Returns the absolute path of shared preferences file. */
diff --git a/packages/SettingsLib/DataStore/tests/Android.bp b/packages/SettingsLib/DataStore/tests/Android.bp
index 8770dfa..5d000eb 100644
--- a/packages/SettingsLib/DataStore/tests/Android.bp
+++ b/packages/SettingsLib/DataStore/tests/Android.bp
@@ -9,11 +9,16 @@
 
 android_robolectric_test {
     name: "SettingsLibDataStoreTest",
-    srcs: ["src/**/*"],
+    srcs: [
+        ":SettingsLibDataStore-srcs", // b/240432457
+        "src/**/*",
+    ],
     static_libs: [
-        "SettingsLibDataStore",
+        "androidx.collection_collection-ktx",
+        "androidx.core_core-ktx",
         "androidx.test.ext.junit",
         "guava",
+        "kotlin-test",
         "mockito-robolectric-prebuilt", // mockito deps order matters!
         "mockito-kotlin2",
     ],
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupCodecTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupCodecTest.kt
new file mode 100644
index 0000000..867831b
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupCodecTest.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
+import kotlin.random.Random
+import kotlin.test.assertFailsWith
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Tests of [BackupCodec]. */
+@RunWith(AndroidJUnit4::class)
+class BackupCodecTest {
+    @Test
+    fun name() {
+        val names = mutableSetOf<String>()
+        for (codec in allCodecs()) {
+            assertThat(names).doesNotContain(codec.name)
+            names.add(codec.name)
+        }
+    }
+
+    @Test
+    fun fromId() {
+        for (codec in allCodecs()) {
+            assertThat(BackupCodec.fromId(codec.id)).isInstanceOf(codec::class.java)
+        }
+    }
+
+    @Test
+    fun fromId_unknownId() {
+        assertFailsWith(IllegalArgumentException::class) { BackupCodec.fromId(-1) }
+    }
+
+    @Test
+    fun encode_decode() {
+        val random = Random.Default
+        fun test(codec: BackupCodec, size: Int) {
+            val data = random.nextBytes(size)
+
+            // encode
+            val outputStream = ByteArrayOutputStream()
+            codec.encode(outputStream).use { it.write(data) }
+
+            // decode
+            val inputStream = ByteArrayInputStream(outputStream.toByteArray())
+            val result = codec.decode(inputStream).use { it.readBytes() }
+
+            assertWithMessage("$size bytes: $data").that(result).isEqualTo(data)
+        }
+
+        for (codec in allCodecs()) {
+            test(codec, 0)
+            repeat(10) { test(codec, random.nextInt(1, 1024)) }
+        }
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreContextTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreContextTest.kt
new file mode 100644
index 0000000..911665a
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreContextTest.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import android.app.backup.BackupDataOutput
+import android.os.Build
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+
+/** Tests of [BackupContext] and [RestoreContext]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreContextTest {
+    @Test
+    fun backupContext_quota() {
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.O) return
+        val data = mock<BackupDataOutput> { on { quota } doReturn 10L }
+        assertThat(BackupContext(data).quota).isEqualTo(10)
+    }
+
+    @Test
+    fun backupContext_transportFlags() {
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.P) return
+        val data = mock<BackupDataOutput> { on { transportFlags } doReturn 5 }
+        assertThat(BackupContext(data).transportFlags).isEqualTo(5)
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileArchiverTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileArchiverTest.kt
new file mode 100644
index 0000000..6cce453
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileArchiverTest.kt
@@ -0,0 +1,275 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import android.app.Application
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
+import java.io.File
+import java.io.IOException
+import java.io.InputStream
+import kotlin.random.Random
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TemporaryFolder
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+
+/** Tests of [BackupRestoreFileArchiver]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreFileArchiverTest {
+    private val random = Random.Default
+    private val application: Application = getApplicationContext()
+    @get:Rule val temporaryFolder = TemporaryFolder(application.dataDirCompat)
+
+    @Test
+    fun createBackupRestoreEntities() {
+        val fileStorages = mutableListOf<BackupRestoreFileStorage>()
+        for (count in 0 until 3) {
+            val fileArchiver = BackupRestoreFileArchiver(application, fileStorages, "")
+            fileArchiver.createBackupRestoreEntities().apply {
+                assertThat(this).hasSize(fileStorages.size)
+                for (index in 0 until count) {
+                    assertThat(get(index).key).isEqualTo(fileStorages[index].storageFilePath)
+                }
+            }
+            fileStorages.add(FileStorage("storage", "path$count"))
+        }
+    }
+
+    @Test
+    fun wrapBackupOutputStream() {
+        val fileArchiver = BackupRestoreFileArchiver(application, listOf(), "")
+        val outputStream = ByteArrayOutputStream()
+        assertThat(fileArchiver.wrapBackupOutputStream(BackupZipCodec.BEST_SPEED, outputStream))
+            .isSameInstanceAs(outputStream)
+    }
+
+    @Test
+    fun wrapRestoreInputStream() {
+        val fileArchiver = BackupRestoreFileArchiver(application, listOf(), "")
+        val inputStream = ByteArrayInputStream(byteArrayOf())
+        assertThat(fileArchiver.wrapRestoreInputStream(BackupZipCodec.BEST_SPEED, inputStream))
+            .isSameInstanceAs(inputStream)
+    }
+
+    @Test
+    fun restoreEntity_disabled() {
+        val file = temporaryFolder.newFile()
+        val key = file.name
+        val fileStorage = FileStorage("fs", key, restoreEnabled = false)
+
+        BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver").apply {
+            restoreEntity(newBackupDataInputStream(key, byteArrayOf()))
+            assertThat(entityStates.asMap()).isEmpty()
+        }
+    }
+
+    @Test
+    fun restoreEntity_raiseIOException() {
+        val key = "key"
+        val fileStorage = FileStorage("fs", key)
+        BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver").apply {
+            restoreEntity(newBackupDataInputStream(key, byteArrayOf(), IOException()))
+            assertThat(entityStates.asMap()).isEmpty()
+        }
+    }
+
+    @Test
+    fun restoreEntity_onRestoreFinished_raiseException() {
+        val key = "key"
+        val fileStorage = FileStorage("fs", key, restoreException = IllegalStateException())
+        BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver").apply {
+            val data = random.nextBytes(random.nextInt(10))
+            val outputStream = ByteArrayOutputStream()
+            fileStorage.wrapBackupOutputStream(fileStorage.defaultCodec(), outputStream).use {
+                it.write(data)
+            }
+            val payload = outputStream.toByteArray()
+            restoreEntity(newBackupDataInputStream(key, payload))
+            assertThat(entityStates.asMap()).isEmpty()
+        }
+    }
+
+    @Test
+    fun restoreEntity_forwardCompatibility() {
+        val key = "key"
+        val fileStorage = FileStorage("fs", key)
+        for (codec in allCodecs()) {
+            BackupRestoreFileArchiver(application, listOf(), "archiver").apply {
+                val data = random.nextBytes(random.nextInt(MAX_DATA_SIZE))
+                val outputStream = ByteArrayOutputStream()
+                fileStorage.wrapBackupOutputStream(codec, outputStream).use { it.write(data) }
+                val payload = outputStream.toByteArray()
+
+                restoreEntity(newBackupDataInputStream(key, payload))
+
+                assertThat(entityStates.asMap()).apply {
+                    hasSize(1)
+                    containsKey(key)
+                }
+                assertThat(fileStorage.restoreFile.readBytes()).isEqualTo(data)
+            }
+        }
+    }
+
+    @Test
+    fun restoreEntity() {
+        val folder = File(application.dataDirCompat, "backup")
+        val file = File(folder, "file")
+        val key = "${folder.name}${File.separator}${file.name}"
+        fun test(codec: BackupCodec, size: Int) {
+            val fileStorage = FileStorage("fs", key, if (size % 2 == 0) codec else null)
+            val data = random.nextBytes(size)
+            val outputStream = ByteArrayOutputStream()
+            fileStorage.wrapBackupOutputStream(codec, outputStream).use { it.write(data) }
+            val payload = outputStream.toByteArray()
+
+            val fileArchiver =
+                BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver")
+            fileArchiver.restoreEntity(newBackupDataInputStream(key, payload))
+
+            assertThat(fileArchiver.entityStates.asMap()).apply {
+                hasSize(1)
+                containsKey(key)
+            }
+            assertThat(file.readBytes()).isEqualTo(data)
+        }
+
+        for (codec in allCodecs()) {
+            for (size in 0 until 100) test(codec, size)
+            repeat(10) { test(codec, random.nextInt(100, MAX_DATA_SIZE)) }
+        }
+    }
+
+    @Test
+    fun onRestoreFinished() {
+        val fileStorage = mock<BackupRestoreFileStorage>()
+        val fileArchiver = BackupRestoreFileArchiver(application, listOf(fileStorage), "")
+
+        fileArchiver.onRestoreFinished()
+
+        verify(fileStorage).onRestoreFinished()
+    }
+
+    @Test
+    fun toBackupRestoreEntity_backup_disabled() {
+        val context = BackupContext(mock())
+        val fileStorage =
+            mock<BackupRestoreFileStorage> { on { enableBackup(context) } doReturn false }
+
+        assertThat(fileStorage.toBackupRestoreEntity().backup(context, ByteArrayOutputStream()))
+            .isEqualTo(EntityBackupResult.INTACT)
+
+        verify(fileStorage, never()).prepareBackup(any())
+    }
+
+    @Test
+    fun toBackupRestoreEntity_backup_fileNotExist() {
+        val context = BackupContext(mock())
+        val file = File("NotExist")
+        val fileStorage =
+            mock<BackupRestoreFileStorage> {
+                on { enableBackup(context) } doReturn true
+                on { backupFile } doReturn file
+            }
+
+        assertThat(fileStorage.toBackupRestoreEntity().backup(context, ByteArrayOutputStream()))
+            .isEqualTo(EntityBackupResult.DELETE)
+
+        verify(fileStorage).prepareBackup(file)
+        verify(fileStorage, never()).defaultCodec()
+    }
+
+    @Test
+    fun toBackupRestoreEntity_backup() {
+        val context = BackupContext(mock())
+        val file = temporaryFolder.newFile()
+
+        fun test(codec: BackupCodec, size: Int) {
+            val data = random.nextBytes(size)
+            file.outputStream().use { it.write(data) }
+
+            val outputStream = ByteArrayOutputStream()
+            val fileStorage =
+                mock<BackupRestoreFileStorage> {
+                    on { enableBackup(context) } doReturn true
+                    on { backupFile } doReturn file
+                    on { defaultCodec() } doReturn codec
+                    on { wrapBackupOutputStream(any(), any()) }.thenCallRealMethod()
+                    on { wrapRestoreInputStream(any(), any()) }.thenCallRealMethod()
+                    on { prepareBackup(any()) }.thenCallRealMethod()
+                    on { onBackupFinished(any()) }.thenCallRealMethod()
+                }
+
+            assertThat(fileStorage.toBackupRestoreEntity().backup(context, outputStream))
+                .isEqualTo(EntityBackupResult.UPDATE)
+
+            verify(fileStorage).prepareBackup(file)
+            verify(fileStorage).onBackupFinished(file)
+
+            val decodedData =
+                fileStorage
+                    .wrapRestoreInputStream(codec, ByteArrayInputStream(outputStream.toByteArray()))
+                    .readBytes()
+            assertThat(decodedData).isEqualTo(data)
+        }
+
+        for (codec in allCodecs()) {
+            // test small data to ensure correctness
+            for (size in 0 until 100) test(codec, size)
+            repeat(10) { test(codec, random.nextInt(100, MAX_DATA_SIZE)) }
+        }
+    }
+
+    @Test
+    fun toBackupRestoreEntity_restore() {
+        val restoreContext = RestoreContext("storage")
+        val inputStream =
+            object : InputStream() {
+                override fun read() = throw IllegalStateException()
+
+                override fun read(b: ByteArray, off: Int, len: Int) = throw IllegalStateException()
+            }
+        FileStorage("storage", "path").toBackupRestoreEntity().restore(restoreContext, inputStream)
+    }
+
+    private open class FileStorage(
+        override val name: String,
+        filePath: String,
+        private val codec: BackupCodec? = null,
+        private val restoreEnabled: Boolean? = null,
+        private val restoreException: Exception? = null,
+    ) : BackupRestoreFileStorage(getApplicationContext(), filePath) {
+
+        override fun defaultCodec() = codec ?: super.defaultCodec()
+
+        override fun enableRestore() = restoreEnabled ?: super.enableRestore()
+
+        override fun onRestoreFinished(file: File) {
+            super.onRestoreFinished(file)
+            if (restoreException != null) throw restoreException
+        }
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileStorageTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileStorageTest.kt
new file mode 100644
index 0000000..422273d
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileStorageTest.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.settingslib.datastore
+
+import android.app.Application
+import android.os.Build
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import java.io.File
+import kotlin.test.assertFailsWith
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Tests of [BackupRestoreFileStorage]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreFileStorageTest {
+    private val application: Application = getApplicationContext()
+
+    @Test
+    fun dataDirCompat() {
+        val expected =
+            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
+                application.dataDir
+            } else {
+                File(application.applicationInfo.dataDir)
+            }
+        assertThat(application.dataDirCompat).isEqualTo(expected)
+    }
+
+    @Test
+    fun backupFile() {
+        assertThat(FileStorage("path").backupFile.toString())
+            .startsWith(application.dataDirCompat.toString())
+    }
+
+    @Test
+    fun restoreFile() {
+        FileStorage("path").apply { assertThat(restoreFile).isEqualTo(backupFile) }
+    }
+
+    @Test
+    fun checkFilePaths() {
+        FileStorage("path").checkFilePaths()
+    }
+
+    @Test
+    fun checkFilePaths_emptyFilePath() {
+        assertFailsWith(IllegalArgumentException::class) { FileStorage("").checkFilePaths() }
+    }
+
+    @Test
+    fun checkFilePaths_absoluteFilePath() {
+        assertFailsWith(IllegalArgumentException::class) {
+            FileStorage("${File.separatorChar}file").checkFilePaths()
+        }
+    }
+
+    @Test
+    fun checkFilePaths_backupFile() {
+        assertFailsWith(IllegalArgumentException::class) {
+            FileStorage("path", fileForBackup = File("path")).checkFilePaths()
+        }
+    }
+
+    @Test
+    fun checkFilePaths_restoreFile() {
+        assertFailsWith(IllegalArgumentException::class) {
+            FileStorage("path", fileForRestore = File("path")).checkFilePaths()
+        }
+    }
+
+    @Test
+    fun createBackupRestoreEntities() {
+        assertThat(FileStorage("path").createBackupRestoreEntities()).isEmpty()
+    }
+
+    private class FileStorage(
+        filePath: String,
+        val fileForBackup: File? = null,
+        val fileForRestore: File? = null,
+    ) : BackupRestoreFileStorage(getApplicationContext(), filePath) {
+        override val name = "storage"
+
+        override val backupFile: File
+            get() = fileForBackup ?: super.backupFile
+
+        override val restoreFile: File
+            get() = fileForRestore ?: super.restoreFile
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt
new file mode 100644
index 0000000..d8f5028
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt
@@ -0,0 +1,237 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import android.app.Application
+import android.app.backup.BackupAgentHelper
+import android.app.backup.BackupHelper
+import android.app.backup.BackupManager
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors.directExecutor
+import kotlin.test.assertFailsWith
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.reset
+import org.mockito.kotlin.verify
+import org.robolectric.Shadows
+import org.robolectric.shadows.ShadowBackupManager
+
+/** Tests of [BackupRestoreStorageManager]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreStorageManagerTest {
+    private val application: Application = getApplicationContext()
+    private val manager = BackupRestoreStorageManager.getInstance(application)
+    private val fileStorage = FileStorage("fileStorage")
+    private val keyedStorage = KeyedStorage("keyedStorage")
+
+    private val storage1 = mock<ObservableBackupRestoreStorage> { on { name } doReturn "1" }
+    private val storage2 = mock<ObservableBackupRestoreStorage> { on { name } doReturn "1" }
+
+    @After
+    fun tearDown() {
+        manager.removeAll()
+        ShadowBackupManager.reset()
+    }
+
+    @Test
+    fun getInstance() {
+        assertThat(BackupRestoreStorageManager.getInstance(application)).isSameInstanceAs(manager)
+    }
+
+    @Test
+    fun addBackupAgentHelpers() {
+        val fs = FileStorage("fs")
+        manager.add(keyedStorage, fileStorage, storage1, fs)
+        val backupAgentHelper = DummyBackupAgentHelper()
+        manager.addBackupAgentHelpers(backupAgentHelper)
+        backupAgentHelper.backupHelpers.apply {
+            assertThat(size).isEqualTo(3)
+            assertThat(remove(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+            assertThat(remove(storage1.name)).isSameInstanceAs(storage1)
+            val fileArchiver = entries.first().value as BackupRestoreFileArchiver
+            assertThat(fileArchiver.fileStorages.toSet()).containsExactly(fs, fileStorage)
+        }
+    }
+
+    @Test
+    fun addBackupAgentHelpers_withoutFileStorage() {
+        manager.add(keyedStorage, storage1)
+        val backupAgentHelper = DummyBackupAgentHelper()
+        manager.addBackupAgentHelpers(backupAgentHelper)
+        backupAgentHelper.backupHelpers.apply {
+            assertThat(size).isEqualTo(3)
+            assertThat(remove(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+            assertThat(remove(storage1.name)).isSameInstanceAs(storage1)
+            val fileArchiver = entries.first().value as BackupRestoreFileArchiver
+            assertThat(fileArchiver.fileStorages).isEmpty()
+        }
+    }
+
+    @Test
+    fun add() {
+        manager.add(keyedStorage, fileStorage, storage1)
+        assertThat(manager.storageWrappers).apply {
+            hasSize(3)
+            containsKey(keyedStorage.name)
+            containsKey(fileStorage.name)
+            containsKey(storage1.name)
+        }
+    }
+
+    @Test
+    fun add_identicalName() {
+        manager.add(storage1)
+        assertFailsWith(IllegalStateException::class) { manager.add(storage1) }
+        assertFailsWith(IllegalStateException::class) { manager.add(storage2) }
+    }
+
+    @Test
+    fun add_nonObservable() {
+        assertFailsWith(IllegalArgumentException::class) {
+            manager.add(mock<BackupRestoreStorage>())
+        }
+    }
+
+    @Test
+    fun removeAll() {
+        add()
+        manager.removeAll()
+        assertThat(manager.storageWrappers).isEmpty()
+    }
+
+    @Test
+    fun remove() {
+        manager.add(keyedStorage, fileStorage)
+        assertThat(manager.remove(storage1.name)).isNull()
+        assertThat(manager.remove(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+        assertThat(manager.remove(fileStorage.name)).isSameInstanceAs(fileStorage)
+    }
+
+    @Test
+    fun get() {
+        manager.add(keyedStorage, fileStorage)
+        assertThat(manager.get(storage1.name)).isNull()
+        assertThat(manager.get(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+        assertThat(manager.get(fileStorage.name)).isSameInstanceAs(fileStorage)
+    }
+
+    @Test
+    fun getOrThrow() {
+        manager.add(keyedStorage, fileStorage)
+        assertFailsWith(NullPointerException::class) { manager.getOrThrow(storage1.name) }
+        assertThat(manager.getOrThrow(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+        assertThat(manager.getOrThrow(fileStorage.name)).isSameInstanceAs(fileStorage)
+    }
+
+    @Test
+    fun notifyRestoreFinished() {
+        manager.add(keyedStorage, fileStorage)
+        val keyedObserver = mock<KeyedObserver<String>>()
+        val anyKeyObserver = mock<KeyedObserver<String?>>()
+        val observer = mock<Observer>()
+        val executor = directExecutor()
+        keyedStorage.addObserver("key", keyedObserver, executor)
+        keyedStorage.addObserver(anyKeyObserver, executor)
+        fileStorage.addObserver(observer, executor)
+
+        manager.onRestoreFinished()
+
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.RESTORE)
+        verify(anyKeyObserver).onKeyChanged(null, ChangeReason.RESTORE)
+        verify(observer).onChanged(ChangeReason.RESTORE)
+        if (isRobolectric()) {
+            Shadows.shadowOf(BackupManager(application)).apply {
+                assertThat(isDataChanged).isFalse()
+                assertThat(dataChangedCount).isEqualTo(0)
+            }
+        }
+    }
+
+    @Test
+    fun notifyBackupManager() {
+        manager.add(keyedStorage, fileStorage)
+        val keyedObserver = mock<KeyedObserver<String>>()
+        val anyKeyObserver = mock<KeyedObserver<String?>>()
+        val observer = mock<Observer>()
+        val executor = directExecutor()
+        keyedStorage.addObserver("key", keyedObserver, executor)
+        keyedStorage.addObserver(anyKeyObserver, executor)
+        fileStorage.addObserver(observer, executor)
+
+        val backupManager =
+            if (isRobolectric()) Shadows.shadowOf(BackupManager(application)) else null
+        backupManager?.apply {
+            assertThat(isDataChanged).isFalse()
+            assertThat(dataChangedCount).isEqualTo(0)
+        }
+
+        fileStorage.notifyChange(ChangeReason.UPDATE)
+        verify(observer).onChanged(ChangeReason.UPDATE)
+        verify(keyedObserver, never()).onKeyChanged(any(), any())
+        verify(anyKeyObserver, never()).onKeyChanged(any(), any())
+        reset(observer)
+        backupManager?.apply {
+            assertThat(isDataChanged).isTrue()
+            assertThat(dataChangedCount).isEqualTo(1)
+        }
+
+        keyedStorage.notifyChange("key", ChangeReason.DELETE)
+        verify(observer, never()).onChanged(any())
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.DELETE)
+        verify(anyKeyObserver).onKeyChanged("key", ChangeReason.DELETE)
+        backupManager?.apply {
+            assertThat(isDataChanged).isTrue()
+            assertThat(dataChangedCount).isEqualTo(2)
+        }
+        reset(keyedObserver)
+
+        // backup manager is not notified for restore event
+        fileStorage.notifyChange(ChangeReason.RESTORE)
+        keyedStorage.notifyChange("key", ChangeReason.RESTORE)
+        verify(observer).onChanged(ChangeReason.RESTORE)
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.RESTORE)
+        verify(anyKeyObserver).onKeyChanged("key", ChangeReason.RESTORE)
+        backupManager?.apply {
+            assertThat(isDataChanged).isTrue()
+            assertThat(dataChangedCount).isEqualTo(2)
+        }
+    }
+
+    private class KeyedStorage(override val name: String) :
+        BackupRestoreStorage(), KeyedObservable<String> by KeyedDataObservable() {
+
+        override fun createBackupRestoreEntities(): List<BackupRestoreEntity> = listOf()
+    }
+
+    private class FileStorage(override val name: String) :
+        BackupRestoreFileStorage(getApplicationContext(), "file"), Observable by DataObservable()
+
+    private class DummyBackupAgentHelper : BackupAgentHelper() {
+        val backupHelpers = mutableMapOf<String, BackupHelper>()
+
+        override fun addHelper(keyPrefix: String, helper: BackupHelper) {
+            backupHelpers[keyPrefix] = helper
+        }
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageTest.kt
new file mode 100644
index 0000000..99998ff
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageTest.kt
@@ -0,0 +1,414 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import android.app.backup.BackupAgentHelper
+import android.app.backup.BackupDataInput
+import android.app.backup.BackupDataInputStream
+import android.app.backup.BackupDataOutput
+import android.os.ParcelFileDescriptor
+import android.os.ParcelFileDescriptor.MODE_APPEND
+import android.os.ParcelFileDescriptor.MODE_READ_ONLY
+import android.os.ParcelFileDescriptor.MODE_WRITE_ONLY
+import androidx.collection.MutableScatterMap
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import java.io.DataOutputStream
+import java.io.File
+import java.io.FileDescriptor
+import java.io.FileOutputStream
+import java.io.InputStream
+import java.io.OutputStream
+import kotlin.random.Random
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TemporaryFolder
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.doThrow
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.reset
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.verify
+
+/** Tests of [BackupRestoreStorage]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreStorageTest {
+    @get:Rule val temporaryFolder = TemporaryFolder()
+
+    private val entity1 = Entity("key1", "value1".toByteArray())
+    private val entity1NoOpCodec = Entity("key1", "value1".toByteArray(), BackupNoOpCodec())
+    private val entity2 = Entity("key2", "value2".toByteArray(), BackupZipCodec.BEST_SPEED)
+
+    @Test
+    fun performBackup_disabled() {
+        val storage = spy(TestStorage().apply { enabled = false })
+        val unused = performBackup { data, newState -> storage.performBackup(null, data, newState) }
+        verify(storage, never()).createBackupRestoreEntities()
+        assertThat(storage.entities).isNull()
+        assertThat(storage.entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun performBackup_enabled() {
+        val storage = spy(TestStorage())
+        val unused = performBackup { data, newState -> storage.performBackup(null, data, newState) }
+        verify(storage).createBackupRestoreEntities()
+        assertThat(storage.entities).isNull()
+        assertThat(storage.entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun performBackup_entityBackupWithException() {
+        val entity =
+            mock<BackupRestoreEntity> {
+                on { key } doReturn ""
+                on { backup(any(), any()) } doThrow IllegalStateException()
+            }
+        val storage = TestStorage(entity, entity1)
+
+        val (_, stateFile) =
+            performBackup { data, newState -> storage.performBackup(null, data, newState) }
+
+        assertThat(storage.readEntityStates(stateFile)).apply {
+            hasSize(1)
+            containsKey(entity1.key)
+        }
+    }
+
+    @Test
+    fun performBackup_update_unchanged() {
+        performBackupTest({}) { entityStates, newEntityStates ->
+            assertThat(entityStates).isEqualTo(newEntityStates)
+        }
+    }
+
+    @Test
+    fun performBackup_intact() {
+        performBackupTest({ entity1.backupResult = EntityBackupResult.INTACT }) {
+            entityStates,
+            newEntityStates ->
+            assertThat(entityStates).isEqualTo(newEntityStates)
+        }
+    }
+
+    @Test
+    fun performBackup_delete() {
+        performBackupTest({ entity1.backupResult = EntityBackupResult.DELETE }) { _, newEntityStates
+            ->
+            assertThat(newEntityStates.size).isEqualTo(1)
+            assertThat(newEntityStates).containsKey(entity2.key)
+        }
+    }
+
+    private fun performBackupTest(
+        update: () -> Unit,
+        verification: (Map<String, Long>, Map<String, Long>) -> Unit,
+    ) {
+        val storage = TestStorage(entity1, entity2)
+        val (_, stateFile) =
+            performBackup { data, newState -> storage.performBackup(null, data, newState) }
+
+        val entityStates = storage.readEntityStates(stateFile)
+        assertThat(entityStates).apply {
+            hasSize(2)
+            containsKey(entity1.key)
+            containsKey(entity2.key)
+        }
+
+        update.invoke()
+        val (_, newStateFile) =
+            performBackup { data, newState ->
+                stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+                    storage.performBackup(it, data, newState)
+                }
+            }
+        verification.invoke(entityStates, storage.readEntityStates(newStateFile))
+    }
+
+    @Test
+    fun restoreEntity_disabled() {
+        val storage = spy(TestStorage().apply { enabled = false })
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.restoreEntity(it.toBackupDataInputStream())
+        }
+        verify(storage, never()).createBackupRestoreEntities()
+        assertThat(storage.entities).isNull()
+        assertThat(storage.entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun restoreEntity_entityNotFound() {
+        val storage = TestStorage()
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            val backupDataInputStream = it.toBackupDataInputStream()
+            backupDataInputStream.setKey("")
+            storage.restoreEntity(backupDataInputStream)
+        }
+    }
+
+    @Test
+    fun restoreEntity_exception() {
+        val storage = TestStorage(entity1)
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            val backupDataInputStream = it.toBackupDataInputStream()
+            backupDataInputStream.setKey(entity1.key)
+            storage.restoreEntity(backupDataInputStream)
+        }
+    }
+
+    @Test
+    fun restoreEntity_codecChanged() {
+        assertThat(entity1.codec()).isNotEqualTo(entity1NoOpCodec.codec())
+        backupAndRestore(entity1) { _, data ->
+            TestStorage(entity1NoOpCodec).apply { restoreEntity(data) }
+        }
+        assertThat(entity1.data).isEqualTo(entity1NoOpCodec.restoredData)
+    }
+
+    @Test
+    fun restoreEntity() {
+        val random = Random.Default
+        fun test(codec: BackupCodec, size: Int) {
+            val entity = Entity("key", random.nextBytes(size), codec)
+            backupAndRestore(entity)
+            entity.verifyRestoredData()
+        }
+        for (codec in allCodecs()) {
+            // test small data to ensure correctness
+            for (size in 0 until 100) test(codec, size)
+            repeat(10) { test(codec, random.nextInt(100, MAX_DATA_SIZE)) }
+        }
+    }
+
+    @Test
+    fun readEntityStates_eof_exception() {
+        val storage = TestStorage()
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates.put("", 0) // add an item to verify that exiting elements are clear
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun readEntityStates_other_exception() {
+        val storage = TestStorage()
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates.put("", 0) // add an item to verify that exiting elements are clear
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).apply {
+            close() // cause exception when read state file
+            storage.readEntityStates(this, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun readEntityStates_unknownVersion() {
+        val storage = TestStorage()
+        val stateFile = temporaryFolder.newFile()
+        stateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+            DataOutputStream(FileOutputStream(it.fileDescriptor))
+                .writeByte(BackupRestoreStorage.STATE_VERSION + 1)
+        }
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates.put("", 0) // add an item to verify that exiting elements are clear
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun writeNewStateDescription() {
+        val storage = spy(TestStorage())
+        // use read only mode to trigger exception when write state file
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.writeNewStateDescription(it)
+        }
+        verify(storage).onRestoreFinished()
+    }
+
+    @Test
+    fun backupAndRestore() {
+        val storage = spy(TestStorage(entity1, entity2))
+        val backupAgentHelper = BackupAgentHelper()
+        backupAgentHelper.addHelper(storage.name, storage)
+
+        // backup
+        val (dataFile, stateFile) =
+            performBackup { data, newState -> backupAgentHelper.onBackup(null, data, newState) }
+        storage.verifyFieldsArePurged()
+
+        // verify state
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates[""] = 1
+        storage.readEntityStates(null, entityStates)
+        assertThat(entityStates.size).isEqualTo(0)
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.asMap()).apply {
+            hasSize(2)
+            containsKey(entity1.key)
+            containsKey(entity2.key)
+        }
+        reset(storage)
+
+        // restore
+        val newStateFile = temporaryFolder.newFile()
+        dataFile.toParcelFileDescriptor(MODE_READ_ONLY).use { dataPfd ->
+            newStateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+                backupAgentHelper.onRestore(dataPfd.toBackupDataInput(), 0, it)
+            }
+        }
+        verify(storage).onRestoreFinished()
+        storage.verifyFieldsArePurged()
+
+        // ShadowBackupDataOutput does not write data to file, so restore is bypassed
+        if (!isRobolectric()) {
+            entity1.verifyRestoredData()
+            entity2.verifyRestoredData()
+            assertThat(entityStates.asMap()).isEqualTo(storage.readEntityStates(newStateFile))
+        }
+    }
+
+    private fun backupAndRestore(
+        entity: BackupRestoreEntity,
+        restoreEntity: (TestStorage, BackupDataInputStream) -> TestStorage = { storage, data ->
+            storage.restoreEntity(data)
+            storage
+        },
+    ) {
+        val storage = TestStorage(entity)
+        val entityKey = argumentCaptor<String>()
+        val entitySize = argumentCaptor<Int>()
+        val entityData = argumentCaptor<ByteArray>()
+        val data = mock<BackupDataOutput>()
+
+        val stateFile = temporaryFolder.newFile()
+        stateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+            storage.performBackup(null, data, it)
+        }
+        val entityStates = MutableScatterMap<String, Long>()
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(1)
+
+        verify(data).writeEntityHeader(entityKey.capture(), entitySize.capture())
+        verify(data).writeEntityData(entityData.capture(), entitySize.capture())
+        assertThat(entityKey.allValues).isEqualTo(listOf(entity.key))
+        assertThat(entityData.allValues).hasSize(1)
+        val payload = entityData.firstValue
+        assertThat(entitySize.allValues).isEqualTo(listOf(payload.size, payload.size))
+
+        val dataFile = temporaryFolder.newFile()
+        dataFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+            FileOutputStream(it.fileDescriptor).write(payload)
+        }
+
+        newBackupDataInputStream(entity.key, payload).apply {
+            restoreEntity.invoke(storage, this).also {
+                assertThat(it.entityStates).isEqualTo(entityStates)
+            }
+        }
+    }
+
+    fun performBackup(backup: (BackupDataOutput, ParcelFileDescriptor) -> Unit): Pair<File, File> {
+        val dataFile = temporaryFolder.newFile()
+        val stateFile = temporaryFolder.newFile()
+        dataFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use { dataPfd ->
+            stateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+                backup.invoke(dataPfd.toBackupDataOutput(), it)
+            }
+        }
+        return dataFile to stateFile
+    }
+
+    private fun BackupRestoreStorage.verifyFieldsArePurged() {
+        assertThat(entities).isNull()
+        assertThat(entityStates.size).isEqualTo(0)
+        assertThat(entityStates.capacity).isEqualTo(0)
+    }
+
+    private fun BackupRestoreStorage.readEntityStates(stateFile: File): Map<String, Long> {
+        val entityStates = MutableScatterMap<String, Long>()
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use { readEntityStates(it, entityStates) }
+        return entityStates.asMap()
+    }
+
+    private fun File.toParcelFileDescriptor(mode: Int) = ParcelFileDescriptor.open(this, mode)
+
+    private fun ParcelFileDescriptor.toBackupDataOutput() = fileDescriptor.toBackupDataOutput()
+
+    private fun ParcelFileDescriptor.toBackupDataInputStream(): BackupDataInputStream =
+        BackupDataInputStream::class.java.newInstance(toBackupDataInput())
+
+    private fun ParcelFileDescriptor.toBackupDataInput() = fileDescriptor.toBackupDataInput()
+
+    private fun FileDescriptor.toBackupDataOutput(): BackupDataOutput =
+        BackupDataOutput::class.java.newInstance(this)
+
+    private fun FileDescriptor.toBackupDataInput(): BackupDataInput =
+        BackupDataInput::class.java.newInstance(this)
+}
+
+private open class TestStorage(vararg val backupRestoreEntities: BackupRestoreEntity) :
+    ObservableBackupRestoreStorage() {
+    var enabled: Boolean? = null
+
+    override val name
+        get() = "TestBackup"
+
+    override fun createBackupRestoreEntities() = backupRestoreEntities.toList()
+
+    override fun enableBackup(backupContext: BackupContext) =
+        enabled ?: super.enableBackup(backupContext)
+
+    override fun enableRestore() = enabled ?: super.enableRestore()
+}
+
+private class Entity(
+    override val key: String,
+    val data: ByteArray,
+    private val codec: BackupCodec? = null,
+) : BackupRestoreEntity {
+    var restoredData: ByteArray? = null
+    var backupResult = EntityBackupResult.UPDATE
+
+    override fun codec() = codec ?: super.codec()
+
+    override fun backup(
+        backupContext: BackupContext,
+        outputStream: OutputStream,
+    ): EntityBackupResult {
+        outputStream.write(data)
+        return backupResult
+    }
+
+    override fun restore(restoreContext: RestoreContext, inputStream: InputStream) {
+        restoredData = inputStream.readBytes()
+        inputStream.close()
+    }
+
+    fun verifyRestoredData() = assertThat(restoredData).isEqualTo(data)
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt
index b52586c..8638b2f 100644
--- a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt
@@ -16,76 +16,58 @@
 
 package com.android.settingslib.datastore
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.google.common.truth.Truth.assertThat
-import com.google.common.util.concurrent.MoreExecutors.directExecutor
+import com.google.common.util.concurrent.MoreExecutors
 import java.util.concurrent.Executor
 import java.util.concurrent.atomic.AtomicInteger
 import org.junit.Assert
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.junit.MockitoJUnit
-import org.mockito.junit.MockitoRule
-import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
 import org.mockito.kotlin.reset
 import org.mockito.kotlin.verify
-import org.robolectric.RobolectricTestRunner
 
-@RunWith(RobolectricTestRunner::class)
+@RunWith(AndroidJUnit4::class)
 class KeyedObserverTest {
-    @get:Rule
-    val mockitoRule: MockitoRule = MockitoJUnit.rule()
+    private val observer1 = mock<KeyedObserver<Any?>>()
+    private val observer2 = mock<KeyedObserver<Any?>>()
+    private val keyedObserver1 = mock<KeyedObserver<Any>>()
+    private val keyedObserver2 = mock<KeyedObserver<Any>>()
 
-    @Mock
-    private lateinit var observer1: KeyedObserver<Any?>
+    private val key1 = Object()
+    private val key2 = Object()
 
-    @Mock
-    private lateinit var observer2: KeyedObserver<Any?>
-
-    @Mock
-    private lateinit var keyedObserver1: KeyedObserver<Any>
-
-    @Mock
-    private lateinit var keyedObserver2: KeyedObserver<Any>
-
-    @Mock
-    private lateinit var key1: Any
-
-    @Mock
-    private lateinit var key2: Any
-
-    @Mock
-    private lateinit var executor: Executor
-
+    private val executor1: Executor = MoreExecutors.directExecutor()
+    private val executor2: Executor = MoreExecutors.newDirectExecutorService()
     private val keyedObservable = KeyedDataObservable<Any>()
 
     @Test
     fun addObserver_sameExecutor() {
-        keyedObservable.addObserver(observer1, executor)
-        keyedObservable.addObserver(observer1, executor)
+        keyedObservable.addObserver(observer1, executor1)
+        keyedObservable.addObserver(observer1, executor1)
     }
 
     @Test
     fun addObserver_keyedObserver_sameExecutor() {
-        keyedObservable.addObserver(key1, keyedObserver1, executor)
-        keyedObservable.addObserver(key1, keyedObserver1, executor)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
     }
 
     @Test
     fun addObserver_differentExecutor() {
-        keyedObservable.addObserver(observer1, executor)
+        keyedObservable.addObserver(observer1, executor1)
         Assert.assertThrows(IllegalStateException::class.java) {
-            keyedObservable.addObserver(observer1, directExecutor())
+            keyedObservable.addObserver(observer1, executor2)
         }
     }
 
     @Test
     fun addObserver_keyedObserver_differentExecutor() {
-        keyedObservable.addObserver(key1, keyedObserver1, executor)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
         Assert.assertThrows(IllegalStateException::class.java) {
-            keyedObservable.addObserver(key1, keyedObserver1, directExecutor())
+            keyedObservable.addObserver(key1, keyedObserver1, executor2)
         }
     }
 
@@ -93,7 +75,7 @@
     fun addObserver_weaklyReferenced() {
         val counter = AtomicInteger()
         var observer: KeyedObserver<Any?>? = KeyedObserver { _, _ -> counter.incrementAndGet() }
-        keyedObservable.addObserver(observer!!, directExecutor())
+        keyedObservable.addObserver(observer!!, executor1)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         assertThat(counter.get()).isEqualTo(1)
@@ -111,7 +93,7 @@
     fun addObserver_keyedObserver_weaklyReferenced() {
         val counter = AtomicInteger()
         var keyObserver: KeyedObserver<Any>? = KeyedObserver { _, _ -> counter.incrementAndGet() }
-        keyedObservable.addObserver(key1, keyObserver!!, directExecutor())
+        keyedObservable.addObserver(key1, keyObserver!!, executor1)
 
         keyedObservable.notifyChange(key1, ChangeReason.UPDATE)
         assertThat(counter.get()).isEqualTo(1)
@@ -127,45 +109,43 @@
 
     @Test
     fun addObserver_notifyObservers_removeObserver() {
-        keyedObservable.addObserver(observer1, directExecutor())
-        keyedObservable.addObserver(observer2, executor)
+        keyedObservable.addObserver(observer1, executor1)
+        keyedObservable.addObserver(observer2, executor2)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         verify(observer1).onKeyChanged(null, ChangeReason.UPDATE)
-        verify(observer2, never()).onKeyChanged(any(), any())
-        verify(executor).execute(any())
+        verify(observer2).onKeyChanged(null, ChangeReason.UPDATE)
 
-        reset(observer1, executor)
+        reset(observer1, observer2)
         keyedObservable.removeObserver(observer2)
 
         keyedObservable.notifyChange(ChangeReason.DELETE)
         verify(observer1).onKeyChanged(null, ChangeReason.DELETE)
-        verify(executor, never()).execute(any())
+        verify(observer2, never()).onKeyChanged(null, ChangeReason.DELETE)
     }
 
     @Test
     fun addObserver_keyedObserver_notifyObservers_removeObserver() {
-        keyedObservable.addObserver(key1, keyedObserver1, directExecutor())
-        keyedObservable.addObserver(key2, keyedObserver2, executor)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
+        keyedObservable.addObserver(key2, keyedObserver2, executor2)
 
         keyedObservable.notifyChange(key1, ChangeReason.UPDATE)
         verify(keyedObserver1).onKeyChanged(key1, ChangeReason.UPDATE)
-        verify(keyedObserver2, never()).onKeyChanged(any(), any())
-        verify(executor, never()).execute(any())
+        verify(keyedObserver2, never()).onKeyChanged(key2, ChangeReason.UPDATE)
 
-        reset(keyedObserver1, executor)
-        keyedObservable.removeObserver(key2, keyedObserver2)
+        reset(keyedObserver1, keyedObserver2)
+        keyedObservable.removeObserver(key1, keyedObserver1)
 
         keyedObservable.notifyChange(key1, ChangeReason.DELETE)
-        verify(keyedObserver1).onKeyChanged(key1, ChangeReason.DELETE)
-        verify(executor, never()).execute(any())
+        verify(keyedObserver1, never()).onKeyChanged(key1, ChangeReason.DELETE)
+        verify(keyedObserver2, never()).onKeyChanged(key2, ChangeReason.DELETE)
     }
 
     @Test
     fun notifyChange_addMoreTypeObservers_checkOnKeyChanged() {
-        keyedObservable.addObserver(observer1, directExecutor())
-        keyedObservable.addObserver(key1, keyedObserver1, directExecutor())
-        keyedObservable.addObserver(key2, keyedObserver2, directExecutor())
+        keyedObservable.addObserver(observer1, executor1)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
+        keyedObservable.addObserver(key2, keyedObserver2, executor1)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         verify(observer1).onKeyChanged(null, ChangeReason.UPDATE)
@@ -191,10 +171,10 @@
     fun notifyChange_addObserverWithinCallback() {
         // ConcurrentModificationException is raised if it is not implemented correctly
         val observer: KeyedObserver<Any?> = KeyedObserver { _, _ ->
-            keyedObservable.addObserver(observer1, executor)
+            keyedObservable.addObserver(observer1, executor1)
         }
 
-        keyedObservable.addObserver(observer, directExecutor())
+        keyedObservable.addObserver(observer, executor1)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         keyedObservable.removeObserver(observer)
@@ -204,12 +184,12 @@
     fun notifyChange_KeyedObserver_addObserverWithinCallback() {
         // ConcurrentModificationException is raised if it is not implemented correctly
         val keyObserver: KeyedObserver<Any?> = KeyedObserver { _, _ ->
-            keyedObservable.addObserver(key1, keyedObserver1, executor)
+            keyedObservable.addObserver(key1, keyedObserver1, executor1)
         }
 
-        keyedObservable.addObserver(key1, keyObserver, directExecutor())
+        keyedObservable.addObserver(key1, keyObserver, executor1)
 
         keyedObservable.notifyChange(key1, ChangeReason.UPDATE)
         keyedObservable.removeObserver(key1, keyObserver)
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
index f065829..173c2b1 100644
--- a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
@@ -22,40 +22,33 @@
 import java.util.concurrent.Executor
 import java.util.concurrent.atomic.AtomicInteger
 import org.junit.Assert.assertThrows
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.junit.MockitoJUnit
-import org.mockito.junit.MockitoRule
-import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
 import org.mockito.kotlin.reset
 import org.mockito.kotlin.verify
 
 @RunWith(AndroidJUnit4::class)
 class ObserverTest {
-    @get:Rule val mockitoRule: MockitoRule = MockitoJUnit.rule()
+    private val observer1 = mock<Observer>()
+    private val observer2 = mock<Observer>()
 
-    @Mock private lateinit var observer1: Observer
-
-    @Mock private lateinit var observer2: Observer
-
-    @Mock private lateinit var executor: Executor
-
+    private val executor1: Executor = MoreExecutors.directExecutor()
+    private val executor2: Executor = MoreExecutors.newDirectExecutorService()
     private val observable = DataObservable()
 
     @Test
     fun addObserver_sameExecutor() {
-        observable.addObserver(observer1, executor)
-        observable.addObserver(observer1, executor)
+        observable.addObserver(observer1, executor1)
+        observable.addObserver(observer1, executor1)
     }
 
     @Test
     fun addObserver_differentExecutor() {
-        observable.addObserver(observer1, executor)
+        observable.addObserver(observer1, executor1)
         assertThrows(IllegalStateException::class.java) {
-            observable.addObserver(observer1, MoreExecutors.directExecutor())
+            observable.addObserver(observer1, executor2)
         }
     }
 
@@ -63,7 +56,7 @@
     fun addObserver_weaklyReferenced() {
         val counter = AtomicInteger()
         var observer: Observer? = Observer { counter.incrementAndGet() }
-        observable.addObserver(observer!!, MoreExecutors.directExecutor())
+        observable.addObserver(observer!!, executor1)
 
         observable.notifyChange(ChangeReason.UPDATE)
         assertThat(counter.get()).isEqualTo(1)
@@ -79,31 +72,27 @@
 
     @Test
     fun addObserver_notifyObservers_removeObserver() {
-        observable.addObserver(observer1, MoreExecutors.directExecutor())
-        observable.addObserver(observer2, executor)
+        observable.addObserver(observer1, executor1)
+        observable.addObserver(observer2, executor2)
 
         observable.notifyChange(ChangeReason.DELETE)
 
         verify(observer1).onChanged(ChangeReason.DELETE)
-        verify(observer2, never()).onChanged(any())
-        verify(executor).execute(any())
+        verify(observer2).onChanged(ChangeReason.DELETE)
 
-        reset(observer1, executor)
+        reset(observer1, observer2)
         observable.removeObserver(observer2)
 
         observable.notifyChange(ChangeReason.UPDATE)
         verify(observer1).onChanged(ChangeReason.UPDATE)
-        verify(executor, never()).execute(any())
+        verify(observer2, never()).onChanged(ChangeReason.UPDATE)
     }
 
     @Test
     fun notifyChange_addObserverWithinCallback() {
         // ConcurrentModificationException is raised if it is not implemented correctly
-        val observer = Observer { observable.addObserver(observer1, executor) }
-        observable.addObserver(
-            observer,
-            MoreExecutors.directExecutor()
-        )
+        val observer = Observer { observable.addObserver(observer1, executor1) }
+        observable.addObserver(observer, executor1)
         observable.notifyChange(ChangeReason.UPDATE)
         observable.removeObserver(observer)
     }
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/SharedPreferencesStorageTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/SharedPreferencesStorageTest.kt
new file mode 100644
index 0000000..fec7d75
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/SharedPreferencesStorageTest.kt
@@ -0,0 +1,175 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import android.app.Application
+import android.content.Context
+import android.content.SharedPreferences
+import android.os.Build
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors
+import java.io.ByteArrayOutputStream
+import java.io.File
+import java.util.concurrent.Executor
+import kotlin.random.Random
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.verify
+
+/** Tests of [SharedPreferencesStorage]. */
+@RunWith(AndroidJUnit4::class)
+class SharedPreferencesStorageTest {
+    private val random = Random.Default
+    private val application: Application = ApplicationProvider.getApplicationContext()
+    private val map =
+        mapOf(
+            "boolean" to true,
+            "float" to random.nextFloat(),
+            "int" to random.nextInt(),
+            "long" to random.nextLong(),
+            "string" to "string",
+            "set" to setOf("string"),
+        )
+
+    @After
+    fun tearDown() {
+        application.getSharedPreferences(NAME, MODE).edit().clear().applySync()
+    }
+
+    @Test
+    fun constructors() {
+        val storage1 = SharedPreferencesStorage(application, NAME, MODE)
+        val storage2 =
+            SharedPreferencesStorage(
+                application,
+                NAME,
+                application.getSharedPreferences(NAME, MODE),
+            )
+        assertThat(storage1.sharedPreferences).isSameInstanceAs(storage2.sharedPreferences)
+    }
+
+    @Test
+    fun observer() {
+        val observer = mock<KeyedObserver<Any?>>()
+        val keyedObserver = mock<KeyedObserver<Any>>()
+        val storage = SharedPreferencesStorage(application, NAME, MODE)
+        val executor: Executor = MoreExecutors.directExecutor()
+        storage.addObserver(observer, executor)
+        storage.addObserver("key", keyedObserver, executor)
+
+        storage.sharedPreferences.edit().putString("key", "string").applySync()
+        verify(observer).onKeyChanged("key", ChangeReason.UPDATE)
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.UPDATE)
+
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
+            storage.sharedPreferences.edit().clear().applySync()
+            verify(observer).onKeyChanged(null, ChangeReason.DELETE)
+            verify(keyedObserver).onKeyChanged("key", ChangeReason.DELETE)
+        }
+    }
+
+    @Test
+    fun prepareBackup_commitFailed() {
+        val editor = mock<SharedPreferences.Editor> { on { commit() } doReturn false }
+        val storage =
+            spy(SharedPreferencesStorage(application, NAME, MODE)) {
+                onGeneric { mergeSharedPreferences(any(), any(), any()) } doReturn editor
+            }
+        storage.prepareBackup(File(""))
+    }
+
+    @Test
+    fun backupAndRestore() {
+        fun test(codec: BackupCodec) {
+            val storage = SharedPreferencesStorage(application, NAME, MODE, codec)
+            storage.mergeSharedPreferences(storage.sharedPreferences, map, "op").commit()
+            assertThat(storage.sharedPreferences.all).isEqualTo(map)
+
+            val outputStream = ByteArrayOutputStream()
+            assertThat(storage.toBackupRestoreEntity().backup(BackupContext(mock()), outputStream))
+                .isEqualTo(EntityBackupResult.UPDATE)
+            val payload = outputStream.toByteArray()
+
+            storage.sharedPreferences.edit().clear().commit()
+            assertThat(storage.sharedPreferences.all).isEmpty()
+
+            BackupRestoreFileArchiver(application, listOf(storage), "archiver")
+                .restoreEntity(newBackupDataInputStream(storage.storageFilePath, payload))
+            assertThat(storage.sharedPreferences.all).isEqualTo(map)
+        }
+
+        for (codec in allCodecs()) test(codec)
+    }
+
+    @Test
+    fun mergeSharedPreferences_filter() {
+        val storage =
+            SharedPreferencesStorage(application, NAME, MODE) { key, value ->
+                key == "float" || value is String
+            }
+        storage.mergeSharedPreferences(storage.sharedPreferences, map, "op").apply()
+        assertThat(storage.sharedPreferences.all)
+            .containsExactly("float", map["float"], "string", map["string"])
+    }
+
+    @Test
+    fun mergeSharedPreferences_invalidSet() {
+        val storage = SharedPreferencesStorage(application, NAME, MODE, verbose = true)
+        storage
+            .mergeSharedPreferences(
+                storage.sharedPreferences,
+                mapOf<String, Any>("set" to setOf(Any())),
+                "op"
+            )
+            .apply()
+        assertThat(storage.sharedPreferences.all).isEmpty()
+    }
+
+    @Test
+    fun mergeSharedPreferences_unknownType() {
+        val storage = SharedPreferencesStorage(application, NAME, MODE)
+        storage
+            .mergeSharedPreferences(storage.sharedPreferences, map + ("key" to Any()), "op")
+            .apply()
+        assertThat(storage.sharedPreferences.all).isEqualTo(map)
+    }
+
+    @Test
+    fun mergeSharedPreferences() {
+        val storage = SharedPreferencesStorage(application, NAME, MODE, verbose = true)
+        storage.mergeSharedPreferences(storage.sharedPreferences, map, "op").apply()
+        assertThat(storage.sharedPreferences.all).isEqualTo(map)
+    }
+
+    private fun SharedPreferences.Editor.applySync() {
+        apply()
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+    }
+
+    companion object {
+        private const val NAME = "pref"
+        private const val MODE = Context.MODE_PRIVATE
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/TestUtils.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/TestUtils.kt
new file mode 100644
index 0000000..823d222
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/TestUtils.kt
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.datastore
+
+import android.app.backup.BackupDataInput
+import android.app.backup.BackupDataInputStream
+import android.os.Build
+import java.io.ByteArrayInputStream
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.mock
+
+internal const val MAX_DATA_SIZE = 1 shl 12
+
+internal fun allCodecs() =
+    arrayOf<BackupCodec>(
+        BackupNoOpCodec(),
+    ) + zipCodecs()
+
+internal fun zipCodecs() =
+    arrayOf<BackupCodec>(
+        BackupZipCodec.DEFAULT_COMPRESSION,
+        BackupZipCodec.BEST_COMPRESSION,
+        BackupZipCodec.BEST_SPEED,
+    )
+
+internal fun <T : Any> Class<T>.newInstance(arg: Any, type: Class<*> = arg.javaClass): T =
+    getDeclaredConstructor(type).apply { isAccessible = true }.newInstance(arg)
+
+internal fun newBackupDataInputStream(
+    key: String,
+    data: ByteArray,
+    e: Exception? = null,
+): BackupDataInputStream {
+    // ShadowBackupDataOutput does not write data to file, so mock for reading data
+    val inputStream = ByteArrayInputStream(data)
+    val backupDataInput =
+        mock<BackupDataInput> {
+            on { readEntityData(any(), any(), any()) } doAnswer
+                {
+                    if (e != null) throw e
+                    val buf = it.arguments[0] as ByteArray
+                    val offset = it.arguments[1] as Int
+                    val size = it.arguments[2] as Int
+                    inputStream.read(buf, offset, size)
+                }
+        }
+    return BackupDataInputStream::class
+        .java
+        .newInstance(backupDataInput, BackupDataInput::class.java)
+        .apply {
+            setKey(key)
+            setDataSize(data.size)
+        }
+}
+
+internal fun BackupDataInputStream.setKey(value: Any) {
+    val field = javaClass.getDeclaredField("key")
+    field.isAccessible = true
+    field.set(this, value)
+}
+
+internal fun BackupDataInputStream.setDataSize(dataSize: Int) {
+    val field = javaClass.getDeclaredField("dataSize")
+    field.isAccessible = true
+    field.setInt(this, dataSize)
+}
+
+internal fun isRobolectric() = Build.FINGERPRINT.contains("robolectric")
diff --git a/packages/SettingsLib/aconfig/OWNERS b/packages/SettingsLib/aconfig/OWNERS
new file mode 100644
index 0000000..ba02d20
--- /dev/null
+++ b/packages/SettingsLib/aconfig/OWNERS
@@ -0,0 +1,2 @@
+# go/android-fwk-media-solutions for info on areas of ownership.
+per-file settingslib_media_flag_declarations.aconfig = file:platform/frameworks/av:/media/janitors/media_solutions_OWNERS
diff --git a/packages/SettingsLib/aconfig/settingslib_media_flag_declarations.aconfig b/packages/SettingsLib/aconfig/settingslib_media_flag_declarations.aconfig
index 4d70aec..7aae1a6 100644
--- a/packages/SettingsLib/aconfig/settingslib_media_flag_declarations.aconfig
+++ b/packages/SettingsLib/aconfig/settingslib_media_flag_declarations.aconfig
@@ -21,3 +21,13 @@
     description: "Enable Output Switcher when no media is playing."
     bug: "284227163"
 }
+
+flag {
+    name: "remove_unnecessary_route_scanning"
+    namespace: "media_solutions"
+    description: "Avoid active scan requests on UI components that only display route status information."
+    bug: "332515672"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java b/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java
index 3e29872..eae58ad 100644
--- a/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java
+++ b/packages/SettingsLib/src/com/android/settingslib/media/InfoMediaManager.java
@@ -150,7 +150,16 @@
         mUserHandle = userHandle;
     }
 
-    /** Creates an instance of InfoMediaManager. */
+    /**
+     * Creates an instance of InfoMediaManager.
+     *
+     * @param context The {@link Context}.
+     * @param packageName The package name of the app for which to control routing, or null if the
+     *     caller is interested in system-level routing only (for example, headsets, built-in
+     *     speakers, as opposed to app-specific routing (for example, casting to another device).
+     * @param userHandle The {@link UserHandle} of the user on which the app to control is running,
+     *     or null if the caller does not need app-specific routing (see {@code packageName}).
+     */
     public static InfoMediaManager createInstance(
             Context context,
             @Nullable String packageName,
@@ -185,11 +194,7 @@
     }
 
     public void startScan() {
-        mMediaDevices.clear();
-        registerRouter();
         startScanOnRouter();
-        updateRouteListingPreference();
-        refreshDevices();
     }
 
     private void updateRouteListingPreference() {
@@ -203,7 +208,6 @@
 
     public final void stopScan() {
         stopScanOnRouter();
-        unregisterRouter();
     }
 
     protected abstract void stopScanOnRouter();
@@ -290,14 +294,37 @@
         return null;
     }
 
-    protected final void registerCallback(MediaDeviceCallback callback) {
+    /**
+     * Registers the specified {@code callback} to receive state updates about routing information.
+     *
+     * <p>As long as there is a registered {@link MediaDeviceCallback}, {@link InfoMediaManager}
+     * will receive state updates from the platform.
+     *
+     * <p>Call {@link #unregisterCallback(MediaDeviceCallback)} once you no longer need platform
+     * updates.
+     */
+    public final void registerCallback(@NonNull MediaDeviceCallback callback) {
+        boolean wasEmpty = mCallbacks.isEmpty();
         if (!mCallbacks.contains(callback)) {
             mCallbacks.add(callback);
+            if (wasEmpty) {
+                mMediaDevices.clear();
+                registerRouter();
+                updateRouteListingPreference();
+                refreshDevices();
+            }
         }
     }
 
-    protected final void unregisterCallback(MediaDeviceCallback callback) {
-        mCallbacks.remove(callback);
+    /**
+     * Unregisters the specified {@code callback}.
+     *
+     * @see #registerCallback(MediaDeviceCallback)
+     */
+    public final void unregisterCallback(@NonNull MediaDeviceCallback callback) {
+        if (mCallbacks.remove(callback) && mCallbacks.isEmpty()) {
+            unregisterRouter();
+        }
     }
 
     private void dispatchDeviceListAdded(@NonNull List<MediaDevice> devices) {
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/LocalMediaManager.java b/packages/SettingsLib/src/com/android/settingslib/media/LocalMediaManager.java
index 0c2414c..a08a3dc 100644
--- a/packages/SettingsLib/src/com/android/settingslib/media/LocalMediaManager.java
+++ b/packages/SettingsLib/src/com/android/settingslib/media/LocalMediaManager.java
@@ -17,7 +17,6 @@
 
 import static android.media.MediaRoute2ProviderService.REASON_UNKNOWN_ERROR;
 
-import android.app.Notification;
 import android.bluetooth.BluetoothAdapter;
 import android.bluetooth.BluetoothDevice;
 import android.content.ComponentName;
@@ -106,14 +105,23 @@
      * Register to start receiving callbacks for MediaDevice events.
      */
     public void registerCallback(DeviceCallback callback) {
-        mCallbacks.add(callback);
+        boolean wasEmpty = mCallbacks.isEmpty();
+        if (!mCallbacks.contains(callback)) {
+            mCallbacks.add(callback);
+            if (wasEmpty) {
+                mInfoMediaManager.registerCallback(mMediaDeviceCallback);
+            }
+        }
     }
 
     /**
      * Unregister to stop receiving callbacks for MediaDevice events
      */
     public void unregisterCallback(DeviceCallback callback) {
-        mCallbacks.remove(callback);
+        if (mCallbacks.remove(callback) && mCallbacks.isEmpty()) {
+            mInfoMediaManager.unregisterCallback(mMediaDeviceCallback);
+            unRegisterDeviceAttributeChangeCallback();
+        }
     }
 
     /**
@@ -125,7 +133,7 @@
      *
      * It will use {@link BluetoothAdapter#getDefaultAdapter()] for setting the bluetooth adapter.
      */
-    public LocalMediaManager(Context context, String packageName, Notification notification) {
+    public LocalMediaManager(Context context, String packageName) {
         mContext = context;
         mPackageName = packageName;
         mLocalBluetoothManager =
@@ -228,10 +236,6 @@
      * Start scan connected MediaDevice
      */
     public void startScan() {
-        synchronized (mMediaDevicesLock) {
-            mMediaDevices.clear();
-        }
-        mInfoMediaManager.registerCallback(mMediaDeviceCallback);
         mInfoMediaManager.startScan();
     }
 
@@ -281,9 +285,7 @@
      * Stop scan MediaDevice
      */
     public void stopScan() {
-        mInfoMediaManager.unregisterCallback(mMediaDeviceCallback);
         mInfoMediaManager.stopScan();
-        unRegisterDeviceAttributeChangeCallback();
     }
 
     /**
diff --git a/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/LocalMediaRepository.kt b/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/LocalMediaRepository.kt
index 724dd51..869fb7f 100644
--- a/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/LocalMediaRepository.kt
+++ b/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/LocalMediaRepository.kt
@@ -17,6 +17,7 @@
 
 import com.android.settingslib.media.LocalMediaManager
 import com.android.settingslib.media.MediaDevice
+import com.android.settingslib.media.flags.Flags
 import com.android.settingslib.volume.shared.AudioManagerEventsReceiver
 import com.android.settingslib.volume.shared.model.AudioManagerEvent
 import kotlinx.coroutines.CoroutineScope
@@ -69,10 +70,14 @@
                         }
                     }
                 localMediaManager.registerCallback(callback)
-                localMediaManager.startScan()
+                if (!Flags.removeUnnecessaryRouteScanning()) {
+                    localMediaManager.startScan()
+                }
 
                 awaitClose {
-                    localMediaManager.stopScan()
+                    if (!Flags.removeUnnecessaryRouteScanning()) {
+                        localMediaManager.stopScan()
+                    }
                     localMediaManager.unregisterCallback(callback)
                 }
             }
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java
index a4b87da..69faddf 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InfoMediaManagerTest.java
@@ -35,6 +35,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -751,35 +752,31 @@
 
     @Test
     public void onTransferred_getAvailableRoutes_shouldAddMediaDevice() {
-        final List<RoutingSessionInfo> routingSessionInfos = new ArrayList<>();
-        final RoutingSessionInfo sessionInfo = mock(RoutingSessionInfo.class);
-        routingSessionInfos.add(sessionInfo);
-        final List<String> selectedRoutes = new ArrayList<>();
-        selectedRoutes.add(TEST_ID);
-        when(sessionInfo.getSelectedRoutes()).thenReturn(selectedRoutes);
-        mShadowRouter2Manager.setRoutingSessions(routingSessionInfos);
+        mInfoMediaManager.mRouterManager = mRouterManager;
+        // Since test is running in Robolectric, return a fake session to avoid NPE.
+        when(mRouterManager.getRoutingSessions(anyString()))
+                .thenReturn(List.of(TEST_SYSTEM_ROUTING_SESSION));
+        when(mRouterManager.getSelectedRoutes(any()))
+                .thenReturn(List.of(TEST_SELECTED_SYSTEM_ROUTE));
 
-        final MediaRoute2Info info = mock(MediaRoute2Info.class);
         mInfoMediaManager.registerCallback(mCallback);
 
-        when(info.getDeduplicationIds()).thenReturn(Set.of());
-        when(info.getId()).thenReturn(TEST_ID);
-        when(info.getClientPackageName()).thenReturn(TEST_PACKAGE_NAME);
+        MediaDevice mediaDevice = mInfoMediaManager.getCurrentConnectedDevice();
+        assertThat(mediaDevice).isNotNull();
+        assertThat(mediaDevice.getId()).isEqualTo(TEST_SYSTEM_ROUTE_ID);
 
-        final List<MediaRoute2Info> routes = new ArrayList<>();
-        routes.add(info);
-        mShadowRouter2Manager.setTransferableRoutes(routes);
+        when(mRouterManager.getRoutingSessions(anyString()))
+                .thenReturn(List.of(TEST_SYSTEM_ROUTING_SESSION, TEST_REMOTE_ROUTING_SESSION));
+        when(mRouterManager.getSelectedRoutes(any())).thenReturn(List.of(TEST_REMOTE_ROUTE));
 
-        final MediaDevice mediaDevice = mInfoMediaManager.findMediaDevice(TEST_ID);
-        assertThat(mediaDevice).isNull();
-
-        mInfoMediaManager.mMediaRouterCallback.onTransferred(sessionInfo, sessionInfo);
+        mInfoMediaManager.mMediaRouterCallback.onTransferred(
+                TEST_SYSTEM_ROUTING_SESSION, TEST_REMOTE_ROUTING_SESSION);
 
         final MediaDevice infoDevice = mInfoMediaManager.mMediaDevices.get(0);
-        assertThat(infoDevice.getId()).isEqualTo(TEST_ID);
+        assertThat(infoDevice).isNotNull();
+        assertThat(infoDevice.getId()).isEqualTo(TEST_REMOTE_ROUTE.getId());
         assertThat(mInfoMediaManager.getCurrentConnectedDevice()).isEqualTo(infoDevice);
-        assertThat(mInfoMediaManager.mMediaDevices).hasSize(routes.size());
-        verify(mCallback).onConnectedDeviceChanged(TEST_ID);
+        verify(mCallback).onConnectedDeviceChanged(TEST_REMOTE_ROUTE.getId());
     }
 
     @Test
@@ -795,7 +792,8 @@
 
         mInfoMediaManager.mMediaRouterCallback.onSessionUpdated(TEST_SYSTEM_ROUTING_SESSION);
 
-        verify(mCallback).onDeviceListAdded(any());
+        // Expecting 1st call after registerCallback() and 2nd call after onSessionUpdated().
+        verify(mCallback, times(2)).onDeviceListAdded(any());
     }
 
     @Test
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/LocalMediaManagerTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/LocalMediaManagerTest.java
index 693b7d0..ddb5419 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/LocalMediaManagerTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/LocalMediaManagerTest.java
@@ -21,6 +21,8 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
@@ -58,6 +60,7 @@
 import org.robolectric.shadow.api.Shadow;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 @RunWith(RobolectricTestRunner.class)
@@ -117,6 +120,13 @@
         mInfoMediaManager = mock(
                 InfoMediaManager.class,
                 withSettings().useConstructor(mContext, TEST_PACKAGE_NAME, mLocalBluetoothManager));
+        doReturn(
+                        List.of(
+                                new RoutingSessionInfo.Builder(TEST_SESSION_ID, TEST_PACKAGE_NAME)
+                                        .addSelectedRoute(TEST_DEVICE_ID_1)
+                                        .build()))
+                .when(mInfoMediaManager)
+                .getRoutingSessionsForPackage();
 
         mInfoMediaDevice1 = spy(new InfoMediaDevice(mContext, mRouteInfo1));
         mInfoMediaDevice2 = new InfoMediaDevice(mContext, mRouteInfo2);
@@ -124,15 +134,16 @@
                 new LocalMediaManager(
                         mContext, mLocalBluetoothManager, mInfoMediaManager, TEST_PACKAGE_NAME);
         mLocalMediaManager.mAudioManager = mAudioManager;
+        mLocalMediaManager.registerCallback(mCallback);
+        clearInvocations(mCallback);
     }
 
     @Test
-    public void startScan_mediaDevicesListShouldBeClear() {
+    public void onDeviceListAdded_shouldClearDeviceList() {
         final MediaDevice device = mock(MediaDevice.class);
         mLocalMediaManager.mMediaDevices.add(device);
-
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(1);
-        mLocalMediaManager.startScan();
+        mLocalMediaManager.mMediaDeviceCallback.onDeviceListAdded(Collections.emptyList());
         assertThat(mLocalMediaManager.mMediaDevices).isEmpty();
     }
 
@@ -147,7 +158,6 @@
         when(device.getId()).thenReturn(TEST_DEVICE_ID_1);
         when(currentDevice.getId()).thenReturn(TEST_CURRENT_DEVICE_ID);
 
-        mLocalMediaManager.registerCallback(mCallback);
         assertThat(mLocalMediaManager.connectDevice(device)).isTrue();
         verify(mInfoMediaManager).connectToDevice(device);
     }
@@ -158,7 +168,6 @@
         mLocalMediaManager.mMediaDevices.add(mInfoMediaDevice2);
         mLocalMediaManager.mCurrentConnectedDevice = mInfoMediaDevice1;
 
-        mLocalMediaManager.registerCallback(mCallback);
         assertThat(mLocalMediaManager.connectDevice(mInfoMediaDevice2)).isTrue();
 
         assertThat(mInfoMediaDevice2.getState()).isEqualTo(LocalMediaManager.MediaDeviceState
@@ -171,7 +180,6 @@
         mLocalMediaManager.mMediaDevices.add(mInfoMediaDevice2);
         mLocalMediaManager.mCurrentConnectedDevice = mInfoMediaDevice1;
 
-        mLocalMediaManager.registerCallback(mCallback);
         assertThat(mLocalMediaManager.connectDevice(mInfoMediaDevice1)).isFalse();
 
         assertThat(mInfoMediaDevice1.getState()).isNotEqualTo(LocalMediaManager.MediaDeviceState
@@ -189,7 +197,6 @@
         when(cachedDevice.isConnected()).thenReturn(false);
         when(cachedDevice.isBusy()).thenReturn(false);
 
-        mLocalMediaManager.registerCallback(mCallback);
         assertThat(mLocalMediaManager.connectDevice(device)).isTrue();
 
         verify(cachedDevice).connect();
@@ -252,7 +259,6 @@
         when(device2.getId()).thenReturn(TEST_DEVICE_ID_2);
 
         assertThat(mLocalMediaManager.mMediaDevices).isEmpty();
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onDeviceListAdded(devices);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(2);
@@ -274,7 +280,6 @@
         when(device3.getId()).thenReturn(TEST_DEVICE_ID_3);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(1);
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onDeviceListAdded(devices);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(2);
@@ -292,7 +297,6 @@
         mLocalMediaManager.mMediaDevices.add(device2);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(2);
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback
                 .onDeviceListRemoved(mLocalMediaManager.mMediaDevices);
 
@@ -302,19 +306,21 @@
 
     @Test
     public void onDeviceListRemoved_phoneDeviceNotLastDeviceAfterRemoveDeviceList_removeList() {
-        final List<MediaDevice> devices = new ArrayList<>();
         final MediaDevice device1 = mock(MediaDevice.class);
         final MediaDevice device2 = mock(MediaDevice.class);
         final MediaDevice device3 = mock(MediaDevice.class);
-        devices.add(device1);
-        devices.add(device3);
+
+        mLocalMediaManager.mMediaDevices.clear();
         mLocalMediaManager.mMediaDevices.add(device1);
         mLocalMediaManager.mMediaDevices.add(device2);
         mLocalMediaManager.mMediaDevices.add(device3);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(3);
-        mLocalMediaManager.registerCallback(mCallback);
-        mLocalMediaManager.mMediaDeviceCallback.onDeviceListRemoved(devices);
+
+        final List<MediaDevice> devicesToRemove = new ArrayList<>();
+        devicesToRemove.add(device1);
+        devicesToRemove.add(device3);
+        mLocalMediaManager.mMediaDeviceCallback.onDeviceListRemoved(devicesToRemove);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(1);
         verify(mCallback).onDeviceListUpdate(any());
@@ -332,7 +338,6 @@
         when(device1.getId()).thenReturn(TEST_DEVICE_ID_1);
         when(device2.getId()).thenReturn(TEST_DEVICE_ID_2);
 
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onConnectedDeviceChanged(TEST_DEVICE_ID_2);
 
         assertThat(mLocalMediaManager.getCurrentConnectedDevice()).isEqualTo(device2);
@@ -352,7 +357,6 @@
         when(device1.getId()).thenReturn(TEST_DEVICE_ID_1);
         when(device2.getId()).thenReturn(TEST_DEVICE_ID_2);
 
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onConnectedDeviceChanged(TEST_DEVICE_ID_1);
 
         verify(mCallback, never()).onDeviceAttributesChanged();
@@ -366,7 +370,6 @@
 
         assertThat(mInfoMediaDevice1.getState()).isEqualTo(LocalMediaManager.MediaDeviceState
                 .STATE_DISCONNECTED);
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onConnectedDeviceChanged(TEST_DEVICE_ID_1);
 
         assertThat(mInfoMediaDevice1.getState()).isEqualTo(LocalMediaManager.MediaDeviceState
@@ -375,7 +378,6 @@
 
     @Test
     public void onConnectedDeviceChanged_nullConnectedDevice_noException() {
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onConnectedDeviceChanged(TEST_DEVICE_ID_2);
     }
 
@@ -392,7 +394,6 @@
         when(cachedDevice.isConnected()).thenReturn(false);
         when(cachedDevice.isBusy()).thenReturn(false);
 
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.connectDevice(device);
 
         mLocalMediaManager.mDeviceAttributeChangeCallback.onDeviceAttributesChanged();
@@ -410,7 +411,6 @@
                 .STATE_CONNECTING);
         assertThat(mInfoMediaDevice2.getState()).isEqualTo(LocalMediaManager.MediaDeviceState
                 .STATE_CONNECTED);
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onRequestFailed(REASON_UNKNOWN_ERROR);
 
         assertThat(mInfoMediaDevice1.getState()).isEqualTo(LocalMediaManager.MediaDeviceState
@@ -421,8 +421,6 @@
 
     @Test
     public void onDeviceAttributesChanged_shouldBeCalled() {
-        mLocalMediaManager.registerCallback(mCallback);
-
         mLocalMediaManager.mDeviceAttributeChangeCallback.onDeviceAttributesChanged();
 
         verify(mCallback).onDeviceAttributesChanged();
@@ -475,7 +473,6 @@
         when(device1.getDeviceType()).thenReturn(MediaDevice.MediaDeviceType.TYPE_PHONE_DEVICE);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(0);
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onDeviceListAdded(devices);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(2);
@@ -497,7 +494,6 @@
         when(cachedDevice.isConnected()).thenReturn(false);
         when(cachedDevice.isBusy()).thenReturn(false);
 
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.connectDevice(device);
 
         verify(cachedDevice).connect();
@@ -512,8 +508,6 @@
 
     @Test
     public void onRequestFailed_shouldDispatchOnRequestFailed() {
-        mLocalMediaManager.registerCallback(mCallback);
-
         mLocalMediaManager.mMediaDeviceCallback.onRequestFailed(1);
 
         verify(mCallback).onRequestFailed(1);
@@ -532,7 +526,6 @@
         mShadowBluetoothAdapter = null;
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(1);
-        mLocalMediaManager.registerCallback(mCallback);
         mLocalMediaManager.mMediaDeviceCallback.onDeviceListAdded(devices);
 
         assertThat(mLocalMediaManager.mMediaDevices).hasSize(2);
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/fold/ui/composable/UnfoldModifiers.kt b/packages/SystemUI/compose/features/src/com/android/systemui/fold/ui/composable/UnfoldModifiers.kt
deleted file mode 100644
index c2a2696..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/fold/ui/composable/UnfoldModifiers.kt
+++ /dev/null
@@ -1,77 +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.fold.ui.composable
-
-import androidx.annotation.FloatRange
-import androidx.compose.runtime.Composable
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.graphics.graphicsLayer
-import androidx.compose.ui.platform.LocalLayoutDirection
-import androidx.compose.ui.unit.Dp
-import androidx.compose.ui.unit.LayoutDirection
-import com.android.compose.modifiers.padding
-import kotlin.math.roundToInt
-
-/**
- * Applies a translation that feeds off of the unfold transition that's active while the device is
- * being folded or unfolded, effectively shifting the element towards the fold hinge.
- *
- * @param startSide `true` if the affected element is on the start side (left-hand side in
- *   left-to-right layouts), `false` otherwise.
- * @param fullTranslation The maximum translation to apply when the element is the most shifted. The
- *   modifier will never apply more than this much translation on the element.
- * @param unfoldProgress A provider for the amount of progress of the unfold transition. This should
- *   be sourced from the `UnfoldTransitionInteractor`, ideally through a view-model.
- */
-@Composable
-fun Modifier.unfoldTranslation(
-    startSide: Boolean,
-    fullTranslation: Dp,
-    @FloatRange(from = 0.0, to = 1.0) unfoldProgress: () -> Float,
-): Modifier {
-    val translateToTheRight = startSide && LocalLayoutDirection.current == LayoutDirection.Ltr
-    return this.graphicsLayer {
-        translationX =
-            fullTranslation.toPx() *
-                if (translateToTheRight) {
-                    1 - unfoldProgress()
-                } else {
-                    unfoldProgress() - 1
-                }
-    }
-}
-
-/**
- * Applies horizontal padding that feeds off of the unfold transition that's active while the device
- * is being folded or unfolded, effectively "squishing" the element on both sides.
- *
- * This is horizontal padding so it's applied on both the start and end sides of the element.
- *
- * @param fullPadding The maximum padding to apply when the element is the most padded. The modifier
- *   will never apply more than this much horizontal padding on the element.
- * @param unfoldProgress A provider for the amount of progress of the unfold transition. This should
- *   be sourced from the `UnfoldTransitionInteractor`, ideally through a view-model.
- */
-@Composable
-fun Modifier.unfoldHorizontalPadding(
-    fullPadding: Dp,
-    @FloatRange(from = 0.0, to = 1.0) unfoldProgress: () -> Float,
-): Modifier {
-    return this.padding(
-        horizontal = { (fullPadding.toPx() * (1 - unfoldProgress())).roundToInt() },
-    )
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
index 01c27a4d..84b1a4b 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
@@ -62,11 +62,10 @@
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.compose.animation.scene.animateSceneFloatAsState
+import com.android.compose.modifiers.padding
 import com.android.compose.modifiers.thenIf
 import com.android.systemui.battery.BatteryMeterViewController
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.fold.ui.composable.unfoldHorizontalPadding
-import com.android.systemui.fold.ui.composable.unfoldTranslation
 import com.android.systemui.media.controls.ui.composable.MediaCarousel
 import com.android.systemui.media.controls.ui.controller.MediaCarouselController
 import com.android.systemui.media.controls.ui.controller.MediaHierarchyManager
@@ -291,7 +290,18 @@
         remember(lifecycleOwner, viewModel) { viewModel.getFooterActionsViewModel(lifecycleOwner) }
     val tileSquishiness by
         animateSceneFloatAsState(value = 1f, key = QuickSettings.SharedValues.TilesSquishiness)
-    val unfoldTransitionProgress by viewModel.unfoldTransitionProgress.collectAsState()
+    val unfoldTranslationXForStartSide by
+        viewModel
+            .unfoldTranslationX(
+                isOnStartSide = true,
+            )
+            .collectAsState(0f)
+    val unfoldTranslationXForEndSide by
+        viewModel
+            .unfoldTranslationX(
+                isOnStartSide = false,
+            )
+            .collectAsState(0f)
 
     val navBarBottomHeight = WindowInsets.navigationBars.asPaddingValues().calculateBottomPadding()
     val density = LocalDensity.current
@@ -340,21 +350,16 @@
                 modifier =
                     Modifier.padding(horizontal = Shade.Dimensions.HorizontalPadding)
                         .then(brightnessMirrorShowingModifier)
-                        .unfoldHorizontalPadding(
-                            fullPadding = dimensionResource(R.dimen.notification_side_paddings),
-                        ) {
-                            unfoldTransitionProgress
-                        }
+                        .padding(
+                            horizontal = { unfoldTranslationXForStartSide.roundToInt() },
+                        )
             )
 
             Row(modifier = Modifier.fillMaxWidth().weight(1f)) {
                 Box(
                     modifier =
-                        Modifier.weight(1f).unfoldTranslation(
-                            startSide = true,
-                            fullTranslation = dimensionResource(R.dimen.notification_side_paddings),
-                        ) {
-                            unfoldTransitionProgress
+                        Modifier.weight(1f).graphicsLayer {
+                            translationX = unfoldTranslationXForStartSide
                         },
                 ) {
                     BrightnessMirror(
@@ -424,15 +429,6 @@
                             .fillMaxHeight()
                             .padding(bottom = navBarBottomHeight)
                             .then(brightnessMirrorShowingModifier)
-                            .unfoldTranslation(
-                                startSide = false,
-                                fullTranslation =
-                                    dimensionResource(
-                                        R.dimen.notification_side_paddings,
-                                    ),
-                            ) {
-                                unfoldTransitionProgress
-                            },
                 )
             }
         }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ButtonComponent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ButtonComponent.kt
index e15d315..0893b9d 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ButtonComponent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ButtonComponent.kt
@@ -69,9 +69,19 @@
                             role = Role.Button
                             contentDescription = label
                         },
-                    color = MaterialTheme.colorScheme.tertiaryContainer,
+                    color =
+                        if (viewModel.isActive) {
+                            MaterialTheme.colorScheme.tertiaryContainer
+                        } else {
+                            MaterialTheme.colorScheme.surface
+                        },
                     shape = RoundedCornerShape(28.dp),
-                    contentColor = MaterialTheme.colorScheme.onTertiaryContainer,
+                    contentColor =
+                        if (viewModel.isActive) {
+                            MaterialTheme.colorScheme.onTertiaryContainer
+                        } else {
+                            MaterialTheme.colorScheme.onSurface
+                        },
                     onClick = onClick,
                 ) {
                     Box(modifier = Modifier.fillMaxSize(), contentAlignment = Alignment.Center) {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ToggleButtonComponent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ToggleButtonComponent.kt
index b2351c4..4f3a6c8 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ToggleButtonComponent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/button/ui/composable/ToggleButtonComponent.kt
@@ -40,14 +40,14 @@
 import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.unit.dp
 import com.android.systemui.common.ui.compose.Icon
-import com.android.systemui.volume.panel.component.button.ui.viewmodel.ToggleButtonViewModel
+import com.android.systemui.volume.panel.component.button.ui.viewmodel.ButtonViewModel
 import com.android.systemui.volume.panel.ui.composable.ComposeVolumePanelUiComponent
 import com.android.systemui.volume.panel.ui.composable.VolumePanelComposeScope
 import kotlinx.coroutines.flow.StateFlow
 
 /** [ComposeVolumePanelUiComponent] implementing a toggleable button from a bottom row. */
 class ToggleButtonComponent(
-    private val viewModelFlow: StateFlow<ToggleButtonViewModel?>,
+    private val viewModelFlow: StateFlow<ButtonViewModel?>,
     private val onCheckedChange: (isChecked: Boolean) -> Unit
 ) : ComposeVolumePanelUiComponent {
 
@@ -64,7 +64,7 @@
         ) {
             BottomComponentButtonSurface {
                 val colors =
-                    if (viewModel.isChecked) {
+                    if (viewModel.isActive) {
                         ButtonDefaults.buttonColors(
                             containerColor = MaterialTheme.colorScheme.tertiaryContainer,
                             contentColor = MaterialTheme.colorScheme.onTertiaryContainer,
@@ -81,7 +81,7 @@
                             role = Role.Switch
                             contentDescription = label
                         },
-                    onClick = { onCheckedChange(!viewModel.isChecked) },
+                    onClick = { onCheckedChange(!viewModel.isActive) },
                     shape = RoundedCornerShape(28.dp),
                     colors = colors
                 ) {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
index f377fa6..12d2bc2 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/spatialaudio/ui/composable/SpatialAudioPopup.kt
@@ -52,7 +52,7 @@
             VolumePanelUiEvent.VOLUME_PANEL_SPATIAL_AUDIO_POP_UP_SHOWN,
             0,
             null,
-            viewModel.spatialAudioButtons.value.indexOfFirst { it.button.isChecked }
+            viewModel.spatialAudioButtons.value.indexOfFirst { it.button.isActive }
         )
         volumePanelPopup.show(expandable, { Title() }, { Content(it) })
     }
@@ -85,7 +85,7 @@
             for (buttonViewModel in enabledModelStates) {
                 val label = buttonViewModel.button.label.toString()
                 item(
-                    isSelected = buttonViewModel.button.isChecked,
+                    isSelected = buttonViewModel.button.isActive,
                     onItemSelected = { viewModel.setEnabled(buttonViewModel.model) },
                     contentDescription = label,
                     icon = { Icon(icon = buttonViewModel.button.icon) },
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/ui/composable/VolumePanelRoot.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/ui/composable/VolumePanelRoot.kt
index 910cd5e..1bf541a 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/ui/composable/VolumePanelRoot.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/ui/composable/VolumePanelRoot.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.volume.panel.ui.composable
 
-import androidx.compose.foundation.clickable
+import androidx.compose.foundation.gestures.detectTapGestures
 import androidx.compose.foundation.isSystemInDarkTheme
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Box
@@ -38,6 +38,7 @@
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.input.pointer.pointerInput
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.platform.LocalLayoutDirection
 import androidx.compose.ui.res.dimensionResource
@@ -75,21 +76,13 @@
                 modifier =
                     modifier
                         .fillMaxSize()
-                        .clickable(onClick = onDismiss)
+                        .volumePanelClick(onDismiss)
                         .volumePanelPaddings(isPortrait = isPortrait),
                 contentAlignment = Alignment.BottomCenter,
             ) {
                 val radius = dimensionResource(R.dimen.volume_panel_corner_radius)
                 Surface(
-                    modifier =
-                        Modifier.clickable(
-                            interactionSource = null,
-                            indication = null,
-                            onClick = {
-                                // prevent windowCloseOnTouchOutside from dismissing when tapped
-                                // on the panel itself.
-                            },
-                        ),
+                    modifier = Modifier.volumePanelClick {},
                     shape = RoundedCornerShape(topStart = radius, topEnd = radius),
                     color = MaterialTheme.colorScheme.surfaceContainer,
                 ) {
@@ -185,3 +178,13 @@
         )
     }
 }
+
+/**
+ * For some reason adding clickable modifier onto the VolumePanel affects the traversal order:
+ * b/331155283.
+ *
+ * TODO(b/334870995) revert this to Modifier.clickable
+ */
+@Composable
+private fun Modifier.volumePanelClick(onClick: () -> Unit) =
+    pointerInput(onClick) { detectTapGestures { onClick() } }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.java
deleted file mode 100644
index 4e18a47..0000000
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.java
+++ /dev/null
@@ -1,559 +0,0 @@
-/*
- * Copyright (C) 2021 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.dreams;
-
-import static android.view.WindowManager.LayoutParams.SYSTEM_FLAG_SHOW_FOR_ALL_USERS;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.clearInvocations;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.doThrow;
-import static org.mockito.Mockito.inOrder;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import android.content.ComponentName;
-import android.content.Intent;
-import android.os.IBinder;
-import android.os.RemoteException;
-import android.service.dreams.IDreamOverlay;
-import android.service.dreams.IDreamOverlayCallback;
-import android.service.dreams.IDreamOverlayClient;
-import android.service.dreams.IDreamOverlayClientCallback;
-import android.view.View;
-import android.view.ViewGroup;
-import android.view.WindowManager;
-import android.view.WindowManagerImpl;
-
-import androidx.lifecycle.Lifecycle;
-import androidx.lifecycle.LifecycleRegistry;
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
-import com.android.internal.logging.UiEventLogger;
-import com.android.keyguard.KeyguardUpdateMonitor;
-import com.android.systemui.SysuiTestCase;
-import com.android.systemui.ambient.touch.TouchMonitor;
-import com.android.systemui.ambient.touch.dagger.AmbientTouchComponent;
-import com.android.systemui.complication.ComplicationLayoutEngine;
-import com.android.systemui.dreams.complication.HideComplicationTouchHandler;
-import com.android.systemui.dreams.complication.dagger.ComplicationComponent;
-import com.android.systemui.dreams.dagger.DreamOverlayComponent;
-import com.android.systemui.touch.TouchInsetManager;
-import com.android.systemui.util.concurrency.FakeExecutor;
-import com.android.systemui.util.time.FakeSystemClock;
-import com.android.systemui.utils.leaks.LeakCheckedTest;
-
-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.InOrder;
-import org.mockito.Mock;
-import org.mockito.Mockito;
-import org.mockito.MockitoAnnotations;
-
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-public class DreamOverlayServiceTest extends SysuiTestCase {
-    private static final ComponentName LOW_LIGHT_COMPONENT = new ComponentName("package",
-            "lowlight");
-
-    private static final ComponentName HOME_CONTROL_PANEL_DREAM_COMPONENT =
-            new ComponentName("package", "homeControlPanel");
-    private static final String DREAM_COMPONENT = "package/dream";
-    private static final String WINDOW_NAME = "test";
-    private final FakeSystemClock mFakeSystemClock = new FakeSystemClock();
-    private final FakeExecutor mMainExecutor = new FakeExecutor(mFakeSystemClock);
-
-    @Mock
-    DreamOverlayLifecycleOwner mLifecycleOwner;
-
-    @Mock
-    LifecycleRegistry mLifecycleRegistry;
-
-    @Rule
-    public final LeakCheckedTest.SysuiLeakCheck mLeakCheck = new LeakCheckedTest.SysuiLeakCheck();
-
-    WindowManager.LayoutParams mWindowParams;
-
-    @Mock
-    IDreamOverlayCallback mDreamOverlayCallback;
-
-    @Mock
-    WindowManagerImpl mWindowManager;
-
-    @Mock
-    com.android.systemui.complication.dagger.ComplicationComponent.Factory
-            mComplicationComponentFactory;
-
-    @Mock
-    com.android.systemui.complication.dagger.ComplicationComponent mComplicationComponent;
-
-    @Mock
-    ComplicationLayoutEngine mComplicationVisibilityController;
-
-    @Mock
-    ComplicationComponent.Factory mDreamComplicationComponentFactory;
-
-    @Mock
-    ComplicationComponent mDreamComplicationComponent;
-
-    @Mock
-    HideComplicationTouchHandler mHideComplicationTouchHandler;
-
-    @Mock
-    DreamOverlayComponent.Factory mDreamOverlayComponentFactory;
-
-    @Mock
-    DreamOverlayComponent mDreamOverlayComponent;
-
-    @Mock
-    AmbientTouchComponent.Factory mAmbientTouchComponentFactory;
-
-    @Mock
-    AmbientTouchComponent mAmbientTouchComponent;
-
-    @Mock
-    DreamOverlayContainerView mDreamOverlayContainerView;
-
-    @Mock
-    DreamOverlayContainerViewController mDreamOverlayContainerViewController;
-
-    @Mock
-    KeyguardUpdateMonitor mKeyguardUpdateMonitor;
-
-    @Mock
-    TouchMonitor mTouchMonitor;
-
-    @Mock
-    DreamOverlayStateController mStateController;
-
-    @Mock
-    ViewGroup mDreamOverlayContainerViewParent;
-
-    @Mock
-    TouchInsetManager mTouchInsetManager;
-
-    @Mock
-    UiEventLogger mUiEventLogger;
-
-    @Mock
-    DreamOverlayCallbackController mDreamOverlayCallbackController;
-
-    @Captor
-    ArgumentCaptor<View> mViewCaptor;
-
-    DreamOverlayService mService;
-
-    @Before
-    public void setup() {
-        MockitoAnnotations.initMocks(this);
-
-        when(mDreamOverlayComponent.getDreamOverlayContainerViewController())
-                .thenReturn(mDreamOverlayContainerViewController);
-        when(mLifecycleOwner.getRegistry())
-                .thenReturn(mLifecycleRegistry);
-        when(mComplicationComponentFactory
-                .create(any(), any(), any(), any()))
-                .thenReturn(mComplicationComponent);
-        when(mComplicationComponent.getVisibilityController())
-                .thenReturn(mComplicationVisibilityController);
-        when(mDreamComplicationComponent.getHideComplicationTouchHandler())
-                .thenReturn(mHideComplicationTouchHandler);
-        when(mDreamComplicationComponentFactory
-                .create(any(), any()))
-                .thenReturn(mDreamComplicationComponent);
-        when(mDreamOverlayComponentFactory
-                .create(any(), any(), any()))
-                .thenReturn(mDreamOverlayComponent);
-        when(mAmbientTouchComponentFactory.create(any(), any())).thenReturn(mAmbientTouchComponent);
-        when(mAmbientTouchComponent.getTouchMonitor())
-                .thenReturn(mTouchMonitor);
-        when(mDreamOverlayContainerViewController.getContainerView())
-                .thenReturn(mDreamOverlayContainerView);
-
-        mWindowParams = new WindowManager.LayoutParams();
-        mService = new DreamOverlayService(
-                mContext,
-                mLifecycleOwner,
-                mMainExecutor,
-                mWindowManager,
-                mComplicationComponentFactory,
-                mDreamComplicationComponentFactory,
-                mDreamOverlayComponentFactory,
-                mAmbientTouchComponentFactory,
-                mStateController,
-                mKeyguardUpdateMonitor,
-                mUiEventLogger,
-                mTouchInsetManager,
-                LOW_LIGHT_COMPONENT,
-                HOME_CONTROL_PANEL_DREAM_COMPONENT,
-                mDreamOverlayCallbackController,
-                WINDOW_NAME);
-    }
-
-    public IDreamOverlayClient getClient() throws RemoteException {
-        final IBinder proxy = mService.onBind(new Intent());
-        final IDreamOverlay overlay = IDreamOverlay.Stub.asInterface(proxy);
-        final IDreamOverlayClientCallback callback =
-                Mockito.mock(IDreamOverlayClientCallback.class);
-        overlay.getClient(callback);
-        final ArgumentCaptor<IDreamOverlayClient> clientCaptor =
-                ArgumentCaptor.forClass(IDreamOverlayClient.class);
-        verify(callback).onDreamOverlayClient(clientCaptor.capture());
-
-        return clientCaptor.getValue();
-    }
-
-    @Test
-    public void testOnStartMetricsLogged() throws Exception {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        verify(mUiEventLogger).log(DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_ENTER_START);
-        verify(mUiEventLogger).log(
-                DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_COMPLETE_START);
-    }
-
-    @Test
-    public void testOverlayContainerViewAddedToWindow() throws Exception {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        verify(mWindowManager).addView(any(), any());
-    }
-
-    // Validates that {@link DreamOverlayService} properly handles the case where the dream's
-    // window is no longer valid by the time start is called.
-    @Test
-    public void testInvalidWindowAddStart() throws Exception {
-        final IDreamOverlayClient client = getClient();
-
-        doThrow(new WindowManager.BadTokenException()).when(mWindowManager).addView(any(), any());
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        verify(mWindowManager).addView(any(), any());
-
-        verify(mStateController).setOverlayActive(false);
-        verify(mStateController).setLowLightActive(false);
-        verify(mStateController).setEntryAnimationsFinished(false);
-
-        verify(mStateController, never()).setOverlayActive(true);
-        verify(mUiEventLogger, never()).log(
-                DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_COMPLETE_START);
-
-        verify(mDreamOverlayCallbackController, never()).onStartDream();
-    }
-
-    @Test
-    public void testDreamOverlayContainerViewControllerInitialized() throws Exception {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        verify(mDreamOverlayContainerViewController).init();
-    }
-
-    @Test
-    public void testDreamOverlayContainerViewRemovedFromOldParentWhenInitialized()
-            throws Exception {
-        when(mDreamOverlayContainerView.getParent())
-                .thenReturn(mDreamOverlayContainerViewParent)
-                .thenReturn(null);
-
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        verify(mDreamOverlayContainerViewParent).removeView(mDreamOverlayContainerView);
-    }
-
-    @Test
-    public void testShouldShowComplicationsSetByStartDream() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                true /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        assertThat(mService.shouldShowComplications()).isTrue();
-    }
-
-    @Test
-    public void testLowLightSetByStartDream() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback,
-                LOW_LIGHT_COMPONENT.flattenToString(), false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        assertThat(mService.getDreamComponent()).isEqualTo(LOW_LIGHT_COMPONENT);
-        verify(mStateController).setLowLightActive(true);
-    }
-
-    @Test
-    public void testHomeControlPanelSetsByStartDream() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback,
-                HOME_CONTROL_PANEL_DREAM_COMPONENT.flattenToString(),
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-        assertThat(mService.getDreamComponent()).isEqualTo(HOME_CONTROL_PANEL_DREAM_COMPONENT);
-        verify(mStateController).setHomeControlPanelActive(true);
-    }
-
-    @Test
-    public void testOnEndDream() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback,
-                LOW_LIGHT_COMPONENT.flattenToString(), false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        // Verify view added.
-        verify(mWindowManager).addView(mViewCaptor.capture(), any());
-
-        // Service destroyed.
-        mService.onEndDream();
-        mMainExecutor.runAllReady();
-
-        // Verify view removed.
-        verify(mWindowManager).removeView(mViewCaptor.getValue());
-
-        // Verify state correctly set.
-        verify(mStateController).setOverlayActive(false);
-        verify(mStateController).setLowLightActive(false);
-        verify(mStateController).setEntryAnimationsFinished(false);
-    }
-
-    @Test
-    public void testImmediateEndDream() throws Exception {
-        final IDreamOverlayClient client = getClient();
-
-        // Start the dream, but don't execute any Runnables put on the executor yet. We delay
-        // executing Runnables as the timing isn't guaranteed and we want to verify that the overlay
-        // starts and finishes in the proper order even if Runnables are delayed.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        // Immediately end the dream.
-        client.endDream();
-        // Run any scheduled Runnables.
-        mMainExecutor.runAllReady();
-
-        // The overlay starts then finishes.
-        InOrder inOrder = inOrder(mWindowManager);
-        inOrder.verify(mWindowManager).addView(mViewCaptor.capture(), any());
-        inOrder.verify(mWindowManager).removeView(mViewCaptor.getValue());
-    }
-
-    @Test
-    public void testEndDreamDuringStartDream() throws Exception {
-        final IDreamOverlayClient client = getClient();
-
-        // Schedule the endDream call in the middle of the startDream implementation, as any
-        // ordering is possible.
-        doAnswer(invocation -> {
-            client.endDream();
-            return null;
-        }).when(mStateController).setOverlayActive(true);
-
-        // Start the dream.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        // The overlay starts then finishes.
-        InOrder inOrder = inOrder(mWindowManager);
-        inOrder.verify(mWindowManager).addView(mViewCaptor.capture(), any());
-        inOrder.verify(mWindowManager).removeView(mViewCaptor.getValue());
-    }
-
-    @Test
-    public void testDestroy() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback,
-                LOW_LIGHT_COMPONENT.flattenToString(), false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        // Verify view added.
-        verify(mWindowManager).addView(mViewCaptor.capture(), any());
-
-        // Service destroyed.
-        mService.onDestroy();
-        mMainExecutor.runAllReady();
-
-        // Verify view removed.
-        verify(mWindowManager).removeView(mViewCaptor.getValue());
-
-        // Verify state correctly set.
-        verify(mKeyguardUpdateMonitor).removeCallback(any());
-        verify(mLifecycleRegistry).setCurrentState(Lifecycle.State.DESTROYED);
-        verify(mStateController).setOverlayActive(false);
-        verify(mStateController).setLowLightActive(false);
-        verify(mStateController).setEntryAnimationsFinished(false);
-    }
-
-    @Test
-    public void testDoNotRemoveViewOnDestroyIfOverlayNotStarted() {
-        // Service destroyed without ever starting dream.
-        mService.onDestroy();
-        mMainExecutor.runAllReady();
-
-        // Verify no view is removed.
-        verify(mWindowManager, never()).removeView(any());
-
-        // Verify state still correctly set.
-        verify(mKeyguardUpdateMonitor).removeCallback(any());
-        verify(mLifecycleRegistry).setCurrentState(Lifecycle.State.DESTROYED);
-        verify(mStateController).setOverlayActive(false);
-        verify(mStateController).setLowLightActive(false);
-    }
-
-    @Test
-    public void testDecorViewNotAddedToWindowAfterDestroy() throws Exception {
-        final IDreamOverlayClient client = getClient();
-
-        // Destroy the service.
-        mService.onDestroy();
-        mMainExecutor.runAllReady();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        verify(mWindowManager, never()).addView(any(), any());
-    }
-
-    @Test
-    public void testNeverRemoveDecorViewIfNotAdded() {
-        // Service destroyed before dream started.
-        mService.onDestroy();
-        mMainExecutor.runAllReady();
-
-        verify(mWindowManager, never()).removeView(any());
-    }
-
-    @Test
-    public void testResetCurrentOverlayWhenConnectedToNewDream() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting. Do not show dream complications.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        // Verify that a new window is added.
-        verify(mWindowManager).addView(mViewCaptor.capture(), any());
-        final View windowDecorView = mViewCaptor.getValue();
-
-        // Assert that the overlay is not showing complications.
-        assertThat(mService.shouldShowComplications()).isFalse();
-
-        clearInvocations(mDreamOverlayComponent);
-        clearInvocations(mAmbientTouchComponent);
-        clearInvocations(mWindowManager);
-
-        // New dream starting with dream complications showing. Note that when a new dream is
-        // binding to the dream overlay service, it receives the same instance of IBinder as the
-        // first one.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                true /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        // Assert that the overlay is showing complications.
-        assertThat(mService.shouldShowComplications()).isTrue();
-
-        // Verify that the old overlay window has been removed, and a new one created.
-        verify(mWindowManager).removeView(windowDecorView);
-        verify(mWindowManager).addView(any(), any());
-
-        // Verify that new instances of overlay container view controller and overlay touch monitor
-        // are created.
-        verify(mDreamOverlayComponent).getDreamOverlayContainerViewController();
-        verify(mAmbientTouchComponent).getTouchMonitor();
-    }
-
-    @Test
-    public void testWakeUp() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                true /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        mService.onWakeUp();
-        verify(mDreamOverlayContainerViewController).wakeUp();
-        verify(mDreamOverlayCallbackController).onWakeUp();
-    }
-
-    @Test
-    public void testWakeUpBeforeStartDoesNothing() {
-        mService.onWakeUp();
-        verify(mDreamOverlayContainerViewController, never()).wakeUp();
-    }
-
-    @Test
-    public void testSystemFlagShowForAllUsersSetOnWindow() throws RemoteException {
-        final IDreamOverlayClient client = getClient();
-
-        // Inform the overlay service of dream starting. Do not show dream complications.
-        client.startDream(mWindowParams, mDreamOverlayCallback, DREAM_COMPONENT,
-                false /*shouldShowComplication*/);
-        mMainExecutor.runAllReady();
-
-        final ArgumentCaptor<WindowManager.LayoutParams> paramsCaptor =
-                ArgumentCaptor.forClass(WindowManager.LayoutParams.class);
-
-        // Verify that a new window is added.
-        verify(mWindowManager).addView(any(), paramsCaptor.capture());
-
-        assertThat((paramsCaptor.getValue().privateFlags & SYSTEM_FLAG_SHOW_FOR_ALL_USERS)
-                == SYSTEM_FLAG_SHOW_FOR_ALL_USERS).isTrue();
-    }
-}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
new file mode 100644
index 0000000..b18a8ec
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
@@ -0,0 +1,573 @@
+/*
+ * 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.dreams
+
+import android.content.ComponentName
+import android.content.Intent
+import android.os.RemoteException
+import android.service.dreams.IDreamOverlay
+import android.service.dreams.IDreamOverlayCallback
+import android.service.dreams.IDreamOverlayClient
+import android.service.dreams.IDreamOverlayClientCallback
+import android.view.View
+import android.view.ViewGroup
+import android.view.WindowManager
+import android.view.WindowManagerImpl
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleRegistry
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.logging.UiEventLogger
+import com.android.keyguard.KeyguardUpdateMonitor
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.ambient.touch.TouchMonitor
+import com.android.systemui.ambient.touch.dagger.AmbientTouchComponent
+import com.android.systemui.complication.ComplicationHostViewController
+import com.android.systemui.complication.ComplicationLayoutEngine
+import com.android.systemui.complication.dagger.ComplicationComponent
+import com.android.systemui.dreams.complication.HideComplicationTouchHandler
+import com.android.systemui.dreams.dagger.DreamOverlayComponent
+import com.android.systemui.touch.TouchInsetManager
+import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.time.FakeSystemClock
+import com.google.common.truth.Truth
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.Mockito
+import org.mockito.MockitoAnnotations
+import org.mockito.invocation.InvocationOnMock
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class DreamOverlayServiceTest : SysuiTestCase() {
+    private val mFakeSystemClock = FakeSystemClock()
+    private val mMainExecutor = FakeExecutor(mFakeSystemClock)
+
+    @Mock lateinit var mLifecycleOwner: DreamOverlayLifecycleOwner
+
+    @Mock lateinit var mLifecycleRegistry: LifecycleRegistry
+
+    lateinit var mWindowParams: WindowManager.LayoutParams
+
+    @Mock lateinit var mDreamOverlayCallback: IDreamOverlayCallback
+
+    @Mock lateinit var mWindowManager: WindowManagerImpl
+
+    @Mock lateinit var mComplicationComponentFactory: ComplicationComponent.Factory
+
+    @Mock lateinit var mComplicationComponent: ComplicationComponent
+
+    @Mock lateinit var mComplicationHostViewController: ComplicationHostViewController
+
+    @Mock lateinit var mComplicationVisibilityController: ComplicationLayoutEngine
+
+    @Mock
+    lateinit var mDreamComplicationComponentFactory:
+        com.android.systemui.dreams.complication.dagger.ComplicationComponent.Factory
+
+    @Mock
+    lateinit var mDreamComplicationComponent:
+        com.android.systemui.dreams.complication.dagger.ComplicationComponent
+
+    @Mock lateinit var mHideComplicationTouchHandler: HideComplicationTouchHandler
+
+    @Mock lateinit var mDreamOverlayComponentFactory: DreamOverlayComponent.Factory
+
+    @Mock lateinit var mDreamOverlayComponent: DreamOverlayComponent
+
+    @Mock lateinit var mAmbientTouchComponentFactory: AmbientTouchComponent.Factory
+
+    @Mock lateinit var mAmbientTouchComponent: AmbientTouchComponent
+
+    @Mock lateinit var mDreamOverlayContainerView: DreamOverlayContainerView
+
+    @Mock lateinit var mDreamOverlayContainerViewController: DreamOverlayContainerViewController
+
+    @Mock lateinit var mKeyguardUpdateMonitor: KeyguardUpdateMonitor
+
+    @Mock lateinit var mTouchMonitor: TouchMonitor
+
+    @Mock lateinit var mStateController: DreamOverlayStateController
+
+    @Mock lateinit var mDreamOverlayContainerViewParent: ViewGroup
+
+    @Mock lateinit var mTouchInsetManager: TouchInsetManager
+
+    @Mock lateinit var mUiEventLogger: UiEventLogger
+
+    @Mock lateinit var mDreamOverlayCallbackController: DreamOverlayCallbackController
+
+    @Captor var mViewCaptor: ArgumentCaptor<View>? = null
+    var mService: DreamOverlayService? = null
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        whenever(mDreamOverlayComponent.getDreamOverlayContainerViewController())
+            .thenReturn(mDreamOverlayContainerViewController)
+        whenever(mComplicationComponent.getComplicationHostViewController())
+            .thenReturn(mComplicationHostViewController)
+        whenever(mLifecycleOwner.registry).thenReturn(mLifecycleRegistry)
+        whenever(mComplicationComponentFactory.create(any(), any(), any(), any()))
+            .thenReturn(mComplicationComponent)
+        whenever(mComplicationComponent.getVisibilityController())
+            .thenReturn(mComplicationVisibilityController)
+        whenever(mDreamComplicationComponent.getHideComplicationTouchHandler())
+            .thenReturn(mHideComplicationTouchHandler)
+        whenever(mDreamComplicationComponentFactory.create(any(), any()))
+            .thenReturn(mDreamComplicationComponent)
+        whenever(mDreamOverlayComponentFactory.create(any(), any(), any()))
+            .thenReturn(mDreamOverlayComponent)
+        whenever(mAmbientTouchComponentFactory.create(any(), any()))
+            .thenReturn(mAmbientTouchComponent)
+        whenever(mAmbientTouchComponent.getTouchMonitor()).thenReturn(mTouchMonitor)
+        whenever(mDreamOverlayContainerViewController.containerView)
+            .thenReturn(mDreamOverlayContainerView)
+        mWindowParams = WindowManager.LayoutParams()
+        mService =
+            DreamOverlayService(
+                mContext,
+                mLifecycleOwner,
+                mMainExecutor,
+                mWindowManager,
+                mComplicationComponentFactory,
+                mDreamComplicationComponentFactory,
+                mDreamOverlayComponentFactory,
+                mAmbientTouchComponentFactory,
+                mStateController,
+                mKeyguardUpdateMonitor,
+                mUiEventLogger,
+                mTouchInsetManager,
+                LOW_LIGHT_COMPONENT,
+                HOME_CONTROL_PANEL_DREAM_COMPONENT,
+                mDreamOverlayCallbackController,
+                WINDOW_NAME
+            )
+    }
+
+    @get:Throws(RemoteException::class)
+    val client: IDreamOverlayClient
+        get() {
+            val proxy = mService!!.onBind(Intent())
+            val overlay = IDreamOverlay.Stub.asInterface(proxy)
+            val callback = Mockito.mock(IDreamOverlayClientCallback::class.java)
+            overlay.getClient(callback)
+            val clientCaptor = ArgumentCaptor.forClass(IDreamOverlayClient::class.java)
+            Mockito.verify(callback).onDreamOverlayClient(clientCaptor.capture())
+            return clientCaptor.value
+        }
+
+    @Test
+    fun testOnStartMetricsLogged() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Mockito.verify(mUiEventLogger)
+            .log(DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_ENTER_START)
+        Mockito.verify(mUiEventLogger)
+            .log(DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_COMPLETE_START)
+    }
+
+    @Test
+    fun testOverlayContainerViewAddedToWindow() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Mockito.verify(mWindowManager).addView(any(), any())
+    }
+
+    // Validates that {@link DreamOverlayService} properly handles the case where the dream's
+    // window is no longer valid by the time start is called.
+    @Test
+    fun testInvalidWindowAddStart() {
+        val client = client
+        Mockito.doThrow(WindowManager.BadTokenException())
+            .`when`(mWindowManager)
+            .addView(any(), any())
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Mockito.verify(mWindowManager).addView(any(), any())
+        Mockito.verify(mStateController).setOverlayActive(false)
+        Mockito.verify(mStateController).setLowLightActive(false)
+        Mockito.verify(mStateController).setEntryAnimationsFinished(false)
+        Mockito.verify(mStateController, Mockito.never()).setOverlayActive(true)
+        Mockito.verify(mUiEventLogger, Mockito.never())
+            .log(DreamOverlayService.DreamOverlayEvent.DREAM_OVERLAY_COMPLETE_START)
+        Mockito.verify(mDreamOverlayCallbackController, Mockito.never()).onStartDream()
+    }
+
+    @Test
+    fun testDreamOverlayContainerViewControllerInitialized() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Mockito.verify(mDreamOverlayContainerViewController).init()
+    }
+
+    @Test
+    fun testDreamOverlayContainerViewRemovedFromOldParentWhenInitialized() {
+        whenever(mDreamOverlayContainerView.parent)
+            .thenReturn(mDreamOverlayContainerViewParent)
+            .thenReturn(null)
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Mockito.verify(mDreamOverlayContainerViewParent).removeView(mDreamOverlayContainerView)
+    }
+
+    @Test
+    fun testShouldShowComplicationsSetByStartDream() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            true /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Truth.assertThat(mService!!.shouldShowComplications()).isTrue()
+    }
+
+    @Test
+    fun testLowLightSetByStartDream() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            LOW_LIGHT_COMPONENT.flattenToString(),
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Truth.assertThat(mService!!.dreamComponent).isEqualTo(LOW_LIGHT_COMPONENT)
+        Mockito.verify(mStateController).setLowLightActive(true)
+    }
+
+    @Test
+    fun testHomeControlPanelSetsByStartDream() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            HOME_CONTROL_PANEL_DREAM_COMPONENT.flattenToString(),
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Truth.assertThat(mService!!.dreamComponent).isEqualTo(HOME_CONTROL_PANEL_DREAM_COMPONENT)
+        Mockito.verify(mStateController).setHomeControlPanelActive(true)
+    }
+
+    @Test
+    fun testOnEndDream() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            LOW_LIGHT_COMPONENT.flattenToString(),
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        // Verify view added.
+        Mockito.verify(mWindowManager).addView(mViewCaptor!!.capture(), any())
+
+        // Service destroyed.
+        mService!!.onEndDream()
+        mMainExecutor.runAllReady()
+
+        // Verify view removed.
+        Mockito.verify(mWindowManager).removeView(mViewCaptor!!.value)
+
+        // Verify state correctly set.
+        Mockito.verify(mStateController).setOverlayActive(false)
+        Mockito.verify(mStateController).setLowLightActive(false)
+        Mockito.verify(mStateController).setEntryAnimationsFinished(false)
+    }
+
+    @Test
+    fun testImmediateEndDream() {
+        val client = client
+
+        // Start the dream, but don't execute any Runnables put on the executor yet. We delay
+        // executing Runnables as the timing isn't guaranteed and we want to verify that the overlay
+        // starts and finishes in the proper order even if Runnables are delayed.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        // Immediately end the dream.
+        client.endDream()
+        // Run any scheduled Runnables.
+        mMainExecutor.runAllReady()
+
+        // The overlay starts then finishes.
+        val inOrder = Mockito.inOrder(mWindowManager)
+        inOrder.verify(mWindowManager).addView(mViewCaptor!!.capture(), any())
+        inOrder.verify(mWindowManager).removeView(mViewCaptor!!.value)
+    }
+
+    @Test
+    fun testEndDreamDuringStartDream() {
+        val client = client
+
+        // Schedule the endDream call in the middle of the startDream implementation, as any
+        // ordering is possible.
+        Mockito.doAnswer { invocation: InvocationOnMock? ->
+                client.endDream()
+                null
+            }
+            .`when`(mStateController)
+            .setOverlayActive(true)
+
+        // Start the dream.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        // The overlay starts then finishes.
+        val inOrder = Mockito.inOrder(mWindowManager)
+        inOrder.verify(mWindowManager).addView(mViewCaptor!!.capture(), any())
+        inOrder.verify(mWindowManager).removeView(mViewCaptor!!.value)
+    }
+
+    @Test
+    fun testDestroy() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            LOW_LIGHT_COMPONENT.flattenToString(),
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        // Verify view added.
+        Mockito.verify(mWindowManager).addView(mViewCaptor!!.capture(), any())
+
+        // Service destroyed.
+        mService!!.onDestroy()
+        mMainExecutor.runAllReady()
+
+        // Verify view removed.
+        Mockito.verify(mWindowManager).removeView(mViewCaptor!!.value)
+
+        // Verify state correctly set.
+        Mockito.verify(mKeyguardUpdateMonitor).removeCallback(any())
+        Mockito.verify(mLifecycleRegistry).currentState = Lifecycle.State.DESTROYED
+        Mockito.verify(mStateController).setOverlayActive(false)
+        Mockito.verify(mStateController).setLowLightActive(false)
+        Mockito.verify(mStateController).setEntryAnimationsFinished(false)
+    }
+
+    @Test
+    fun testDoNotRemoveViewOnDestroyIfOverlayNotStarted() {
+        // Service destroyed without ever starting dream.
+        mService!!.onDestroy()
+        mMainExecutor.runAllReady()
+
+        // Verify no view is removed.
+        Mockito.verify(mWindowManager, Mockito.never()).removeView(any())
+
+        // Verify state still correctly set.
+        Mockito.verify(mKeyguardUpdateMonitor).removeCallback(any())
+        Mockito.verify(mLifecycleRegistry).currentState = Lifecycle.State.DESTROYED
+        Mockito.verify(mStateController).setOverlayActive(false)
+        Mockito.verify(mStateController).setLowLightActive(false)
+    }
+
+    @Test
+    fun testDecorViewNotAddedToWindowAfterDestroy() {
+        val client = client
+
+        // Destroy the service.
+        mService!!.onDestroy()
+        mMainExecutor.runAllReady()
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        Mockito.verify(mWindowManager, Mockito.never()).addView(any(), any())
+    }
+
+    @Test
+    fun testNeverRemoveDecorViewIfNotAdded() {
+        // Service destroyed before dream started.
+        mService!!.onDestroy()
+        mMainExecutor.runAllReady()
+        Mockito.verify(mWindowManager, Mockito.never()).removeView(any())
+    }
+
+    @Test
+    fun testResetCurrentOverlayWhenConnectedToNewDream() {
+        val client = client
+
+        // Inform the overlay service of dream starting. Do not show dream complications.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        // Verify that a new window is added.
+        Mockito.verify(mWindowManager).addView(mViewCaptor!!.capture(), any())
+        val windowDecorView = mViewCaptor!!.value
+
+        // Assert that the overlay is not showing complications.
+        Truth.assertThat(mService!!.shouldShowComplications()).isFalse()
+        Mockito.clearInvocations(mDreamOverlayComponent)
+        Mockito.clearInvocations(mAmbientTouchComponent)
+        Mockito.clearInvocations(mWindowManager)
+
+        // New dream starting with dream complications showing. Note that when a new dream is
+        // binding to the dream overlay service, it receives the same instance of IBinder as the
+        // first one.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            true /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        // Assert that the overlay is showing complications.
+        Truth.assertThat(mService!!.shouldShowComplications()).isTrue()
+
+        // Verify that the old overlay window has been removed, and a new one created.
+        Mockito.verify(mWindowManager).removeView(windowDecorView)
+        Mockito.verify(mWindowManager).addView(any(), any())
+
+        // Verify that new instances of overlay container view controller and overlay touch monitor
+        // are created.
+        Mockito.verify(mDreamOverlayComponent).getDreamOverlayContainerViewController()
+        Mockito.verify(mAmbientTouchComponent).getTouchMonitor()
+    }
+
+    @Test
+    fun testWakeUp() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            true /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        mService!!.onWakeUp()
+        Mockito.verify(mDreamOverlayContainerViewController).wakeUp()
+        Mockito.verify(mDreamOverlayCallbackController).onWakeUp()
+    }
+
+    @Test
+    fun testWakeUpBeforeStartDoesNothing() {
+        mService!!.onWakeUp()
+        Mockito.verify(mDreamOverlayContainerViewController, Mockito.never()).wakeUp()
+    }
+
+    @Test
+    fun testSystemFlagShowForAllUsersSetOnWindow() {
+        val client = client
+
+        // Inform the overlay service of dream starting. Do not show dream complications.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+        val paramsCaptor = ArgumentCaptor.forClass(WindowManager.LayoutParams::class.java)
+
+        // Verify that a new window is added.
+        Mockito.verify(mWindowManager).addView(any(), paramsCaptor.capture())
+        Truth.assertThat(
+                paramsCaptor.value.privateFlags and
+                    WindowManager.LayoutParams.SYSTEM_FLAG_SHOW_FOR_ALL_USERS ==
+                    WindowManager.LayoutParams.SYSTEM_FLAG_SHOW_FOR_ALL_USERS
+            )
+            .isTrue()
+    }
+
+    companion object {
+        private val LOW_LIGHT_COMPONENT = ComponentName("package", "lowlight")
+        private val HOME_CONTROL_PANEL_DREAM_COMPONENT =
+            ComponentName("package", "homeControlPanel")
+        private const val DREAM_COMPONENT = "package/dream"
+        private const val WINDOW_NAME = "test"
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt
index 2727af6..2397de6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModelTest.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
+import com.android.systemui.common.ui.data.repository.fakeConfigurationRepository
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
@@ -49,7 +50,9 @@
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
+import java.util.Locale
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -261,22 +264,55 @@
     @Test
     fun unfoldTransitionProgress() =
         testScope.runTest {
-            val unfoldProvider = kosmos.fakeUnfoldTransitionProgressProvider
-            val progress by collectLastValue(underTest.unfoldTransitionProgress)
+            val maxTranslation = prepareConfiguration()
+            val translations by
+                collectLastValue(
+                    combine(
+                        underTest.unfoldTranslationX(isOnStartSide = true),
+                        underTest.unfoldTranslationX(isOnStartSide = false),
+                    ) { start, end ->
+                        Translations(
+                            start = start,
+                            end = end,
+                        )
+                    }
+                )
 
+            val unfoldProvider = kosmos.fakeUnfoldTransitionProgressProvider
             unfoldProvider.onTransitionStarted()
-            assertThat(progress).isEqualTo(1f)
+            assertThat(translations?.start).isEqualTo(0f)
+            assertThat(translations?.end).isEqualTo(-0f)
 
             repeat(10) { repetition ->
                 val transitionProgress = 0.1f * (repetition + 1)
                 unfoldProvider.onTransitionProgress(transitionProgress)
-                assertThat(progress).isEqualTo(transitionProgress)
+                assertThat(translations?.start).isEqualTo((1 - transitionProgress) * maxTranslation)
+                assertThat(translations?.end).isEqualTo(-(1 - transitionProgress) * maxTranslation)
             }
 
             unfoldProvider.onTransitionFinishing()
-            assertThat(progress).isEqualTo(1f)
+            assertThat(translations?.start).isEqualTo(0f)
+            assertThat(translations?.end).isEqualTo(-0f)
 
             unfoldProvider.onTransitionFinished()
-            assertThat(progress).isEqualTo(1f)
+            assertThat(translations?.start).isEqualTo(0f)
+            assertThat(translations?.end).isEqualTo(-0f)
         }
+
+    private fun prepareConfiguration(): Int {
+        val configuration = context.resources.configuration
+        configuration.setLayoutDirection(Locale.US)
+        kosmos.fakeConfigurationRepository.onConfigurationChange(configuration)
+        val maxTranslation = 10
+        kosmos.fakeConfigurationRepository.setDimensionPixelSize(
+            R.dimen.notification_side_paddings,
+            maxTranslation
+        )
+        return maxTranslation
+    }
+
+    private data class Translations(
+        val start: Float,
+        val end: Float,
+    )
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorTest.kt
index 3b4cce4..12f08a3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorTest.kt
@@ -18,13 +18,17 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.common.ui.data.repository.fakeConfigurationRepository
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.res.R
 import com.android.systemui.testKosmos
 import com.android.systemui.unfold.fakeUnfoldTransitionProgressProvider
 import com.google.common.truth.Truth.assertThat
+import java.util.Locale
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.async
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -79,24 +83,102 @@
         }
 
     @Test
-    fun unfoldProgress() =
+    fun unfoldTranslationX_leftToRight() =
         testScope.runTest {
-            val progress by collectLastValue(underTest.unfoldProgress)
+            val maxTranslation = prepareConfiguration(isLeftToRight = true)
+            val translations by
+                collectLastValue(
+                    combine(
+                        underTest.unfoldTranslationX(isOnStartSide = true),
+                        underTest.unfoldTranslationX(isOnStartSide = false),
+                    ) { start, end ->
+                        Translations(
+                            start = start,
+                            end = end,
+                        )
+                    }
+                )
             runCurrent()
 
             unfoldTransitionProgressProvider.onTransitionStarted()
-            assertThat(progress).isEqualTo(1f)
+            assertThat(translations?.start).isEqualTo(0f)
+            assertThat(translations?.end).isEqualTo(-0f)
 
             repeat(10) { repetition ->
-                val transitionProgress = 0.1f * (repetition + 1)
+                val transitionProgress = 1 - 0.1f * (repetition + 1)
                 unfoldTransitionProgressProvider.onTransitionProgress(transitionProgress)
-                assertThat(progress).isEqualTo(transitionProgress)
+                assertThat(translations?.start).isEqualTo((1 - transitionProgress) * maxTranslation)
+                assertThat(translations?.end).isEqualTo(-(1 - transitionProgress) * maxTranslation)
             }
 
             unfoldTransitionProgressProvider.onTransitionFinishing()
-            assertThat(progress).isEqualTo(1f)
+            assertThat(translations?.start).isEqualTo(maxTranslation)
+            assertThat(translations?.end).isEqualTo(-maxTranslation)
 
             unfoldTransitionProgressProvider.onTransitionFinished()
-            assertThat(progress).isEqualTo(1f)
+            assertThat(translations?.start).isEqualTo(0f)
+            assertThat(translations?.end).isEqualTo(-0f)
         }
+
+    @Test
+    fun unfoldTranslationX_rightToLeft() =
+        testScope.runTest {
+            val maxTranslation = prepareConfiguration(isLeftToRight = false)
+            val translations by
+                collectLastValue(
+                    combine(
+                        underTest.unfoldTranslationX(isOnStartSide = true),
+                        underTest.unfoldTranslationX(isOnStartSide = false),
+                    ) { start, end ->
+                        Translations(
+                            start = start,
+                            end = end,
+                        )
+                    }
+                )
+            runCurrent()
+
+            unfoldTransitionProgressProvider.onTransitionStarted()
+            assertThat(translations?.start).isEqualTo(-0f)
+            assertThat(translations?.end).isEqualTo(0f)
+
+            repeat(10) { repetition ->
+                val transitionProgress = 1 - 0.1f * (repetition + 1)
+                unfoldTransitionProgressProvider.onTransitionProgress(transitionProgress)
+                assertThat(translations?.start)
+                    .isEqualTo(-(1 - transitionProgress) * maxTranslation)
+                assertThat(translations?.end).isEqualTo((1 - transitionProgress) * maxTranslation)
+            }
+
+            unfoldTransitionProgressProvider.onTransitionFinishing()
+            assertThat(translations?.start).isEqualTo(-maxTranslation)
+            assertThat(translations?.end).isEqualTo(maxTranslation)
+
+            unfoldTransitionProgressProvider.onTransitionFinished()
+            assertThat(translations?.start).isEqualTo(-0f)
+            assertThat(translations?.end).isEqualTo(0f)
+        }
+
+    private fun prepareConfiguration(
+        isLeftToRight: Boolean,
+    ): Int {
+        val configuration = context.resources.configuration
+        if (isLeftToRight) {
+            configuration.setLayoutDirection(Locale.US)
+        } else {
+            configuration.setLayoutDirection(Locale("he", "il"))
+        }
+        kosmos.fakeConfigurationRepository.onConfigurationChange(configuration)
+        val maxTranslation = 10
+        kosmos.fakeConfigurationRepository.setDimensionPixelSize(
+            R.dimen.notification_side_paddings,
+            maxTranslation
+        )
+        return maxTranslation
+    }
+
+    private data class Translations(
+        val start: Float,
+        val end: Float,
+    )
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt
index fdeded8..4cf924a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModelTest.kt
@@ -64,7 +64,7 @@
                 val buttonViewModel by collectLastValue(underTest.buttonViewModel)
                 runCurrent()
 
-                assertThat(buttonViewModel!!.isChecked).isFalse()
+                assertThat(buttonViewModel!!.isActive).isFalse()
             }
         }
     }
@@ -78,7 +78,7 @@
                 val buttonViewModel by collectLastValue(underTest.buttonViewModel)
                 runCurrent()
 
-                assertThat(buttonViewModel!!.isChecked).isTrue()
+                assertThat(buttonViewModel!!.isActive).isTrue()
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java
index 458a21c..75d925d 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java
@@ -107,7 +107,10 @@
         }
     }
 
-    private void updateMessageAreaVisibility() {
+    /**
+     * Determines whether to show the message area controlled by MessageAreaController.
+     */
+    public void updateMessageAreaVisibility() {
         if (mMessageAreaController == null) return;
         if (Flags.revampedBouncerMessages()) {
             mMessageAreaController.disable();
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java
index 558679e..3ef3418 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java
@@ -118,6 +118,12 @@
     }
 
     @Override
+    public void updateMessageAreaVisibility() {
+        if (mMessageAreaController == null) return;
+        mMessageAreaController.setIsVisible(true);
+    }
+
+    @Override
     void resetState() {
         super.resetState();
         if (DEBUG) Log.v(TAG, "Resetting state");
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java
index cb1c4b3..46225c7 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java
@@ -115,6 +115,12 @@
     }
 
     @Override
+    public void updateMessageAreaVisibility() {
+        if (mMessageAreaController == null) return;
+        mMessageAreaController.setIsVisible(true);
+    }
+
+    @Override
     public void onResume(int reason) {
         super.onResume(reason);
         if (mShowDefaultMessage) {
diff --git a/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt b/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt
index 5f6ff82..638af58 100644
--- a/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/common/ui/domain/interactor/ConfigurationInteractor.kt
@@ -56,6 +56,13 @@
     val naturalMaxBounds: Flow<Rect> =
         repository.configurationValues.map { it.naturalScreenBounds }.distinctUntilChanged()
 
+    /**
+     * The layout direction. Will be either `View#LAYOUT_DIRECTION_LTR` or
+     * `View#LAYOUT_DIRECTION_RTL`.
+     */
+    val layoutDirection: Flow<Int> =
+        repository.configurationValues.map { it.layoutDirection }.distinctUntilChanged()
+
     /** Given [resourceId], emit the dimension pixel size on config change */
     fun dimensionPixelSize(resourceId: Int): Flow<Int> {
         return onAnyConfigurationChange.mapLatest { repository.getDimensionPixelSize(resourceId) }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt
index 06a0c72..5a559fc 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt
@@ -18,61 +18,29 @@
 package com.android.systemui.keyguard.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.keyguard.domain.interactor.FromAlternateBouncerTransitionInteractor.Companion.TRANSITION_DURATION_MS
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
-import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager
-import com.android.wm.shell.animation.Interpolators
 import javax.inject.Inject
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.merge
 
 @ExperimentalCoroutinesApi
 class AlternateBouncerViewModel
 @Inject
 constructor(
     private val statusBarKeyguardViewManager: StatusBarKeyguardViewManager,
-    animationFlow: KeyguardTransitionAnimationFlow,
+    keyguardTransitionInteractor: KeyguardTransitionInteractor,
 ) {
     // When we're fully transitioned to the AlternateBouncer, the alpha of the scrim should be:
     private val alternateBouncerScrimAlpha = .66f
-    private val toAlternateBouncerTransition =
-        animationFlow
-            .setup(
-                duration = TRANSITION_DURATION_MS,
-                from = null,
-                to = ALTERNATE_BOUNCER,
-            )
-            .sharedFlow(
-                duration = TRANSITION_DURATION_MS,
-                onStep = { it },
-                onFinish = { 1f },
-                // Reset on cancel
-                onCancel = { 0f },
-                interpolator = Interpolators.FAST_OUT_SLOW_IN,
-            )
-    private val fromAlternateBouncerTransition =
-        animationFlow
-            .setup(
-                TRANSITION_DURATION_MS,
-                from = ALTERNATE_BOUNCER,
-                to = null,
-            )
-            .sharedFlow(
-                duration = TRANSITION_DURATION_MS,
-                onStep = { 1f - it },
-                // Reset on cancel
-                onCancel = { 0f },
-                interpolator = Interpolators.FAST_OUT_SLOW_IN,
-            )
 
     /** Progress to a fully transitioned alternate bouncer. 1f represents fully transitioned. */
     val transitionToAlternateBouncerProgress =
-        merge(fromAlternateBouncerTransition, toAlternateBouncerTransition)
+        keyguardTransitionInteractor.transitionValue(ALTERNATE_BOUNCER)
 
     val forcePluginOpen: Flow<Boolean> =
         transitionToAlternateBouncerProgress.map { it > 0f }.distinctUntilChanged()
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt
index c7cfb0b..0e2814b 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt
@@ -35,6 +35,7 @@
 import com.android.settingslib.media.LocalMediaManager
 import com.android.settingslib.media.MediaDevice
 import com.android.settingslib.media.PhoneMediaDevice
+import com.android.settingslib.media.flags.Flags
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.media.controls.shared.model.MediaData
@@ -178,7 +179,9 @@
             bgExecutor.execute {
                 if (!started) {
                     localMediaManager.registerCallback(this)
-                    localMediaManager.startScan()
+                    if (!Flags.removeUnnecessaryRouteScanning()) {
+                        localMediaManager.startScan()
+                    }
                     muteAwaitConnectionManager.startListening()
                     playbackType = controller?.playbackInfo?.playbackType ?: PLAYBACK_TYPE_UNKNOWN
                     playbackVolumeControlId = controller?.playbackInfo?.volumeControlId
@@ -195,7 +198,9 @@
                 if (started) {
                     started = false
                     controller?.unregisterCallback(this)
-                    localMediaManager.stopScan()
+                    if (!Flags.removeUnnecessaryRouteScanning()) {
+                        localMediaManager.stopScan()
+                    }
                     localMediaManager.unregisterCallback(this)
                     muteAwaitConnectionManager.stopListening()
                     configurationController.removeCallback(configListener)
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowViewController.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowViewController.java
index 00bc752..a763641 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowViewController.java
@@ -659,6 +659,9 @@
             mTouchCancelled = true;
         }
         mAmbientState.setSwipingUp(false);
+        if (MigrateClocksToBlueprint.isEnabled()) {
+            mDragDownHelper.stopDragging();
+        }
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt
index 6800c61..5b76acb 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneViewModel.kt
@@ -65,7 +65,7 @@
     private val footerActionsViewModelFactory: FooterActionsViewModel.Factory,
     private val footerActionsController: FooterActionsController,
     private val sceneInteractor: SceneInteractor,
-    unfoldTransitionInteractor: UnfoldTransitionInteractor,
+    private val unfoldTransitionInteractor: UnfoldTransitionInteractor,
 ) {
     val destinationScenes: StateFlow<Map<UserAction, UserActionResult>> =
         combine(
@@ -109,14 +109,12 @@
     val shadeMode: StateFlow<ShadeMode> = shadeInteractor.shadeMode
 
     /**
-     * The unfold transition progress. When fully-unfolded, this is `1` and fully folded, it's `0`.
+     * Amount of X-axis translation to apply to various elements as the unfolded foldable is folded
+     * slightly, in pixels.
      */
-    val unfoldTransitionProgress: StateFlow<Float> =
-        unfoldTransitionInteractor.unfoldProgress.stateIn(
-            scope = applicationScope,
-            started = SharingStarted.WhileSubscribed(),
-            initialValue = 1f
-        )
+    fun unfoldTranslationX(isOnStartSide: Boolean): Flow<Float> {
+        return unfoldTransitionInteractor.unfoldTranslationX(isOnStartSide)
+    }
 
     /** Notifies that some content in the shade was clicked. */
     fun onContentClicked() {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt
index 519d719..e3db626 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/LockscreenShadeTransitionController.kt
@@ -959,7 +959,7 @@
         anim.start()
     }
 
-    private fun stopDragging() {
+    fun stopDragging() {
         if (startingChild != null) {
             cancelChildExpansion(startingChild!!)
             startingChild = null
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 9e0b16c..0c8518f 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
@@ -1767,7 +1767,8 @@
      */
     public ExpandableNotificationRow(Context context, AttributeSet attrs) {
         this(context, attrs, context);
-        Log.e(TAG, "This constructor shouldn't be called");
+        // NOTE(b/317503801): Always crash when using the insecure constructor.
+        throw new UnsupportedOperationException("Insecure constructor");
     }
 
     /**
@@ -2804,12 +2805,7 @@
     }
 
     public boolean isExpanded(boolean allowOnKeyguard) {
-        if (DEBUG) {
-            if (!mShowingPublicInitialized && !allowOnKeyguard) {
-                Log.d(TAG, "mShowingPublic is not initialized.");
-            }
-        }
-        return !mShowingPublic && (!mOnKeyguard || allowOnKeyguard)
+        return (!shouldShowPublic()) && (!mOnKeyguard || allowOnKeyguard)
                 && (!hasUserChangedExpansion() && (isSystemExpanded() || isSystemChildExpanded())
                 || isUserExpanded());
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
index 5fbcebd..35afda7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
@@ -33,6 +33,8 @@
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.util.time.SystemClock;
 
+import java.util.concurrent.Executor;
+
 import javax.inject.Inject;
 
 /**
@@ -58,10 +60,22 @@
     }
 
     /**
-     * Inflates a new notificationView. This should not be called twice on this object
+     * Inflates a new notificationView asynchronously, calling the {@code listener} on the main
+     * thread when done. This should not be called twice on this object.
      */
     public void inflate(Context context, ViewGroup parent, NotificationEntry entry,
             RowInflationFinishedListener listener) {
+        inflate(context, parent, entry, null, listener);
+    }
+
+    /**
+     * Inflates a new notificationView asynchronously, calling the {@code listener} on the supplied
+     * {@code listenerExecutor} (or the main thread if null) when done. This should not be called
+     * twice on this object.
+     */
+    @VisibleForTesting
+    public void inflate(Context context, ViewGroup parent, NotificationEntry entry,
+            @Nullable Executor listenerExecutor, RowInflationFinishedListener listener) {
         if (TRACE_ORIGIN) {
             mInflateOrigin = new Throwable("inflate requested here");
         }
@@ -72,7 +86,7 @@
 
         mLogger.logInflateStart(entry);
         mInflateStartTimeMs = mSystemClock.elapsedRealtime();
-        inflater.inflate(R.layout.status_bar_notification_row, parent, this);
+        inflater.inflate(R.layout.status_bar_notification_row, parent, listenerExecutor, this);
     }
 
     private RowAsyncLayoutInflater makeRowInflater(NotificationEntry entry) {
@@ -80,12 +94,12 @@
     }
 
     @VisibleForTesting
-    static class RowAsyncLayoutInflater implements AsyncLayoutFactory {
+    public static class RowAsyncLayoutInflater implements AsyncLayoutFactory {
         private final NotificationEntry mEntry;
         private final SystemClock mSystemClock;
         private final RowInflaterTaskLogger mLogger;
 
-        RowAsyncLayoutInflater(NotificationEntry entry, SystemClock systemClock,
+        public RowAsyncLayoutInflater(NotificationEntry entry, SystemClock systemClock,
                 RowInflaterTaskLogger logger) {
             mEntry = entry;
             mSystemClock = systemClock;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
index 5099682..37bbbd0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
@@ -66,6 +66,7 @@
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.statusbar.notification.stack.domain.interactor.NotificationStackAppearanceInteractor
 import com.android.systemui.statusbar.notification.stack.domain.interactor.SharedNotificationContainerInteractor
+import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractor
 import com.android.systemui.util.kotlin.BooleanFlowOperators.and
 import com.android.systemui.util.kotlin.BooleanFlowOperators.or
 import com.android.systemui.util.kotlin.FlowDumperImpl
@@ -80,6 +81,7 @@
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.combineTransform
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flow
@@ -128,6 +130,7 @@
     private val primaryBouncerToLockscreenTransitionViewModel:
         PrimaryBouncerToLockscreenTransitionViewModel,
     private val aodBurnInViewModel: AodBurnInViewModel,
+    unfoldTransitionInteractor: UnfoldTransitionInteractor,
 ) : FlowDumperImpl(dumpManager) {
     private val statesForConstrainedNotifications: Set<KeyguardState> =
         setOf(AOD, LOCKSCREEN, DOZING, ALTERNATE_BOUNCER, PRIMARY_BOUNCER)
@@ -577,14 +580,20 @@
             .dumpWhileCollecting("translationY")
     }
 
-    /**
-     * The container may need to be translated in the x direction as the keyguard fades out, such as
-     * when swiping open the glanceable hub from the lockscreen.
-     */
+    /** Horizontal translation to apply to the container. */
     val translationX: Flow<Float> =
         merge(
+                // The container may need to be translated along the X axis as the keyguard fades
+                // out, such as when swiping open the glanceable hub from the lockscreen.
                 lockscreenToGlanceableHubTransitionViewModel.notificationTranslationX,
                 glanceableHubToLockscreenTransitionViewModel.notificationTranslationX,
+                if (SceneContainerFlag.isEnabled) {
+                    // The container may need to be translated along the X axis as the unfolded
+                    // foldable is folded slightly.
+                    unfoldTransitionInteractor.unfoldTranslationX(isOnStartSide = false)
+                } else {
+                    emptyFlow()
+                }
             )
             .dumpWhileCollecting("translationX")
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
index ebaeb39..68d54e7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/LegacyActivityStarterInternalImpl.kt
@@ -131,7 +131,7 @@
 
         val runnable = Runnable {
             assistManagerLazy.get().hideAssist()
-            intent.flags = Intent.FLAG_ACTIVITY_NEW_TASK
+            intent.flags = Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_CLEAR_TOP
             intent.addFlags(flags)
             val result = intArrayOf(ActivityManager.START_CANCELED)
             activityTransitionAnimator.startIntentWithAnimation(
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
index 1b56702..0c2abd9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
@@ -153,18 +153,20 @@
             // Use default duration, like we did before AvalancheController existed
             return autoDismissMs
         }
+
         val showingList: MutableList<HeadsUpEntry> = mutableListOf()
         headsUpEntryShowing?.let { showingList.add(it) }
 
+        nextList.sort()
         val entryList = showingList + nextList
-        if (entryList.indexOf(entry) == entryList.size - 1) {
-            // Use default duration if last entry
+        val thisEntryIndex = entryList.indexOf(entry)
+        val nextEntryIndex = thisEntryIndex + 1
+
+        // If last entry, use default duration
+        if (nextEntryIndex >= entryList.size) {
             return autoDismissMs
         }
-
-        nextList.sort()
-        val nextEntry = nextList[0]
-
+        val nextEntry = entryList[nextEntryIndex]
         if (nextEntry.compareNonTimeFields(entry) == -1) {
             // Next entry is higher priority
             return 500
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
index 3522850..37ef1f2 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
@@ -32,8 +32,6 @@
 import com.android.systemui.unfold.data.repository.FoldStateRepositoryImpl
 import com.android.systemui.unfold.data.repository.UnfoldTransitionRepository
 import com.android.systemui.unfold.data.repository.UnfoldTransitionRepositoryImpl
-import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractor
-import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractorImpl
 import com.android.systemui.unfold.system.SystemUnfoldSharedModule
 import com.android.systemui.unfold.updates.FoldProvider
 import com.android.systemui.unfold.updates.FoldStateProvider
@@ -186,8 +184,6 @@
     interface Bindings {
         @Binds fun bindRepository(impl: UnfoldTransitionRepositoryImpl): UnfoldTransitionRepository
 
-        @Binds fun bindInteractor(impl: UnfoldTransitionInteractorImpl): UnfoldTransitionInteractor
-
         @Binds fun bindFoldStateRepository(impl: FoldStateRepositoryImpl): FoldStateRepository
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractor.kt
index a8e4496..03499cb 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractor.kt
@@ -15,51 +15,79 @@
  */
 package com.android.systemui.unfold.domain.interactor
 
+import android.view.View
+import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.res.R
 import com.android.systemui.unfold.data.repository.UnfoldTransitionRepository
 import com.android.systemui.unfold.data.repository.UnfoldTransitionStatus.TransitionFinished
 import com.android.systemui.unfold.data.repository.UnfoldTransitionStatus.TransitionInProgress
 import com.android.systemui.unfold.data.repository.UnfoldTransitionStatus.TransitionStarted
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
 
 /**
  * Contains business-logic related to fold-unfold transitions while interacting with
  * [UnfoldTransitionRepository]
  */
-interface UnfoldTransitionInteractor {
+@SysUISingleton
+class UnfoldTransitionInteractor
+@Inject
+constructor(
+    private val repository: UnfoldTransitionRepository,
+    private val configurationInteractor: ConfigurationInteractor,
+) {
     /** Returns availability of fold/unfold transitions on the device */
     val isAvailable: Boolean
-
-    val unfoldProgress: Flow<Float>
-
-    /** Suspends and waits for a fold/unfold transition to finish */
-    suspend fun waitForTransitionFinish()
-
-    /** Suspends and waits for a fold/unfold transition to start */
-    suspend fun waitForTransitionStart()
-}
-
-class UnfoldTransitionInteractorImpl
-@Inject
-constructor(private val repository: UnfoldTransitionRepository) : UnfoldTransitionInteractor {
-
-    override val isAvailable: Boolean
         get() = repository.isAvailable
 
-    override val unfoldProgress: Flow<Float> =
+    /**
+     * This mapping emits 1 when the device is completely unfolded and 0.0 when the device is
+     * completely folded.
+     */
+    private val unfoldProgress: Flow<Float> =
         repository.transitionStatus
             .map { (it as? TransitionInProgress)?.progress ?: 1f }
+            .onStart { emit(1f) }
             .distinctUntilChanged()
 
-    override suspend fun waitForTransitionFinish() {
+    /**
+     * Amount of X-axis translation to apply to various elements as the unfolded foldable is folded
+     * slightly, in pixels.
+     *
+     * @param isOnStartSide Whether the consumer wishes to get a translation amount that's suitable
+     *   for an element that's on the start-side (left hand-side in left-to-right layouts); if
+     *   `true`, the values will provide positive translations to push the left-hand-side element
+     *   towards the foldable hinge; if `false`, the values will be inverted to provide negative
+     *   translations to push the right-hand-side element towards the foldable hinge. Note that this
+     *   method already accounts for left-to-right vs. right-to-left layout directions.
+     */
+    fun unfoldTranslationX(isOnStartSide: Boolean): Flow<Float> {
+        return combine(
+            unfoldProgress,
+            configurationInteractor.dimensionPixelSize(R.dimen.notification_side_paddings),
+            configurationInteractor.layoutDirection.map {
+                if (it == View.LAYOUT_DIRECTION_RTL) -1 else 1
+            },
+        ) { unfoldedAmount, max, layoutDirectionMultiplier ->
+            val sideMultiplier = if (isOnStartSide) 1 else -1
+            max * (1 - unfoldedAmount) * sideMultiplier * layoutDirectionMultiplier
+        }
+    }
+
+    /** Suspends and waits for a fold/unfold transition to finish */
+    suspend fun waitForTransitionFinish() {
         repository.transitionStatus.filter { it is TransitionFinished }.first()
     }
 
-    override suspend fun waitForTransitionStart() {
+    /** Suspends and waits for a fold/unfold transition to start */
+    suspend fun waitForTransitionStart() {
         repository.transitionStatus.filter { it is TransitionStarted }.first()
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/button/ui/viewmodel/ButtonViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/button/ui/viewmodel/ButtonViewModel.kt
index 754d258..4d11f44 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/button/ui/viewmodel/ButtonViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/button/ui/viewmodel/ButtonViewModel.kt
@@ -1,17 +1,17 @@
 /*
  * 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
+ * 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
+ *      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.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 package com.android.systemui.volume.panel.component.button.ui.viewmodel
@@ -22,4 +22,5 @@
 data class ButtonViewModel(
     val icon: Icon,
     val label: CharSequence,
+    val isActive: Boolean = true,
 )
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/button/ui/viewmodel/ToggleButtonViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/button/ui/viewmodel/ToggleButtonViewModel.kt
deleted file mode 100644
index 6c47aec..0000000
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/button/ui/viewmodel/ToggleButtonViewModel.kt
+++ /dev/null
@@ -1,28 +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.volume.panel.component.button.ui.viewmodel
-
-import com.android.systemui.common.shared.model.Icon
-
-data class ToggleButtonViewModel(
-    val isChecked: Boolean,
-    val icon: Icon,
-    val label: CharSequence,
-)
-
-fun ToggleButtonViewModel.toButtonViewModel(): ButtonViewModel =
-    ButtonViewModel(icon = icon, label = label)
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt
index 01421f8..ca5aef8 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/captioning/ui/viewmodel/CaptioningViewModel.kt
@@ -21,7 +21,7 @@
 import com.android.settingslib.view.accessibility.domain.interactor.CaptioningInteractor
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.res.R
-import com.android.systemui.volume.panel.component.button.ui.viewmodel.ToggleButtonViewModel
+import com.android.systemui.volume.panel.component.button.ui.viewmodel.ButtonViewModel
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
 import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
 import javax.inject.Inject
@@ -43,11 +43,11 @@
     private val uiEventLogger: UiEventLogger,
 ) {
 
-    val buttonViewModel: StateFlow<ToggleButtonViewModel?> =
+    val buttonViewModel: StateFlow<ButtonViewModel?> =
         captioningInteractor.isSystemAudioCaptioningEnabled
             .map { isEnabled ->
-                ToggleButtonViewModel(
-                    isChecked = isEnabled,
+                ButtonViewModel(
+                    isActive = isEnabled,
                     icon =
                         Icon.Resource(
                             if (isEnabled) R.drawable.ic_volume_odi_captions
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioButtonViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioButtonViewModel.kt
index e5c5a65..f7a602e 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioButtonViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioButtonViewModel.kt
@@ -16,10 +16,10 @@
 
 package com.android.systemui.volume.panel.component.spatial.ui.viewmodel
 
-import com.android.systemui.volume.panel.component.button.ui.viewmodel.ToggleButtonViewModel
+import com.android.systemui.volume.panel.component.button.ui.viewmodel.ButtonViewModel
 import com.android.systemui.volume.panel.component.spatial.domain.model.SpatialAudioEnabledModel
 
 data class SpatialAudioButtonViewModel(
     val model: SpatialAudioEnabledModel,
-    val button: ToggleButtonViewModel,
+    val button: ButtonViewModel,
 )
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
index b5e9ed2..4b2d26a 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/ui/viewmodel/SpatialAudioViewModel.kt
@@ -22,8 +22,6 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.res.R
 import com.android.systemui.volume.panel.component.button.ui.viewmodel.ButtonViewModel
-import com.android.systemui.volume.panel.component.button.ui.viewmodel.ToggleButtonViewModel
-import com.android.systemui.volume.panel.component.button.ui.viewmodel.toButtonViewModel
 import com.android.systemui.volume.panel.component.spatial.domain.SpatialAudioAvailabilityCriteria
 import com.android.systemui.volume.panel.component.spatial.domain.interactor.SpatialAudioComponentInteractor
 import com.android.systemui.volume.panel.component.spatial.domain.model.SpatialAudioAvailabilityModel
@@ -53,8 +51,8 @@
     val spatialAudioButton: StateFlow<ButtonViewModel?> =
         interactor.isEnabled
             .map {
-                it.toViewModel(true)
-                    .toButtonViewModel()
+                val isChecked = it is SpatialAudioEnabledModel.SpatialAudioEnabled
+                it.toViewModel(isChecked)
                     .copy(label = context.getString(R.string.volume_panel_spatial_audio_title))
             }
             .stateIn(scope, SharingStarted.Eagerly, null)
@@ -76,8 +74,7 @@
                     }
                     .map { isEnabled ->
                         val isChecked = isEnabled == currentIsEnabled
-                        val buttonViewModel: ToggleButtonViewModel =
-                            isEnabled.toViewModel(isChecked)
+                        val buttonViewModel: ButtonViewModel = isEnabled.toViewModel(isChecked)
                         SpatialAudioButtonViewModel(button = buttonViewModel, model = isEnabled)
                     }
             }
@@ -100,26 +97,26 @@
         scope.launch { interactor.setEnabled(model) }
     }
 
-    private fun SpatialAudioEnabledModel.toViewModel(isChecked: Boolean): ToggleButtonViewModel {
+    private fun SpatialAudioEnabledModel.toViewModel(isChecked: Boolean): ButtonViewModel {
         if (this is SpatialAudioEnabledModel.HeadTrackingEnabled) {
-            return ToggleButtonViewModel(
-                isChecked = isChecked,
+            return ButtonViewModel(
+                isActive = isChecked,
                 icon = Icon.Resource(R.drawable.ic_head_tracking, contentDescription = null),
                 label = context.getString(R.string.volume_panel_spatial_audio_tracking)
             )
         }
 
         if (this is SpatialAudioEnabledModel.SpatialAudioEnabled) {
-            return ToggleButtonViewModel(
-                isChecked = isChecked,
+            return ButtonViewModel(
+                isActive = isChecked,
                 icon = Icon.Resource(R.drawable.ic_spatial_audio, contentDescription = null),
                 label = context.getString(R.string.volume_panel_spatial_audio_fixed)
             )
         }
 
         if (this is SpatialAudioEnabledModel.Disabled) {
-            return ToggleButtonViewModel(
-                isChecked = isChecked,
+            return ButtonViewModel(
+                isActive = isChecked,
                 icon = Icon.Resource(R.drawable.ic_spatial_audio_off, contentDescription = null),
                 label = context.getString(R.string.volume_panel_spatial_audio_off)
             )
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt
index 9b5364e..7151c42 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt
@@ -115,6 +115,7 @@
     @Test
     fun onViewAttached() {
         underTest.onViewAttached()
+        verify(keyguardMessageAreaController).setIsVisible(true)
         verify(keyguardMessageAreaController)
             .setMessage(context.resources.getString(R.string.keyguard_enter_your_pin), false)
         verify(keyguardUpdateMonitor)
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt
index e71490c..acae913 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt
@@ -103,7 +103,9 @@
 
     @Test
     fun onViewAttached() {
+        Mockito.reset(keyguardMessageAreaController)
         underTest.onViewAttached()
+        Mockito.verify(keyguardMessageAreaController).setIsVisible(true)
         Mockito.verify(keyguardUpdateMonitor)
             .registerCallback(any(KeyguardUpdateMonitorCallback::class.java))
         Mockito.verify(keyguardMessageAreaController)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt
index d410dac..f1c93c4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt
@@ -32,7 +32,6 @@
 import com.google.common.collect.Range
 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
@@ -67,7 +66,7 @@
     fun transitionToAlternateBouncer_scrimAlphaUpdate() =
         testScope.runTest {
             val scrimAlphas by collectValues(underTest.scrimAlpha)
-            runCurrent()
+            assertThat(scrimAlphas.size).isEqualTo(1) // initial value is 0f
 
             transitionRepository.sendTransitionSteps(
                 listOf(
@@ -79,7 +78,7 @@
                 testScope,
             )
 
-            assertThat(scrimAlphas.size).isEqualTo(4)
+            assertThat(scrimAlphas.size).isEqualTo(5)
             scrimAlphas.forEach { assertThat(it).isIn(Range.closed(0f, 1f)) }
         }
 
@@ -87,7 +86,7 @@
     fun transitionFromAlternateBouncer_scrimAlphaUpdate() =
         testScope.runTest {
             val scrimAlphas by collectValues(underTest.scrimAlpha)
-            runCurrent()
+            assertThat(scrimAlphas.size).isEqualTo(1) // initial value is 0f
 
             transitionRepository.sendTransitionSteps(
                 listOf(
@@ -98,7 +97,7 @@
                 ),
                 testScope,
             )
-            assertThat(scrimAlphas.size).isEqualTo(4)
+            assertThat(scrimAlphas.size).isEqualTo(5)
             scrimAlphas.forEach { assertThat(it).isIn(Range.closed(0f, 1f)) }
         }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt
index b04503b..8c5a4d0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt
@@ -578,6 +578,14 @@
         assertEquals(keyEvent, falsingCollector.lastKeyEvent)
     }
 
+    @Test
+    fun cancelCurrentTouch_callsDragDownHelper() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT)
+        underTest.cancelCurrentTouch()
+
+        verify(dragDownHelper).stopDragging()
+    }
+
     companion object {
         private val DOWN_EVENT = MotionEvent.obtain(0L, 0L, MotionEvent.ACTION_DOWN, 0f, 0f, 0)
         private val MOVE_EVENT = MotionEvent.obtain(0L, 0L, MotionEvent.ACTION_MOVE, 0f, 0f, 0)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt
index eb418fd..4679a58 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt
@@ -19,11 +19,12 @@
 import android.content.Context
 import android.content.res.Resources
 import android.content.res.TypedArray
+import android.platform.test.annotations.PlatinumTest
 import android.testing.AndroidTestingRunner
 import android.util.AttributeSet
 import androidx.test.filters.SmallTest
-import com.android.systemui.shared.R
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.shared.R
 import com.android.systemui.shared.shadow.DoubleShadowTextClock
 import com.android.systemui.util.mockito.whenever
 import junit.framework.Assert.assertTrue
@@ -36,6 +37,7 @@
 import org.mockito.junit.MockitoJUnit
 import org.mockito.junit.MockitoRule
 
+@PlatinumTest(focusArea = "sysui")
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
 class DoubleShadowTextClockTest : SysuiTestCase() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java
index eb692eb..0e24ed4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java
@@ -40,6 +40,9 @@
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.RowContentBindParams;
 import com.android.systemui.statusbar.notification.row.RowContentBindStage;
+import com.android.systemui.statusbar.notification.row.RowInflaterTask;
+import com.android.systemui.statusbar.notification.row.RowInflaterTaskLogger;
+import com.android.systemui.util.time.FakeSystemClock;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -114,20 +117,25 @@
 
     private NotificationEntry addGroup(int size) {
         NotificationEntry summary = new NotificationEntryBuilder().build();
-        summary.setRow(createRow());
+        summary.setRow(createRow(summary));
         ArrayList<NotificationEntry> children = new ArrayList<>();
         for (int i = 0; i < size; i++) {
             NotificationEntry child = new NotificationEntryBuilder().build();
-            child.setRow(createRow());
+            child.setRow(createRow(child));
             children.add(child);
         }
         mGroupNotifs.put(summary, children);
         return summary;
     }
 
-    private ExpandableNotificationRow createRow() {
+    private ExpandableNotificationRow createRow(NotificationEntry entry) {
+        LayoutInflater inflater = LayoutInflater.from(mContext);
+        inflater.setFactory2(
+                new RowInflaterTask.RowAsyncLayoutInflater(entry, new FakeSystemClock(), mock(
+                        RowInflaterTaskLogger.class)));
+
         ExpandableNotificationRow row = (ExpandableNotificationRow)
-                LayoutInflater.from(mContext).inflate(R.layout.status_bar_notification_row, null);
+                inflater.inflate(R.layout.status_bar_notification_row, null);
         row.getPrivateLayout().setContractedChild(new View(mContext));
         row.getPrivateLayout().setExpandedChild(new View(mContext));
         return row;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
index 01492f6..aa79c23 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
@@ -772,8 +772,7 @@
         row.setUserExpanded(true);
         row.setOnKeyguard(false);
         row.setSensitive(/* sensitive= */true, /* hideSensitive= */false);
-        row.setHideSensitive(/* hideSensitive= */true, /* animated= */false,
-                /* delay= */0L, /* duration= */0L);
+        row.setHideSensitiveForIntrinsicHeight(/* hideSensitive= */true);
 
         // THEN
         assertThat(row.isExpanded()).isFalse();
@@ -787,8 +786,7 @@
         row.setUserExpanded(true);
         row.setOnKeyguard(false);
         row.setSensitive(/* sensitive= */true, /* hideSensitive= */false);
-        row.setHideSensitive(/* hideSensitive= */false, /* animated= */false,
-                /* delay= */0L, /* duration= */0L);
+        row.setHideSensitiveForIntrinsicHeight(/* hideSensitive= */false);
 
         // THEN
         assertThat(row.isExpanded()).isTrue();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.kt
index 91e4666..7332bc3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.kt
@@ -79,10 +79,11 @@
         initMocks(this)
         fakeParent =
             spy(FrameLayout(mContext, /* attrs= */ null).also { it.visibility = View.GONE })
+        val mockEntry = createMockNotificationEntry()
         row =
             spy(
-                ExpandableNotificationRow(mContext, /* attrs= */ null).apply {
-                    entry = createMockNotificationEntry()
+                ExpandableNotificationRow(mContext, /* attrs= */ null, mockEntry).apply {
+                    entry = mockEntry
                 }
             )
         ViewUtils.attachView(fakeParent)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
index 9a7b8ec..745d20d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
@@ -1,5 +1,6 @@
 package com.android.systemui.statusbar.notification.stack
 
+import android.service.notification.StatusBarNotification
 import android.testing.AndroidTestingRunner
 import android.testing.TestableLooper.RunWithLooper
 import android.view.LayoutInflater
@@ -14,6 +15,7 @@
 import com.android.systemui.shade.transition.LargeScreenShadeInterpolator
 import com.android.systemui.statusbar.NotificationShelf
 import com.android.systemui.statusbar.StatusBarIconView
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.ExpandableView
 import com.android.systemui.statusbar.notification.shared.NotificationIconContainerRefactor
@@ -457,8 +459,13 @@
         expansionFraction: Float,
         expectedAlpha: Float
     ) {
+        val sbnMock: StatusBarNotification = mock()
+        val mockEntry = mock<NotificationEntry>().apply {
+            whenever(this.sbn).thenReturn(sbnMock)
+        }
+        val row = ExpandableNotificationRow(mContext, null, mockEntry)
         whenever(ambientState.lastVisibleBackgroundChild)
-            .thenReturn(ExpandableNotificationRow(mContext, null))
+            .thenReturn(row)
         whenever(ambientState.isExpansionChanging).thenReturn(true)
         whenever(ambientState.expansionFraction).thenReturn(expansionFraction)
         whenever(hostLayoutController.speedBumpIndex).thenReturn(0)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
index df82df8..e2ac203 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
@@ -33,7 +33,7 @@
 import com.android.systemui.statusbar.policy.FakeConfigurationController
 import com.android.systemui.unfold.FakeUnfoldTransitionProvider
 import com.android.systemui.unfold.data.repository.UnfoldTransitionRepositoryImpl
-import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractorImpl
+import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractor
 import com.android.systemui.util.animation.data.repository.FakeAnimationStatusRepository
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
@@ -69,11 +69,6 @@
             statusBarStateController = mock()
         )
 
-    private val unfoldTransitionRepository =
-        UnfoldTransitionRepositoryImpl(Optional.of(unfoldTransitionProgressProvider))
-    private val unfoldTransitionInteractor =
-        UnfoldTransitionInteractorImpl(unfoldTransitionRepository)
-
     private val configurationRepository =
         ConfigurationRepositoryImpl(
             configurationController,
@@ -83,6 +78,11 @@
         )
     private val configurationInteractor = ConfigurationInteractor(configurationRepository)
 
+    private val unfoldTransitionRepository =
+        UnfoldTransitionRepositoryImpl(Optional.of(unfoldTransitionProgressProvider))
+    private val unfoldTransitionInteractor =
+        UnfoldTransitionInteractor(unfoldTransitionRepository, configurationInteractor)
+
     private lateinit var configuration: Configuration
     private lateinit var underTest: HideNotificationsInteractor
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/DisplaySwitchLatencyTrackerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/DisplaySwitchLatencyTrackerTest.kt
index 383f4a3..2cdc8d8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/unfold/DisplaySwitchLatencyTrackerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/DisplaySwitchLatencyTrackerTest.kt
@@ -22,6 +22,8 @@
 import androidx.test.filters.SmallTest
 import com.android.internal.R
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.common.ui.data.repository.ConfigurationRepositoryImpl
+import com.android.systemui.common.ui.domain.interactor.ConfigurationInteractor
 import com.android.systemui.display.data.repository.DeviceStateRepository
 import com.android.systemui.display.data.repository.DeviceStateRepository.DeviceState
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
@@ -31,11 +33,12 @@
 import com.android.systemui.power.shared.model.WakefulnessModel
 import com.android.systemui.power.shared.model.WakefulnessState
 import com.android.systemui.shared.system.SysUiStatsLog
+import com.android.systemui.statusbar.policy.FakeConfigurationController
 import com.android.systemui.unfold.DisplaySwitchLatencyTracker.Companion.FOLDABLE_DEVICE_STATE_CLOSED
 import com.android.systemui.unfold.DisplaySwitchLatencyTracker.Companion.FOLDABLE_DEVICE_STATE_HALF_OPEN
 import com.android.systemui.unfold.DisplaySwitchLatencyTracker.DisplaySwitchLatencyEvent
 import com.android.systemui.unfold.data.repository.UnfoldTransitionRepositoryImpl
-import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractorImpl
+import com.android.systemui.unfold.domain.interactor.UnfoldTransitionInteractor
 import com.android.systemui.util.animation.data.repository.AnimationStatusRepository
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.capture
@@ -86,11 +89,20 @@
     private val areAnimationEnabled = MutableStateFlow(true)
     private val lastWakefulnessEvent = MutableStateFlow(WakefulnessModel())
     private val systemClock = FakeSystemClock()
+    private val configurationController = FakeConfigurationController()
+    private val configurationRepository =
+        ConfigurationRepositoryImpl(
+            configurationController,
+            context,
+            testScope.backgroundScope,
+            mock()
+        )
+    private val configurationInteractor = ConfigurationInteractor(configurationRepository)
     private val unfoldTransitionProgressProvider = FakeUnfoldTransitionProvider()
     private val unfoldTransitionRepository =
         UnfoldTransitionRepositoryImpl(Optional.of(unfoldTransitionProgressProvider))
     private val unfoldTransitionInteractor =
-        UnfoldTransitionInteractorImpl(unfoldTransitionRepository)
+        UnfoldTransitionInteractor(unfoldTransitionRepository, configurationInteractor)
 
     @Before
     fun setup() {
@@ -155,7 +167,10 @@
     fun unfold_progressUnavailable_logsLatencyTillScreenTurnedOn() {
         testScope.runTest {
             val unfoldTransitionInteractorWithEmptyProgressProvider =
-                UnfoldTransitionInteractorImpl(UnfoldTransitionRepositoryImpl(Optional.empty()))
+                UnfoldTransitionInteractor(
+                    UnfoldTransitionRepositoryImpl(Optional.empty()),
+                    configurationInteractor,
+                )
             displaySwitchLatencyTracker =
                 DisplaySwitchLatencyTracker(
                     mockContext,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt
index b4f1218..bdd4afa 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt
@@ -18,7 +18,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
-import com.android.systemui.keyguard.ui.keyguardTransitionAnimationFlow
+import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.statusbar.phone.statusBarKeyguardViewManager
@@ -27,6 +27,6 @@
 val Kosmos.alternateBouncerViewModel by Fixture {
     AlternateBouncerViewModel(
         statusBarKeyguardViewManager = statusBarKeyguardViewManager,
-        animationFlow = keyguardTransitionAnimationFlow,
+        keyguardTransitionInteractor = keyguardTransitionInteractor,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelKosmos.kt
index 45b28b1..cbba80b 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelKosmos.kt
@@ -46,6 +46,7 @@
 import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.statusbar.notification.stack.domain.interactor.notificationStackAppearanceInteractor
 import com.android.systemui.statusbar.notification.stack.domain.interactor.sharedNotificationContainerInteractor
+import com.android.systemui.unfold.domain.interactor.unfoldTransitionInteractor
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -81,5 +82,6 @@
         primaryBouncerToLockscreenTransitionViewModel =
             primaryBouncerToLockscreenTransitionViewModel,
         aodBurnInViewModel = aodBurnInViewModel,
+        unfoldTransitionInteractor = unfoldTransitionInteractor,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorKosmos.kt
index d03616a..6faa3b4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/unfold/domain/interactor/UnfoldTransitionInteractorKosmos.kt
@@ -16,10 +16,14 @@
 
 package com.android.systemui.unfold.domain.interactor
 
+import com.android.systemui.common.ui.domain.interactor.configurationInteractor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.unfold.data.repository.unfoldTransitionRepository
 
 val Kosmos.unfoldTransitionInteractor by Fixture {
-    UnfoldTransitionInteractorImpl(repository = unfoldTransitionRepository)
+    UnfoldTransitionInteractor(
+        repository = unfoldTransitionRepository,
+        configurationInteractor = configurationInteractor,
+    )
 }
diff --git a/services/core/java/com/android/server/audio/AudioService.java b/services/core/java/com/android/server/audio/AudioService.java
index 0e22ef1..c11fbe1 100644
--- a/services/core/java/com/android/server/audio/AudioService.java
+++ b/services/core/java/com/android/server/audio/AudioService.java
@@ -7910,6 +7910,7 @@
         DEVICE_MEDIA_UNMUTED_ON_PLUG_SET.add(AudioSystem.DEVICE_OUT_WIRED_HEADSET);
         DEVICE_MEDIA_UNMUTED_ON_PLUG_SET.add(AudioSystem.DEVICE_OUT_WIRED_HEADPHONE);
         DEVICE_MEDIA_UNMUTED_ON_PLUG_SET.add(AudioSystem.DEVICE_OUT_LINE);
+        DEVICE_MEDIA_UNMUTED_ON_PLUG_SET.add(AudioSystem.DEVICE_OUT_HEARING_AID);
         DEVICE_MEDIA_UNMUTED_ON_PLUG_SET.addAll(AudioSystem.DEVICE_OUT_ALL_A2DP_SET);
         DEVICE_MEDIA_UNMUTED_ON_PLUG_SET.addAll(AudioSystem.DEVICE_OUT_ALL_BLE_SET);
         DEVICE_MEDIA_UNMUTED_ON_PLUG_SET.addAll(AudioSystem.DEVICE_OUT_ALL_USB_SET);
diff --git a/services/core/java/com/android/server/display/DisplayDeviceConfig.java b/services/core/java/com/android/server/display/DisplayDeviceConfig.java
index 1c169a0..5c93181 100644
--- a/services/core/java/com/android/server/display/DisplayDeviceConfig.java
+++ b/services/core/java/com/android/server/display/DisplayDeviceConfig.java
@@ -1766,7 +1766,8 @@
                 loadDensityMapping(config);
                 loadBrightnessDefaultFromDdcXml(config);
                 loadBrightnessConstraintsFromConfigXml();
-                if (mFlags.isEvenDimmerEnabled()) {
+                if (mFlags.isEvenDimmerEnabled() && mContext.getResources().getBoolean(
+                        com.android.internal.R.bool.config_evenDimmerEnabled)) {
                     mEvenDimmerBrightnessData = EvenDimmerBrightnessData.loadConfig(config);
                 }
                 loadBrightnessMap(config);
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 8f1277b..4e709a7 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -5271,5 +5271,13 @@
         public ExternalDisplayStatsService getExternalDisplayStatsService() {
             return mExternalDisplayStatsService;
         }
+
+        /**
+         * Called on external display is ready to be enabled.
+         */
+        @Override
+        public void onExternalDisplayReadyToBeEnabled(int displayId) {
+            mDisplayModeDirector.onExternalDisplayReadyToBeEnabled(displayId);
+        }
     }
 }
diff --git a/services/core/java/com/android/server/display/ExternalDisplayPolicy.java b/services/core/java/com/android/server/display/ExternalDisplayPolicy.java
index b24caf4..3c2918f 100644
--- a/services/core/java/com/android/server/display/ExternalDisplayPolicy.java
+++ b/services/core/java/com/android/server/display/ExternalDisplayPolicy.java
@@ -91,6 +91,8 @@
 
         @NonNull
         ExternalDisplayStatsService getExternalDisplayStatsService();
+
+        void onExternalDisplayReadyToBeEnabled(int displayId);
     }
 
     @NonNull
@@ -185,6 +187,10 @@
             return;
         }
 
+        if (enabled) {
+            mInjector.onExternalDisplayReadyToBeEnabled(logicalDisplay.getDisplayIdLocked());
+        }
+
         mLogicalDisplayMapper.setDisplayEnabledLocked(logicalDisplay, enabled);
     }
 
@@ -217,6 +223,7 @@
         if ((Build.IS_ENG || Build.IS_USERDEBUG)
                 && SystemProperties.getBoolean(ENABLE_ON_CONNECT, false)) {
             Slog.w(TAG, "External display is enabled by default, bypassing user consent.");
+            mInjector.onExternalDisplayReadyToBeEnabled(logicalDisplay.getDisplayIdLocked());
             mInjector.sendExternalDisplayEventLocked(logicalDisplay, EVENT_DISPLAY_CONNECTED);
             return;
         } else {
diff --git a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
index a862b6e..9064763 100644
--- a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
+++ b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
@@ -694,6 +694,13 @@
     }
 
     /**
+     * Called when external display is ready to be enabled.
+     */
+    public void onExternalDisplayReadyToBeEnabled(int displayId) {
+        mDisplayObserver.onExternalDisplayReadyToBeEnabled(displayId);
+    }
+
+    /**
      * Listens for changes refresh rate coordination.
      */
     public interface DesiredDisplayModeSpecsListener {
@@ -1379,6 +1386,13 @@
             }
         }
 
+
+        void onExternalDisplayReadyToBeEnabled(int displayId) {
+            DisplayInfo displayInfo = getDisplayInfo(displayId);
+            updateDisplaysPeakRefreshRateAndResolution(displayInfo);
+            addDisplaysSynchronizedPeakRefreshRate(displayInfo);
+        }
+
         @Override
         public void onDisplayAdded(int displayId) {
             updateVrrStatus(displayId);
@@ -1386,8 +1400,6 @@
             updateDisplayModes(displayId, displayInfo);
             updateLayoutLimitedFrameRate(displayId, displayInfo);
             updateUserSettingDisplayPreferredSize(displayInfo);
-            updateDisplaysPeakRefreshRateAndResolution(displayInfo);
-            addDisplaysSynchronizedPeakRefreshRate(displayInfo);
         }
 
         @Override
diff --git a/services/core/java/com/android/server/security/FileIntegrityService.java b/services/core/java/com/android/server/security/FileIntegrityService.java
index bb4876b..5b501e1 100644
--- a/services/core/java/com/android/server/security/FileIntegrityService.java
+++ b/services/core/java/com/android/server/security/FileIntegrityService.java
@@ -170,6 +170,10 @@
         @Override
         public int setupFsverity(android.os.IInstalld.IFsveritySetupAuthToken authToken,
                 String filePath, String packageName) throws RemoteException {
+            getContext().enforceCallingPermission(android.Manifest.permission.SETUP_FSVERITY,
+                    "Permission android.permission.SETUP_FSVERITY not grantted to access "
+                    + "FileIntegrityManager#setupFsverity");
+
             Objects.requireNonNull(authToken);
             Objects.requireNonNull(filePath);
             Objects.requireNonNull(packageName);
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
index 5897d76..0877146 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
@@ -920,6 +920,7 @@
     @Test
     public void testEvenDimmer() throws IOException {
         when(mFlags.isEvenDimmerEnabled()).thenReturn(true);
+        when(mResources.getBoolean(R.bool.config_evenDimmerEnabled)).thenReturn(true);
         setupDisplayDeviceConfigFromDisplayConfigFile(getContent(getValidLuxThrottling(),
                 getValidProxSensor(), /* includeIdleMode= */ false, /* enableEvenDimmer */ true));
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/ExternalDisplayPolicyTest.java b/services/tests/displayservicetests/src/com/android/server/display/ExternalDisplayPolicyTest.java
index ea08be4..fe7bbe0 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/ExternalDisplayPolicyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/ExternalDisplayPolicyTest.java
@@ -157,6 +157,7 @@
         verify(mMockedLogicalDisplayMapper, never()).setDisplayEnabledLocked(any(), anyBoolean());
         verify(mMockedDisplayNotificationManager, times(2))
                 .onHighTemperatureExternalDisplayNotAllowed();
+        verify(mMockedInjector, never()).onExternalDisplayReadyToBeEnabled(anyInt());
     }
 
     @Test
@@ -167,6 +168,7 @@
         verify(mMockedLogicalDisplayMapper, never()).setDisplayEnabledLocked(any(), anyBoolean());
         verify(mMockedDisplayNotificationManager, never())
                 .onHighTemperatureExternalDisplayNotAllowed();
+        verify(mMockedInjector, never()).onExternalDisplayReadyToBeEnabled(anyInt());
     }
 
     @Test
@@ -184,6 +186,7 @@
         // Expected only 1 invocation, upon critical temperature.
         verify(mMockedDisplayNotificationManager).onHighTemperatureExternalDisplayNotAllowed();
         verify(mMockedExternalDisplayStatsService).onDisplayDisabled(eq(EXTERNAL_DISPLAY_ID));
+        verify(mMockedInjector, never()).onExternalDisplayReadyToBeEnabled(anyInt());
     }
 
     @Test
@@ -191,6 +194,7 @@
         mExternalDisplayPolicy.setExternalDisplayEnabledLocked(mMockedLogicalDisplay,
                 /*enabled=*/ true);
         assertDisplaySetEnabled(/*enabled=*/ true);
+        verify(mMockedInjector).onExternalDisplayReadyToBeEnabled(eq(EXTERNAL_DISPLAY_ID));
     }
 
     @Test
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java
index 0efd046..d670b13 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java
@@ -1945,7 +1945,7 @@
         SparseArray<SparseArray<Vote>> votesByDisplay = new SparseArray<>();
         votesByDisplay.put(DISPLAY_ID_2, votes);
 
-        director.getDisplayObserver().onDisplayAdded(DISPLAY_ID_2);
+        director.getDisplayObserver().onExternalDisplayReadyToBeEnabled(DISPLAY_ID_2);
         director.injectVotesByDisplay(votesByDisplay);
 
         var desiredSpecs = director.getDesiredDisplayModeSpecs(DISPLAY_ID_2);
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java
index d0dd921..1c192ef 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java
@@ -38,7 +38,6 @@
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.res.Resources;
-import android.hardware.display.DisplayManager;
 import android.hardware.display.DisplayManagerInternal;
 import android.os.Handler;
 import android.os.Looper;
@@ -109,7 +108,7 @@
     private Context mContext;
     private DisplayModeDirector.Injector mInjector;
     private Handler mHandler;
-    private DisplayManager.DisplayListener mObserver;
+    private DisplayModeDirector.DisplayObserver mObserver;
     private Resources mResources;
     @Mock
     private DisplayManagerFlags mDisplayManagerFlags;
@@ -161,6 +160,7 @@
                 .isEqualTo(null);
 
         // Testing that the vote is not added when display is added because feature is disabled
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
                 .isEqualTo(null);
@@ -194,6 +194,7 @@
         init();
         assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
                 .isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
                 .isEqualTo(null);
@@ -245,6 +246,7 @@
         init();
         assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
                 .isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
                 .isEqualTo(null);
@@ -277,6 +279,7 @@
         init();
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
                 .isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(DEFAULT_DISPLAY);
         mObserver.onDisplayAdded(DEFAULT_DISPLAY);
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
                 .isEqualTo(expectedResolutionVote);
@@ -298,6 +301,7 @@
                 .thenReturn(MAX_HEIGHT);
         init();
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(DEFAULT_DISPLAY);
         mObserver.onDisplayAdded(DEFAULT_DISPLAY);
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
         assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
@@ -317,6 +321,7 @@
                 .thenReturn(MAX_HEIGHT);
         init();
         assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
         assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(
@@ -336,6 +341,7 @@
         when(mDisplayManagerFlags.isExternalDisplayLimitModeEnabled()).thenReturn(true);
         init();
         assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
         assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_LIMIT_MODE)).isEqualTo(null);
@@ -358,6 +364,7 @@
                 .thenReturn(true);
         init();
         assertThat(getVote(GLOBAL_ID, PRIORITY_SYNCHRONIZED_REFRESH_RATE)).isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(GLOBAL_ID, PRIORITY_SYNCHRONIZED_REFRESH_RATE)).isEqualTo(
                 Vote.forPhysicalRefreshRates(
@@ -381,6 +388,7 @@
                 .thenReturn(true);
         init();
         assertThat(getVote(GLOBAL_ID, PRIORITY_SYNCHRONIZED_REFRESH_RATE)).isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(GLOBAL_ID, PRIORITY_SYNCHRONIZED_REFRESH_RATE)).isEqualTo(null);
     }
@@ -395,6 +403,7 @@
         when(mDisplayManagerFlags.isDisplaysRefreshRatesSynchronizationEnabled()).thenReturn(true);
         init();
         assertThat(getVote(GLOBAL_ID, PRIORITY_SYNCHRONIZED_REFRESH_RATE)).isEqualTo(null);
+        mObserver.onExternalDisplayReadyToBeEnabled(EXTERNAL_DISPLAY);
         mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
         assertThat(getVote(GLOBAL_ID, PRIORITY_SYNCHRONIZED_REFRESH_RATE)).isEqualTo(null);
     }
diff --git a/services/tests/servicestests/src/com/android/server/pm/TEST_MAPPING b/services/tests/servicestests/src/com/android/server/pm/TEST_MAPPING
index 861562d..305108e 100644
--- a/services/tests/servicestests/src/com/android/server/pm/TEST_MAPPING
+++ b/services/tests/servicestests/src/com/android/server/pm/TEST_MAPPING
@@ -1,31 +1,11 @@
 {
   "presubmit": [
     {
-      "name": "FrameworksServicesTests",
-      "options": [
-        {
-          "include-filter": "com.android.server.pm."
-        },
-        {
-          "include-annotation": "android.platform.test.annotations.Presubmit"
-        },
-        {
-          "exclude-annotation": "androidx.test.filters.FlakyTest"
-        },
-        {
-          "exclude-annotation": "org.junit.Ignore"
-        }
-      ]
+      "name": "FrameworksServicesTests_pm_presubmit"
     }
   ],
   "postsubmit": [
     {
-      // Presubmit is intentional here while testing with SLO checker.
-      // Tests are flaky, waiting to bypass.
-      "name": "FrameworksServicesTests_pm_presubmit"
-    },
-    {
-      // Leave postsubmit here when migrating
       "name": "FrameworksServicesTests_pm_postsubmit"
     }
   ]
diff --git a/telephony/java/android/telephony/euicc/EuiccManager.java b/telephony/java/android/telephony/euicc/EuiccManager.java
index ebabbf9..ca4a643 100644
--- a/telephony/java/android/telephony/euicc/EuiccManager.java
+++ b/telephony/java/android/telephony/euicc/EuiccManager.java
@@ -1039,17 +1039,18 @@
      * subscription on the
      * current eUICC and the subscription to be downloaded according to the subscription metadata.
      * Without the former, an {@link #EMBEDDED_SUBSCRIPTION_RESULT_RESOLVABLE_ERROR} will be
-     * eturned in the callback intent to prompt the user to accept the download.
+     * returned in the callback intent to prompt the user to accept the download.
      *
      * <p> Starting from Android {@link android.os.Build.VERSION_CODES#VANILLA_ICE_CREAM},
      * if the caller has the
      * {@code android.Manifest.permission#MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS} permission or
-     * is a profile owner or device owner, and
-     * {@code switchAfterDownload} is {@code false}, then the downloaded subscription
-     * will be managed by that caller. If {@code switchAfterDownload} is true,
-     * an {@link #EMBEDDED_SUBSCRIPTION_RESULT_RESOLVABLE_ERROR} will be
-     * returned in the callback intent to prompt the user to accept the download and the
-     * subscription will not be managed.
+     * is a profile owner or device owner, then the downloaded subscription
+     * will be managed by that caller.
+     * In case the caller is device owner or profile owner of an organization-owned device, {@code
+     * switchAfterDownload} can be set to true to automatically enable the subscription after
+     * download. If the caller is a profile owner on non organization owned device
+     * {@code switchAfterDownload} should be false otherwise the operation will fail with
+     * {@link #EMBEDDED_SUBSCRIPTION_RESULT_ERROR}.
      *
      * <p>On a multi-active SIM device, requires the
      * {@code android.Manifest.permission#WRITE_EMBEDDED_SUBSCRIPTIONS} permission, or a calling app