Merge "Extract TestDisplayWindowSettingsProvider class" into main
diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 8a2616e..0ca9789 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -78,6 +78,7 @@
         "android.view.inputmethod.flags-aconfig-java",
         "android.webkit.flags-aconfig-java",
         "android.widget.flags-aconfig-java",
+        "art_exported_aconfig_flags_lib",
         "backstage_power_flags_lib",
         "backup_flags_lib",
         "camera_platform_flags_core_java_lib",
@@ -140,6 +141,14 @@
     libs: ["fake_device_config"],
 }
 
+// ART
+java_aconfig_library {
+    name: "art_exported_aconfig_flags_lib",
+    aconfig_declarations: "art-aconfig-flags",
+    mode: "exported",
+    defaults: ["framework-minus-apex-aconfig-java-defaults"],
+}
+
 // Camera
 java_aconfig_library {
     name: "camera_platform_flags_core_java_lib",
diff --git a/Android.bp b/Android.bp
index 258440f..5b9f2cb 100644
--- a/Android.bp
+++ b/Android.bp
@@ -427,6 +427,7 @@
         "modules-utils-expresslog",
         "perfetto_trace_javastream_protos_jarjar",
         "libaconfig_java_proto_nano",
+        "aconfig_device_paths_java",
     ],
 }
 
diff --git a/cmds/screencap/Android.bp b/cmds/screencap/Android.bp
index c009c1f..16026ec 100644
--- a/cmds/screencap/Android.bp
+++ b/cmds/screencap/Android.bp
@@ -17,6 +17,7 @@
         "libutils",
         "libbinder",
         "libjnigraphics",
+        "libhwui",
         "libui",
         "libgui",
     ],
diff --git a/cmds/screencap/screencap.cpp b/cmds/screencap/screencap.cpp
index 01b20f4..12de82a 100644
--- a/cmds/screencap/screencap.cpp
+++ b/cmds/screencap/screencap.cpp
@@ -15,36 +15,28 @@
  */
 
 #include <android/bitmap.h>
+#include <android/graphics/bitmap.h>
 #include <android/gui/DisplayCaptureArgs.h>
 #include <binder/ProcessState.h>
 #include <errno.h>
-#include <unistd.h>
-#include <stdio.h>
 #include <fcntl.h>
-#include <stdlib.h>
-#include <string.h>
-#include <getopt.h>
-
-#include <linux/fb.h>
-#include <sys/ioctl.h>
-#include <sys/mman.h>
-#include <sys/wait.h>
-
-#include <android/bitmap.h>
-
-#include <binder/ProcessState.h>
-
 #include <ftl/concat.h>
 #include <ftl/optional.h>
+#include <getopt.h>
 #include <gui/ISurfaceComposer.h>
 #include <gui/SurfaceComposerClient.h>
 #include <gui/SyncScreenCaptureListener.h>
-
+#include <linux/fb.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/ioctl.h>
+#include <sys/mman.h>
+#include <sys/wait.h>
+#include <system/graphics.h>
 #include <ui/GraphicTypes.h>
 #include <ui/PixelFormat.h>
 
-#include <system/graphics.h>
-
 using namespace android;
 
 #define COLORSPACE_UNKNOWN    0
@@ -85,11 +77,12 @@
 };
 }
 
-static const struct option LONG_OPTIONS[] = {
-        {"png", no_argument, nullptr, 'p'},
-        {"help", no_argument, nullptr, 'h'},
-        {"hint-for-seamless", no_argument, nullptr, LongOpts::HintForSeamless},
-        {0, 0, 0, 0}};
+static const struct option LONG_OPTIONS[] = {{"png", no_argument, nullptr, 'p'},
+                                             {"jpeg", no_argument, nullptr, 'j'},
+                                             {"help", no_argument, nullptr, 'h'},
+                                             {"hint-for-seamless", no_argument, nullptr,
+                                              LongOpts::HintForSeamless},
+                                             {0, 0, 0, 0}};
 
 static int32_t flinger2bitmapFormat(PixelFormat f)
 {
@@ -170,10 +163,11 @@
     return 0;
 }
 
-status_t saveImage(const char* fn, bool png, const ScreenCaptureResults& captureResults) {
+status_t saveImage(const char* fn, std::optional<AndroidBitmapCompressFormat> format,
+                   const ScreenCaptureResults& captureResults) {
     void* base = nullptr;
     ui::Dataspace dataspace = captureResults.capturedDataspace;
-    sp<GraphicBuffer> buffer = captureResults.buffer;
+    const sp<GraphicBuffer>& buffer = captureResults.buffer;
 
     status_t result = buffer->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &base);
 
@@ -188,22 +182,48 @@
         return 1;
     }
 
+    void* gainmapBase = nullptr;
+    sp<GraphicBuffer> gainmap = captureResults.optionalGainMap;
+
+    if (gainmap) {
+        result = gainmap->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &gainmapBase);
+        if (gainmapBase == nullptr || result != NO_ERROR) {
+            fprintf(stderr, "Failed to capture gainmap with error code (%d)\n", result);
+            gainmapBase = nullptr;
+            // Fall-through: just don't attempt to write the gainmap
+        }
+    }
+
     int fd = -1;
     if (fn == nullptr) {
         fd = dup(STDOUT_FILENO);
         if (fd == -1) {
             fprintf(stderr, "Error writing to stdout. (%s)\n", strerror(errno));
+            if (gainmapBase) {
+                gainmap->unlock();
+            }
+
+            if (base) {
+                buffer->unlock();
+            }
             return 1;
         }
     } else {
         fd = open(fn, O_WRONLY | O_CREAT | O_TRUNC, 0664);
         if (fd == -1) {
             fprintf(stderr, "Error opening file: %s (%s)\n", fn, strerror(errno));
+            if (gainmapBase) {
+                gainmap->unlock();
+            }
+
+            if (base) {
+                buffer->unlock();
+            }
             return 1;
         }
     }
 
-    if (png) {
+    if (format) {
         AndroidBitmapInfo info;
         info.format = flinger2bitmapFormat(buffer->getPixelFormat());
         info.flags = ANDROID_BITMAP_FLAGS_ALPHA_PREMUL;
@@ -211,16 +231,31 @@
         info.height = buffer->getHeight();
         info.stride = buffer->getStride() * bytesPerPixel(buffer->getPixelFormat());
 
-        int result = AndroidBitmap_compress(&info, static_cast<int32_t>(dataspace), base,
-                                            ANDROID_BITMAP_COMPRESS_FORMAT_PNG, 100, &fd,
+        int result;
+
+        if (gainmapBase) {
+            result = ABitmap_compressWithGainmap(&info, static_cast<ADataSpace>(dataspace), base,
+                                                 gainmapBase, captureResults.hdrSdrRatio, *format,
+                                                 100, &fd,
+                                                 [](void* fdPtr, const void* data,
+                                                    size_t size) -> bool {
+                                                     int bytesWritten =
+                                                             write(*static_cast<int*>(fdPtr), data,
+                                                                   size);
+                                                     return bytesWritten == size;
+                                                 });
+        } else {
+            result = AndroidBitmap_compress(&info, static_cast<int32_t>(dataspace), base, *format,
+                                            100, &fd,
                                             [](void* fdPtr, const void* data, size_t size) -> bool {
                                                 int bytesWritten = write(*static_cast<int*>(fdPtr),
                                                                          data, size);
                                                 return bytesWritten == size;
                                             });
+        }
 
         if (result != ANDROID_BITMAP_RESULT_SUCCESS) {
-            fprintf(stderr, "Failed to compress PNG (error code: %d)\n", result);
+            fprintf(stderr, "Failed to compress (error code: %d)\n", result);
         }
 
         if (fn != NULL) {
@@ -245,6 +280,14 @@
     }
     close(fd);
 
+    if (gainmapBase) {
+        gainmap->unlock();
+    }
+
+    if (base) {
+        buffer->unlock();
+    }
+
     return 0;
 }
 
@@ -262,13 +305,17 @@
     gui::CaptureArgs captureArgs;
     const char* pname = argv[0];
     bool png = false;
+    bool jpeg = false;
     bool all = false;
     int c;
-    while ((c = getopt_long(argc, argv, "aphd:", LONG_OPTIONS, nullptr)) != -1) {
+    while ((c = getopt_long(argc, argv, "apjhd:", LONG_OPTIONS, nullptr)) != -1) {
         switch (c) {
             case 'p':
                 png = true;
                 break;
+            case 'j':
+                jpeg = true;
+                break;
             case 'd': {
                 errno = 0;
                 char* end = nullptr;
@@ -325,6 +372,14 @@
             baseName = filename.substr(0, filename.size()-4);
             suffix = ".png";
             png = true;
+        } else if (filename.ends_with(".jpeg")) {
+            baseName = filename.substr(0, filename.size() - 5);
+            suffix = ".jpeg";
+            jpeg = true;
+        } else if (filename.ends_with(".jpg")) {
+            baseName = filename.substr(0, filename.size() - 4);
+            suffix = ".jpg";
+            jpeg = true;
         } else {
             baseName = filename;
         }
@@ -350,6 +405,20 @@
         }
     }
 
+    if (png && jpeg) {
+        fprintf(stderr, "Ambiguous file type");
+        return 1;
+    }
+
+    std::optional<AndroidBitmapCompressFormat> format = std::nullopt;
+
+    if (png) {
+        format = ANDROID_BITMAP_COMPRESS_FORMAT_PNG;
+    } else if (jpeg) {
+        format = ANDROID_BITMAP_COMPRESS_FORMAT_JPEG;
+        captureArgs.attachGainmap = true;
+    }
+
     // setThreadPoolMaxThreadCount(0) actually tells the kernel it's
     // not allowed to spawn any additional threads, but we still spawn
     // a binder thread from userspace when we call startThreadPool().
@@ -385,7 +454,7 @@
         if (!filename.empty()) {
             fn = filename.c_str();
         }
-        if (const status_t saveImageStatus = saveImage(fn, png, result) != 0) {
+        if (const status_t saveImageStatus = saveImage(fn, format, result) != 0) {
             fprintf(stderr, "Saving image failed.\n");
             return saveImageStatus;
         }
diff --git a/core/api/current.txt b/core/api/current.txt
index dada20e..520c7d1 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -6396,6 +6396,7 @@
     method public android.graphics.drawable.Icon getLargeIcon();
     method @Nullable public android.content.LocusId getLocusId();
     method public CharSequence getSettingsText();
+    method @FlaggedApi("android.app.api_rich_ongoing") @Nullable public String getShortCriticalText();
     method public String getShortcutId();
     method public android.graphics.drawable.Icon getSmallIcon();
     method public String getSortKey();
@@ -6719,6 +6720,7 @@
     method @NonNull public android.app.Notification.Builder setPublicVersion(android.app.Notification);
     method @NonNull public android.app.Notification.Builder setRemoteInputHistory(CharSequence[]);
     method @NonNull public android.app.Notification.Builder setSettingsText(CharSequence);
+    method @FlaggedApi("android.app.api_rich_ongoing") @NonNull public android.app.Notification.Builder setShortCriticalText(@Nullable String);
     method @NonNull public android.app.Notification.Builder setShortcutId(String);
     method @NonNull public android.app.Notification.Builder setShowWhen(boolean);
     method @NonNull public android.app.Notification.Builder setSmallIcon(@DrawableRes int);
@@ -19168,7 +19170,7 @@
     method @NonNull public java.util.List<android.hardware.camera2.CaptureRequest.Key<?>> getAvailableCaptureRequestKeys();
     method @NonNull public java.util.List<android.hardware.camera2.CaptureResult.Key<?>> getAvailableCaptureResultKeys();
     method public java.util.List<android.hardware.camera2.CaptureRequest.Key<?>> getAvailablePhysicalCameraRequestKeys();
-    method @FlaggedApi("com.android.internal.camera.flags.feature_combination_query") @NonNull public java.util.List<android.hardware.camera2.CameraCharacteristics.Key<?>> getAvailableSessionCharacteristicsKeys();
+    method @NonNull public java.util.List<android.hardware.camera2.CameraCharacteristics.Key<?>> getAvailableSessionCharacteristicsKeys();
     method public java.util.List<android.hardware.camera2.CaptureRequest.Key<?>> getAvailableSessionKeys();
     method @NonNull public java.util.List<android.hardware.camera2.CameraCharacteristics.Key<?>> getKeys();
     method @NonNull public java.util.List<android.hardware.camera2.CameraCharacteristics.Key<?>> getKeysNeedingPermission();
@@ -19211,7 +19213,7 @@
     field @FlaggedApi("com.android.internal.camera.flags.camera_manual_flash_strength_control") @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<java.lang.Integer> FLASH_TORCH_STRENGTH_MAX_LEVEL;
     field @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<int[]> HOT_PIXEL_AVAILABLE_HOT_PIXEL_MODES;
     field @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<android.hardware.camera2.params.DeviceStateSensorOrientationMap> INFO_DEVICE_STATE_SENSOR_ORIENTATION_MAP;
-    field @FlaggedApi("com.android.internal.camera.flags.feature_combination_query") @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<java.lang.Integer> INFO_SESSION_CONFIGURATION_QUERY_VERSION;
+    field @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<java.lang.Integer> INFO_SESSION_CONFIGURATION_QUERY_VERSION;
     field @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<java.lang.Integer> INFO_SUPPORTED_HARDWARE_LEVEL;
     field @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<java.lang.String> INFO_VERSION;
     field @NonNull public static final android.hardware.camera2.CameraCharacteristics.Key<android.util.Size[]> JPEG_AVAILABLE_THUMBNAIL_SIZES;
@@ -20081,14 +20083,14 @@
 
   public final class ExtensionSessionConfiguration {
     ctor public ExtensionSessionConfiguration(int, @NonNull java.util.List<android.hardware.camera2.params.OutputConfiguration>, @NonNull java.util.concurrent.Executor, @NonNull android.hardware.camera2.CameraExtensionSession.StateCallback);
-    method @FlaggedApi("com.android.internal.camera.flags.extension_10_bit") public void clearColorSpace();
-    method @FlaggedApi("com.android.internal.camera.flags.extension_10_bit") @Nullable public android.graphics.ColorSpace getColorSpace();
+    method public void clearColorSpace();
+    method @Nullable public android.graphics.ColorSpace getColorSpace();
     method @NonNull public java.util.concurrent.Executor getExecutor();
     method public int getExtension();
     method @NonNull public java.util.List<android.hardware.camera2.params.OutputConfiguration> getOutputConfigurations();
     method @Nullable public android.hardware.camera2.params.OutputConfiguration getPostviewOutputConfiguration();
     method @NonNull public android.hardware.camera2.CameraExtensionSession.StateCallback getStateCallback();
-    method @FlaggedApi("com.android.internal.camera.flags.extension_10_bit") public void setColorSpace(@NonNull android.graphics.ColorSpace.Named);
+    method public void setColorSpace(@NonNull android.graphics.ColorSpace.Named);
     method public void setPostviewOutputConfiguration(@Nullable android.hardware.camera2.params.OutputConfiguration);
   }
 
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 475b1e2..9c807af 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -4972,12 +4972,12 @@
 
   @FlaggedApi("com.android.internal.camera.flags.concert_mode") public final class CameraOutputSurface {
     ctor @FlaggedApi("com.android.internal.camera.flags.concert_mode") public CameraOutputSurface(@NonNull android.view.Surface, @NonNull android.util.Size);
-    method @FlaggedApi("com.android.internal.camera.flags.extension_10_bit") public int getColorSpace();
-    method @FlaggedApi("com.android.internal.camera.flags.extension_10_bit") public long getDynamicRangeProfile();
+    method public int getColorSpace();
+    method public long getDynamicRangeProfile();
     method @FlaggedApi("com.android.internal.camera.flags.concert_mode") public int getImageFormat();
     method @FlaggedApi("com.android.internal.camera.flags.concert_mode") @NonNull public android.util.Size getSize();
     method @FlaggedApi("com.android.internal.camera.flags.concert_mode") @NonNull public android.view.Surface getSurface();
-    method @FlaggedApi("com.android.internal.camera.flags.extension_10_bit") public void setDynamicRangeProfile(long);
+    method public void setDynamicRangeProfile(long);
   }
 
   @FlaggedApi("com.android.internal.camera.flags.concert_mode") public class CharacteristicsMap {
@@ -4987,7 +4987,7 @@
 
   @FlaggedApi("com.android.internal.camera.flags.concert_mode") public class ExtensionConfiguration {
     ctor @FlaggedApi("com.android.internal.camera.flags.concert_mode") public ExtensionConfiguration(int, int, @NonNull java.util.List<android.hardware.camera2.extension.ExtensionOutputConfiguration>, @Nullable android.hardware.camera2.CaptureRequest);
-    method @FlaggedApi("com.android.internal.camera.flags.extension_10_bit") public void setColorSpace(int);
+    method public void setColorSpace(int);
   }
 
   @FlaggedApi("com.android.internal.camera.flags.concert_mode") public class ExtensionOutputConfiguration {
diff --git a/core/java/android/app/Instrumentation.java b/core/java/android/app/Instrumentation.java
index 45852c7..93a9489 100644
--- a/core/java/android/app/Instrumentation.java
+++ b/core/java/android/app/Instrumentation.java
@@ -2408,9 +2408,9 @@
      * @hide
      */
     @android.ravenwood.annotation.RavenwoodKeep
-    public final void basicInit(Context context) {
-        mInstrContext = context;
-        mAppContext = context;
+    public final void basicInit(Context instrContext, Context appContext) {
+        mInstrContext = instrContext;
+        mAppContext = appContext;
     }
 
     /** @hide */
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index 7a36fbb..81d2c89 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -1281,6 +1281,15 @@
     public static final String EXTRA_BIG_TEXT = "android.bigText";
 
     /**
+     * {@link #extras} key: very short text summarizing the most critical information contained in
+     * the notification.
+     *
+     * @hide
+     */
+    @FlaggedApi(Flags.FLAG_API_RICH_ONGOING)
+    public static final String EXTRA_SHORT_CRITICAL_TEXT = "android.shortCriticalText";
+
+    /**
      * {@link #extras} key: this is the resource ID of the notification's main small icon, as
      * supplied to {@link Builder#setSmallIcon(int)}.
      *
@@ -4050,6 +4059,17 @@
         return String.join("|", defaultStrings);
     }
 
+
+    /**
+     * Returns the very short text summarizing the most critical information contained in the
+     * notification, or null if this field was not set.
+     */
+    @Nullable
+    @FlaggedApi(Flags.FLAG_API_RICH_ONGOING)
+    public String getShortCriticalText() {
+        return extras.getString(EXTRA_SHORT_CRITICAL_TEXT);
+    }
+
     /**
      * @hide
      */
@@ -4991,6 +5011,18 @@
         }
 
         /**
+         * Sets a very short string summarizing the most critical information contained in the
+         * notification. Suggested max length is 5 characters, and there is no guarantee how much or
+         * how little of this text will be shown.
+         */
+        @FlaggedApi(Flags.FLAG_API_RICH_ONGOING)
+        @NonNull
+        public Builder setShortCriticalText(@Nullable String shortCriticalText) {
+            mN.extras.putString(EXTRA_SHORT_CRITICAL_TEXT, shortCriticalText);
+            return this;
+        }
+
+        /**
          * Set the progress this notification represents.
          *
          * The platform template will represent this using a {@link ProgressBar}.
diff --git a/core/java/android/content/pm/PackageManager.java b/core/java/android/content/pm/PackageManager.java
index 26bb6e4..fb2655c 100644
--- a/core/java/android/content/pm/PackageManager.java
+++ b/core/java/android/content/pm/PackageManager.java
@@ -8842,6 +8842,7 @@
 
         try {
             ParsedPackage pp = parser2.parsePackage(apkFile, parserFlags, false);
+            pp.hideAsFinal();
 
             return PackageInfoCommonUtils.generate(pp, flagsBits, UserHandle.myUserId());
         } catch (PackageParserException e) {
diff --git a/core/java/android/hardware/camera2/CameraCharacteristics.java b/core/java/android/hardware/camera2/CameraCharacteristics.java
index 056ca93..7a8a16f 100644
--- a/core/java/android/hardware/camera2/CameraCharacteristics.java
+++ b/core/java/android/hardware/camera2/CameraCharacteristics.java
@@ -567,7 +567,6 @@
      * @see #INFO_SESSION_CONFIGURATION_QUERY_VERSION
      */
     @NonNull
-    @FlaggedApi(Flags.FLAG_FEATURE_COMBINATION_QUERY)
     public List<CameraCharacteristics.Key<?>> getAvailableSessionCharacteristicsKeys() {
         if (mAvailableSessionCharacteristicsKeys != null) {
             return mAvailableSessionCharacteristicsKeys;
@@ -5210,7 +5209,6 @@
      */
     @PublicKey
     @NonNull
-    @FlaggedApi(Flags.FLAG_FEATURE_COMBINATION_QUERY)
     public static final Key<Integer> INFO_SESSION_CONFIGURATION_QUERY_VERSION =
             new Key<Integer>("android.info.sessionConfigurationQueryVersion", int.class);
 
diff --git a/core/java/android/hardware/camera2/extension/CameraOutputSurface.java b/core/java/android/hardware/camera2/extension/CameraOutputSurface.java
index 001b794..32139b8 100644
--- a/core/java/android/hardware/camera2/extension/CameraOutputSurface.java
+++ b/core/java/android/hardware/camera2/extension/CameraOutputSurface.java
@@ -107,7 +107,6 @@
      * {@link android.hardware.camera2.params.DynamicRangeProfiles.STANDARD}
      * unless specified by CameraOutputSurface.setDynamicRangeProfile.
      */
-    @FlaggedApi(Flags.FLAG_EXTENSION_10_BIT)
     public @DynamicRangeProfiles.Profile long getDynamicRangeProfile() {
         return mOutputSurface.dynamicRangeProfile;
     }
@@ -118,7 +117,6 @@
      * unless specified by CameraOutputSurface.setColorSpace.
      */
     @SuppressLint("MethodNameUnits")
-    @FlaggedApi(Flags.FLAG_EXTENSION_10_BIT)
     public int getColorSpace() {
         return mOutputSurface.colorSpace;
     }
@@ -128,7 +126,6 @@
      * will be {@link android.hardware.camera2.params.DynamicRangeProfiles.STANDARD}
      * unless explicitly set using this method.
      */
-    @FlaggedApi(Flags.FLAG_EXTENSION_10_BIT)
     public void setDynamicRangeProfile(
             @DynamicRangeProfiles.Profile long dynamicRangeProfile) {
         mOutputSurface.dynamicRangeProfile = dynamicRangeProfile;
diff --git a/core/java/android/hardware/camera2/extension/ExtensionConfiguration.java b/core/java/android/hardware/camera2/extension/ExtensionConfiguration.java
index 84b7a7f..32de1ce 100644
--- a/core/java/android/hardware/camera2/extension/ExtensionConfiguration.java
+++ b/core/java/android/hardware/camera2/extension/ExtensionConfiguration.java
@@ -83,7 +83,6 @@
      * The default will be -1, indicating an unspecified ColorSpace,
      * unless explicitly set using this method.
      */
-    @FlaggedApi(Flags.FLAG_EXTENSION_10_BIT)
     public void setColorSpace(int colorSpace) {
         mColorSpace = colorSpace;
     }
@@ -98,11 +97,7 @@
         ret.sessionTemplateId = mSessionTemplateId;
         ret.sessionType = mSessionType;
         ret.outputConfigs = new ArrayList<>(mOutputs.size());
-        if (Flags.extension10Bit()) {
-            ret.colorSpace = mColorSpace;
-        } else {
-            ret.colorSpace = ColorSpaceProfiles.UNSPECIFIED;
-        }
+        ret.colorSpace = mColorSpace;
         for (ExtensionOutputConfiguration outputConfig : mOutputs) {
             ret.outputConfigs.add(outputConfig.getOutputConfig());
         }
diff --git a/core/java/android/hardware/camera2/extension/ExtensionOutputConfiguration.java b/core/java/android/hardware/camera2/extension/ExtensionOutputConfiguration.java
index 3a67d61..8a47430 100644
--- a/core/java/android/hardware/camera2/extension/ExtensionOutputConfiguration.java
+++ b/core/java/android/hardware/camera2/extension/ExtensionOutputConfiguration.java
@@ -80,11 +80,7 @@
         config.outputId = new OutputConfigId();
         config.outputId.id = mOutputConfigId;
         config.surfaceGroupId = mSurfaceGroupId;
-        if (Flags.extension10Bit()) {
-            config.dynamicRangeProfile = surface.getDynamicRangeProfile();
-        } else {
-            config.dynamicRangeProfile = DynamicRangeProfiles.STANDARD;
-        }
+        config.dynamicRangeProfile = surface.getDynamicRangeProfile();
     }
 
     @Nullable CameraOutputConfig getOutputConfig() {
diff --git a/core/java/android/hardware/camera2/impl/CameraAdvancedExtensionSessionImpl.java b/core/java/android/hardware/camera2/impl/CameraAdvancedExtensionSessionImpl.java
index 2e1e90c..323712d 100644
--- a/core/java/android/hardware/camera2/impl/CameraAdvancedExtensionSessionImpl.java
+++ b/core/java/android/hardware/camera2/impl/CameraAdvancedExtensionSessionImpl.java
@@ -148,7 +148,7 @@
 
         for (OutputConfiguration c : config.getOutputConfigurations()) {
             if (c.getDynamicRangeProfile() != DynamicRangeProfiles.STANDARD) {
-                if (Flags.extension10Bit() && Flags.cameraExtensionsCharacteristicsGet()) {
+                if (Flags.cameraExtensionsCharacteristicsGet()) {
                     DynamicRangeProfiles dynamicProfiles = extensionChars.get(
                             config.getExtension(),
                             CameraCharacteristics.REQUEST_AVAILABLE_DYNAMIC_RANGE_PROFILES);
@@ -190,9 +190,7 @@
                 new IntArray(CameraExtensionUtils.SUPPORTED_CAPTURE_OUTPUT_FORMATS.length);
         supportedCaptureOutputFormats.addAll(
                 CameraExtensionUtils.SUPPORTED_CAPTURE_OUTPUT_FORMATS);
-        if (Flags.extension10Bit()) {
-            supportedCaptureOutputFormats.add(ImageFormat.YCBCR_P010);
-        }
+        supportedCaptureOutputFormats.add(ImageFormat.YCBCR_P010);
         for (int format : supportedCaptureOutputFormats.toArray()) {
             List<Size> supportedSizes = extensionChars.getExtensionSupportedSizes(
                     config.getExtension(), format);
@@ -359,22 +357,20 @@
             cameraOutput.setTimestampBase(OutputConfiguration.TIMESTAMP_BASE_SENSOR);
             cameraOutput.setReadoutTimestampEnabled(false);
             cameraOutput.setPhysicalCameraId(output.physicalCameraId);
-            if (Flags.extension10Bit()) {
-                boolean validDynamicRangeProfile = false;
-                for (long profile = DynamicRangeProfiles.STANDARD;
-                        profile < DynamicRangeProfiles.PUBLIC_MAX; profile <<= 1) {
-                    if (output.dynamicRangeProfile == profile) {
-                        validDynamicRangeProfile = true;
-                        break;
-                    }
+            boolean validDynamicRangeProfile = false;
+            for (long profile = DynamicRangeProfiles.STANDARD;
+                    profile < DynamicRangeProfiles.PUBLIC_MAX; profile <<= 1) {
+                if (output.dynamicRangeProfile == profile) {
+                    validDynamicRangeProfile = true;
+                    break;
                 }
-                if (validDynamicRangeProfile) {
-                    cameraOutput.setDynamicRangeProfile(output.dynamicRangeProfile);
-                } else {
-                    Log.e(TAG, "Extension configured dynamic range profile "
-                            + output.dynamicRangeProfile
-                            + " is not valid, using default DynamicRangeProfile.STANDARD");
-                }
+            }
+            if (validDynamicRangeProfile) {
+                cameraOutput.setDynamicRangeProfile(output.dynamicRangeProfile);
+            } else {
+                Log.e(TAG, "Extension configured dynamic range profile "
+                        + output.dynamicRangeProfile
+                        + " is not valid, using default DynamicRangeProfile.STANDARD");
             }
             outputList.add(cameraOutput);
             mCameraConfigMap.put(cameraOutput.getSurface(), output);
@@ -390,15 +386,13 @@
         SessionConfiguration sessionConfiguration = new SessionConfiguration(sessionType,
                 outputList, new CameraExtensionUtils.HandlerExecutor(mHandler),
                 new SessionStateHandler());
-        if (Flags.extension10Bit()) {
-            if (sessionConfig.colorSpace >= 0
-                    && sessionConfig.colorSpace < ColorSpace.Named.values().length) {
-                sessionConfiguration.setColorSpace(
-                        ColorSpace.Named.values()[sessionConfig.colorSpace]);
-            } else {
-                Log.e(TAG, "Extension configured color space " + sessionConfig.colorSpace
-                        + " is not valid, using default unspecified color space");
-            }
+        if (sessionConfig.colorSpace >= 0
+                && sessionConfig.colorSpace < ColorSpace.Named.values().length) {
+            sessionConfiguration.setColorSpace(
+                    ColorSpace.Named.values()[sessionConfig.colorSpace]);
+        } else {
+            Log.e(TAG, "Extension configured color space " + sessionConfig.colorSpace
+                    + " is not valid, using default unspecified color space");
         }
         if ((sessionConfig.sessionParameter != null) &&
                 (!sessionConfig.sessionParameter.isEmpty())) {
@@ -459,16 +453,11 @@
             ret.size.height = surfaceSize.getHeight();
             ret.imageFormat = SurfaceUtils.getSurfaceFormat(s);
 
-            if (Flags.extension10Bit()) {
-                ret.dynamicRangeProfile = o.getDynamicRangeProfile();
-                ColorSpace colorSpace = o.getColorSpace();
-                if (colorSpace != null) {
-                    ret.colorSpace = colorSpace.getId();
-                } else {
-                    ret.colorSpace = ColorSpaceProfiles.UNSPECIFIED;
-                }
+            ret.dynamicRangeProfile = o.getDynamicRangeProfile();
+            ColorSpace colorSpace = o.getColorSpace();
+            if (colorSpace != null) {
+                ret.colorSpace = colorSpace.getId();
             } else {
-                ret.dynamicRangeProfile = DynamicRangeProfiles.STANDARD;
                 ret.colorSpace = ColorSpaceProfiles.UNSPECIFIED;
             }
         } else {
diff --git a/core/java/android/hardware/camera2/impl/CameraExtensionJpegProcessor.java b/core/java/android/hardware/camera2/impl/CameraExtensionJpegProcessor.java
index a10e250..ab4ff72 100644
--- a/core/java/android/hardware/camera2/impl/CameraExtensionJpegProcessor.java
+++ b/core/java/android/hardware/camera2/impl/CameraExtensionJpegProcessor.java
@@ -208,10 +208,6 @@
     }
 
     public void onOutputSurface(Surface surface, int format) throws RemoteException {
-        if (!Flags.extension10Bit() && format != ImageFormat.JPEG) {
-            Log.e(TAG, "Unsupported output format: " + format);
-            return;
-        }
         CameraExtensionUtils.SurfaceInfo surfaceInfo = CameraExtensionUtils.querySurface(surface);
         mCaptureFormat = surfaceInfo.mFormat;
         mOutputSurface = surface;
@@ -221,10 +217,6 @@
     public void onPostviewOutputSurface(Surface surface) throws RemoteException {
         CameraExtensionUtils.SurfaceInfo postviewSurfaceInfo =
                 CameraExtensionUtils.querySurface(surface);
-        if (!Flags.extension10Bit() && postviewSurfaceInfo.mFormat != ImageFormat.JPEG) {
-            Log.e(TAG, "Unsupported output format: " + postviewSurfaceInfo.mFormat);
-            return;
-        }
         mPostviewFormat = postviewSurfaceInfo.mFormat;
         mPostviewOutputSurface = surface;
         initializePostviewPipeline();
@@ -240,10 +232,6 @@
     }
 
     public void onImageFormatUpdate(int format) throws RemoteException {
-        if (!Flags.extension10Bit() && format != ImageFormat.YUV_420_888) {
-            Log.e(TAG, "Unsupported input format: " + format);
-            return;
-        }
         mFormat = format;
         initializePipeline();
     }
@@ -251,7 +239,7 @@
     private void initializePipeline() throws RemoteException {
         if ((mFormat != -1) && (mOutputSurface != null) && (mResolution != null) &&
                 (mYuvReader == null)) {
-            if (Flags.extension10Bit() && mCaptureFormat == ImageFormat.YUV_420_888) {
+            if (mCaptureFormat == ImageFormat.YUV_420_888) {
                 // For the case when postview is JPEG and capture is YUV
                 mProcessor.onOutputSurface(mOutputSurface, mCaptureFormat);
             } else {
@@ -274,7 +262,7 @@
     private void initializePostviewPipeline() throws RemoteException {
         if ((mFormat != -1) && (mPostviewOutputSurface != null) && (mPostviewResolution != null)
                 && (mPostviewYuvReader == null)) {
-            if (Flags.extension10Bit() && mPostviewFormat == ImageFormat.YUV_420_888) {
+            if (mPostviewFormat == ImageFormat.YUV_420_888) {
                 // For the case when postview is YUV and capture is JPEG
                 mProcessor.onPostviewOutputSurface(mPostviewOutputSurface);
             } else {
diff --git a/core/java/android/hardware/camera2/impl/CameraExtensionSessionImpl.java b/core/java/android/hardware/camera2/impl/CameraExtensionSessionImpl.java
index a4ae398..ce1609d 100644
--- a/core/java/android/hardware/camera2/impl/CameraExtensionSessionImpl.java
+++ b/core/java/android/hardware/camera2/impl/CameraExtensionSessionImpl.java
@@ -190,9 +190,7 @@
                 new IntArray(CameraExtensionUtils.SUPPORTED_CAPTURE_OUTPUT_FORMATS.length);
         supportedCaptureOutputFormats.addAll(
                 CameraExtensionUtils.SUPPORTED_CAPTURE_OUTPUT_FORMATS);
-        if (Flags.extension10Bit()) {
-            supportedCaptureOutputFormats.add(ImageFormat.YCBCR_P010);
-        }
+        supportedCaptureOutputFormats.add(ImageFormat.YCBCR_P010);
         for (int format : supportedCaptureOutputFormats.toArray()) {
             List<Size> supportedSizes = extensionChars.getExtensionSupportedSizes(
                     config.getExtension(), format);
@@ -401,7 +399,7 @@
                 if (surfaceInfo.mFormat == ImageFormat.JPEG) {
                     mImageJpegProcessor = new CameraExtensionJpegProcessor(mImageProcessor);
                     mImageProcessor = mImageJpegProcessor;
-                } else if (Flags.extension10Bit() && mClientPostviewSurface != null) {
+                } else if (mClientPostviewSurface != null) {
                     // Handles case when postview is JPEG and capture is YUV
                     CameraExtensionUtils.SurfaceInfo postviewSurfaceInfo =
                             CameraExtensionUtils.querySurface(mClientPostviewSurface);
diff --git a/core/java/android/hardware/camera2/impl/CameraExtensionUtils.java b/core/java/android/hardware/camera2/impl/CameraExtensionUtils.java
index 40f0477..f91d277 100644
--- a/core/java/android/hardware/camera2/impl/CameraExtensionUtils.java
+++ b/core/java/android/hardware/camera2/impl/CameraExtensionUtils.java
@@ -113,32 +113,13 @@
 
         SurfaceInfo surfaceInfo = querySurface(outputConfig.getSurface());
 
-        if (Flags.extension10Bit()) {
-            Size postviewSize = new Size(surfaceInfo.mWidth, surfaceInfo.mHeight);
-            if (supportedPostviewSizes.get(surfaceInfo.mFormat)
-                    .contains(postviewSize)) {
-                return outputConfig.getSurface();
-            } else {
-                throw new IllegalArgumentException("Postview size not supported!");
-            }
+        Size postviewSize = new Size(surfaceInfo.mWidth, surfaceInfo.mHeight);
+        if (supportedPostviewSizes.get(surfaceInfo.mFormat)
+                .contains(postviewSize)) {
+            return outputConfig.getSurface();
         } else {
-            if (surfaceInfo.mFormat == captureFormat) {
-                if (supportedPostviewSizes.containsKey(captureFormat)) {
-                    Size postviewSize = new Size(surfaceInfo.mWidth, surfaceInfo.mHeight);
-                    if (supportedPostviewSizes.get(surfaceInfo.mFormat)
-                            .contains(postviewSize)) {
-                        return outputConfig.getSurface();
-                    } else {
-                        throw new IllegalArgumentException("Postview size not supported!");
-                    }
-                }
-            } else {
-                throw new IllegalArgumentException("Postview format should be equivalent to "
-                        + " the capture format!");
-            }
+            throw new IllegalArgumentException("Postview size not supported!");
         }
-
-        return null;
     }
 
     public static Surface getBurstCaptureSurface(
@@ -148,9 +129,7 @@
                 new IntArray(CameraExtensionUtils.SUPPORTED_CAPTURE_OUTPUT_FORMATS.length);
         supportedCaptureOutputFormats.addAll(
                 CameraExtensionUtils.SUPPORTED_CAPTURE_OUTPUT_FORMATS);
-        if (Flags.extension10Bit()) {
-            supportedCaptureOutputFormats.add(ImageFormat.YCBCR_P010);
-        }
+        supportedCaptureOutputFormats.add(ImageFormat.YCBCR_P010);
         for (OutputConfiguration config : outputConfigs) {
             SurfaceInfo surfaceInfo = querySurface(config.getSurface());
             for (int supportedFormat : supportedCaptureOutputFormats.toArray()) {
diff --git a/core/java/android/hardware/camera2/params/ExtensionSessionConfiguration.java b/core/java/android/hardware/camera2/params/ExtensionSessionConfiguration.java
index bf3f59f..73f4ff4 100644
--- a/core/java/android/hardware/camera2/params/ExtensionSessionConfiguration.java
+++ b/core/java/android/hardware/camera2/params/ExtensionSessionConfiguration.java
@@ -136,7 +136,6 @@
      * or the color space implied by the dataSpace passed into an {@link ImageReader}'s
      * constructor.</p>
      */
-    @FlaggedApi(Flags.FLAG_EXTENSION_10_BIT)
     public void setColorSpace(@NonNull ColorSpace.Named colorSpace) {
         mColorSpace = colorSpace.ordinal();
         for (OutputConfiguration outputConfiguration : mOutputs) {
@@ -150,7 +149,6 @@
     /**
      * Clear the color space, such that the default color space will be used.
      */
-    @FlaggedApi(Flags.FLAG_EXTENSION_10_BIT)
     public void clearColorSpace() {
         mColorSpace = ColorSpaceProfiles.UNSPECIFIED;
         for (OutputConfiguration outputConfiguration : mOutputs) {
@@ -167,7 +165,6 @@
      * @return the currently set color space, or null
      *         if not set
      */
-    @FlaggedApi(Flags.FLAG_EXTENSION_10_BIT)
     @SuppressLint("MethodNameUnits")
     public @Nullable ColorSpace getColorSpace() {
         if (mColorSpace != ColorSpaceProfiles.UNSPECIFIED) {
diff --git a/core/java/android/hardware/input/VirtualInputDeviceConfig.java b/core/java/android/hardware/input/VirtualInputDeviceConfig.java
index a87980c..e8ef8cd 100644
--- a/core/java/android/hardware/input/VirtualInputDeviceConfig.java
+++ b/core/java/android/hardware/input/VirtualInputDeviceConfig.java
@@ -57,7 +57,7 @@
         mVendorId = builder.mVendorId;
         mProductId = builder.mProductId;
         mAssociatedDisplayId = builder.mAssociatedDisplayId;
-        mInputDeviceName = Objects.requireNonNull(builder.mInputDeviceName);
+        mInputDeviceName = Objects.requireNonNull(builder.mInputDeviceName, "Missing device name");
 
         if (mAssociatedDisplayId == Display.INVALID_DISPLAY) {
             throw new IllegalArgumentException(
@@ -77,7 +77,7 @@
         mVendorId = in.readInt();
         mProductId = in.readInt();
         mAssociatedDisplayId = in.readInt();
-        mInputDeviceName = Objects.requireNonNull(in.readString8());
+        mInputDeviceName = Objects.requireNonNull(in.readString8(), "Missing device name");
     }
 
     /**
diff --git a/core/java/android/os/BatteryStats.java b/core/java/android/os/BatteryStats.java
index c4d12d4..996a288 100644
--- a/core/java/android/os/BatteryStats.java
+++ b/core/java/android/os/BatteryStats.java
@@ -2066,9 +2066,11 @@
         public static final int EVENT_LONG_WAKE_LOCK = 0x0014;
         // Event for reporting change of some device states, triggered by a specific UID
         public static final int EVENT_STATE_CHANGE = 0x0015;
+        // Event for reporting change of screen states.
+        public static final int EVENT_DISPLAY_STATE_CHANGED = 0x0016;
 
         // Number of event types.
-        public static final int EVENT_COUNT = 0x0016;
+        public static final int EVENT_COUNT = 0x0017;
         // Mask to extract out only the type part of the event.
         public static final int EVENT_TYPE_MASK = ~(EVENT_FLAG_START|EVENT_FLAG_FINISH);
 
@@ -3079,13 +3081,14 @@
     public static final String[] HISTORY_EVENT_NAMES = new String[] {
             "null", "proc", "fg", "top", "sync", "wake_lock_in", "job", "user", "userfg", "conn",
             "active", "pkginst", "pkgunin", "alarm", "stats", "pkginactive", "pkgactive",
-            "tmpwhitelist", "screenwake", "wakeupap", "longwake", "state"
+            "tmpwhitelist", "screenwake", "wakeupap", "longwake", "state",
+            "display_state_changed"
     };
 
     public static final String[] HISTORY_EVENT_CHECKIN_NAMES = new String[] {
             "Enl", "Epr", "Efg", "Etp", "Esy", "Ewl", "Ejb", "Eur", "Euf", "Ecn",
             "Eac", "Epi", "Epu", "Eal", "Est", "Eai", "Eaa", "Etw",
-            "Esw", "Ewa", "Elw", "Eec", "Esc"
+            "Esw", "Ewa", "Elw", "Eec", "Esc", "Eds"
     };
 
     @FunctionalInterface
diff --git a/core/java/android/os/SystemVibratorManager.java b/core/java/android/os/SystemVibratorManager.java
index 58ab5b6..cfbf528 100644
--- a/core/java/android/os/SystemVibratorManager.java
+++ b/core/java/android/os/SystemVibratorManager.java
@@ -138,11 +138,14 @@
             Log.w(TAG, "Failed to vibrate; no vibrator manager service.");
             return;
         }
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "vibrate, reason=" + reason);
         try {
             mService.vibrate(uid, mContext.getDeviceId(), opPkg, effect, attributes, reason,
                     mToken);
         } catch (RemoteException e) {
             Log.w(TAG, "Failed to vibrate.", e);
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -152,11 +155,14 @@
             Log.w(TAG, "Failed to perform haptic feedback; no vibrator manager service.");
             return;
         }
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "performHapticFeedback, reason=" + reason);
         try {
             mService.performHapticFeedback(mUid, mContext.getDeviceId(), mPackageName, constant,
                     reason, flags, privFlags);
         } catch (RemoteException e) {
             Log.w(TAG, "Failed to perform haptic feedback.", e);
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -168,11 +174,15 @@
                             + " no vibrator manager service.");
             return;
         }
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR,
+                "performHapticFeedbackForInputDevice, reason=" + reason);
         try {
             mService.performHapticFeedbackForInputDevice(mUid, mContext.getDeviceId(), mPackageName,
                     constant, inputDeviceId, inputSource, reason, flags, privFlags);
         } catch (RemoteException e) {
             Log.w(TAG, "Failed to perform haptic feedback for input device.", e);
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
diff --git a/core/java/android/service/notification/ZenModeConfig.java b/core/java/android/service/notification/ZenModeConfig.java
index 3d8d933..752f174 100644
--- a/core/java/android/service/notification/ZenModeConfig.java
+++ b/core/java/android/service/notification/ZenModeConfig.java
@@ -2553,7 +2553,7 @@
         if (!Flags.modesUi()) {
             return manualRule != null;
         }
-        return manualRule != null && manualRule.isAutomaticActive();
+        return manualRule != null && manualRule.isActive();
     }
 
     public static class ZenRule implements Parcelable {
@@ -2932,8 +2932,7 @@
             }
         }
 
-        // TODO: b/363193376 - Rename to isActive()
-        public boolean isAutomaticActive() {
+        public boolean isActive() {
             if (Flags.modesApi() && Flags.modesUi()) {
                 if (!enabled || getPkg() == null) {
                     return false;
@@ -3173,7 +3172,7 @@
 
         // DND turned on by an automatic rule
         for (ZenRule automaticRule : config.automaticRules.values()) {
-            if (automaticRule.isAutomaticActive()) {
+            if (automaticRule.isActive()) {
                 if (isValidEventConditionId(automaticRule.conditionId)
                         || isValidScheduleConditionId(automaticRule.conditionId)) {
                     // set text if automatic rule end time is the latest active rule end time
diff --git a/core/java/android/service/notification/ZenModeDiff.java b/core/java/android/service/notification/ZenModeDiff.java
index 05c2a9c..60a7d6b 100644
--- a/core/java/android/service/notification/ZenModeDiff.java
+++ b/core/java/android/service/notification/ZenModeDiff.java
@@ -495,8 +495,8 @@
 
             // Even if added or removed, there may be a change in whether or not it was active.
             // This only applies to automatic rules.
-            boolean fromActive = from != null ? from.isAutomaticActive() : false;
-            boolean toActive = to != null ? to.isAutomaticActive() : false;
+            boolean fromActive = from != null ? from.isActive() : false;
+            boolean toActive = to != null ? to.isActive() : false;
             if (fromActive != toActive) {
                 mActiveDiff = new FieldDiff<>(fromActive, toActive);
             }
diff --git a/core/java/android/view/flags/scroll_feedback_flags.aconfig b/core/java/android/view/flags/scroll_feedback_flags.aconfig
index e9c85684..658aa29 100644
--- a/core/java/android/view/flags/scroll_feedback_flags.aconfig
+++ b/core/java/android/view/flags/scroll_feedback_flags.aconfig
@@ -21,4 +21,5 @@
     name: "enable_touch_scroll_feedback"
     description: "Enables touchscreen haptic scroll feedback"
     bug: "331830899"
+    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 6e89c49..0f401d3 100644
--- a/core/java/android/window/flags/lse_desktop_experience.aconfig
+++ b/core/java/android/window/flags/lse_desktop_experience.aconfig
@@ -145,6 +145,13 @@
 }
 
 flag {
+    name: "enable_resizing_metrics"
+    namespace: "lse_desktop_experience"
+    description: "Whether to enable log collection for task resizing in desktop windowing mode"
+    bug: "341319100"
+}
+
+flag {
     name: "enable_caption_compat_inset_force_consumption"
     namespace: "lse_desktop_experience"
     description: "Enables force-consumption of caption bar insets for immersive apps in freeform"
diff --git a/core/java/com/android/internal/app/IBatteryStats.aidl b/core/java/com/android/internal/app/IBatteryStats.aidl
index ebcae27..a5e166b 100644
--- a/core/java/com/android/internal/app/IBatteryStats.aidl
+++ b/core/java/com/android/internal/app/IBatteryStats.aidl
@@ -144,9 +144,9 @@
     @EnforcePermission("UPDATE_DEVICE_STATS")
     void noteGpsSignalQuality(int signalLevel);
     @EnforcePermission("UPDATE_DEVICE_STATS")
-    void noteScreenState(int state);
+    void noteScreenState(int displayId, int state, int reason);
     @EnforcePermission("UPDATE_DEVICE_STATS")
-    void noteScreenBrightness(int brightness);
+    void noteScreenBrightness(int displayId, int brightness);
     @EnforcePermission("UPDATE_DEVICE_STATS")
     void noteUserActivity(int uid, int event);
     @EnforcePermission("UPDATE_DEVICE_STATS")
diff --git a/core/java/com/android/internal/pm/pkg/component/AconfigFlags.java b/core/java/com/android/internal/pm/pkg/component/AconfigFlags.java
index b873175..39aadfb 100644
--- a/core/java/com/android/internal/pm/pkg/component/AconfigFlags.java
+++ b/core/java/com/android/internal/pm/pkg/component/AconfigFlags.java
@@ -18,6 +18,7 @@
 
 import static com.android.internal.pm.pkg.parsing.ParsingUtils.ANDROID_RES_NAMESPACE;
 
+import android.aconfig.DeviceProtos;
 import android.aconfig.nano.Aconfig;
 import android.aconfig.nano.Aconfig.parsed_flag;
 import android.aconfig.nano.Aconfig.parsed_flags;
@@ -40,7 +41,7 @@
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
-import java.util.List;
+import java.util.Arrays;
 import java.util.Map;
 
 /**
@@ -54,12 +55,6 @@
 public class AconfigFlags {
     private static final String LOG_TAG = "AconfigFlags";
 
-    private static final List<String> sTextProtoFilesOnDevice = List.of(
-            "/system/etc/aconfig_flags.pb",
-            "/system_ext/etc/aconfig_flags.pb",
-            "/product/etc/aconfig_flags.pb",
-            "/vendor/etc/aconfig_flags.pb");
-
     public enum Permission {
         READ_WRITE,
         READ_ONLY
@@ -73,7 +68,10 @@
             Slog.v(LOG_TAG, "Feature disabled, skipped all loading");
             return;
         }
-        for (String fileName : sTextProtoFilesOnDevice) {
+        final var defaultFlagProtoFiles =
+                (Process.myUid() == Process.SYSTEM_UID) ? DeviceProtos.parsedFlagsProtoPaths()
+                        : Arrays.asList(DeviceProtos.PATHS);
+        for (String fileName : defaultFlagProtoFiles) {
             try (var inputStream = new FileInputStream(fileName)) {
                 loadAconfigDefaultValues(inputStream.readAllBytes());
             } catch (IOException e) {
diff --git a/core/tests/coretests/src/android/app/NotificationTest.java b/core/tests/coretests/src/android/app/NotificationTest.java
index ef6ff05..0837b45 100644
--- a/core/tests/coretests/src/android/app/NotificationTest.java
+++ b/core/tests/coretests/src/android/app/NotificationTest.java
@@ -215,6 +215,25 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_API_RICH_ONGOING)
+    public void testGetShortCriticalText_noneSet() {
+        Notification n = new Notification.Builder(mContext, "test")
+                .build();
+
+        assertSame(n.getShortCriticalText(), null);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_API_RICH_ONGOING)
+    public void testGetShortCriticalText_isSet() {
+        Notification n = new Notification.Builder(mContext, "test")
+                .setShortCriticalText("short critical text here")
+                .build();
+
+        assertSame(n.getShortCriticalText(), "short critical text here");
+    }
+
+    @Test
     public void largeIconMultipleReferences_keptAfterParcelling() {
         Icon originalIcon = Icon.createWithBitmap(BitmapFactory.decodeResource(
                 mContext.getResources(), com.android.frameworks.coretests.R.drawable.test128x96));
diff --git a/core/tests/overlaytests/handle_config_change/test-apps/OverlayResApp/Android.bp b/core/tests/overlaytests/handle_config_change/test-apps/OverlayResApp/Android.bp
index e0f1012..74c7b4c 100644
--- a/core/tests/overlaytests/handle_config_change/test-apps/OverlayResApp/Android.bp
+++ b/core/tests/overlaytests/handle_config_change/test-apps/OverlayResApp/Android.bp
@@ -32,8 +32,8 @@
         "truth",
     ],
     libs: [
-        "android.test.runner",
-        "android.test.base",
+        "android.test.runner.stubs.system",
+        "android.test.base.stubs.system",
     ],
     test_suites: [
         "device-tests",
diff --git a/libs/WindowManager/Shell/OWNERS b/libs/WindowManager/Shell/OWNERS
index c6044a4..394093c6 100644
--- a/libs/WindowManager/Shell/OWNERS
+++ b/libs/WindowManager/Shell/OWNERS
@@ -1,4 +1,6 @@
 xutan@google.com
+pbdr@google.com
+pragyabajoria@google.com
 
 # Give submodule owners in shell resource approval
 per-file res*/*/*.xml = atsjenk@google.com, hwwang@google.com, jorgegil@google.com, lbill@google.com, madym@google.com, vaniadesmonda@google.com, pbdr@google.com, tkachenkoi@google.com, mpodolian@google.com, liranb@google.com, pragyabajoria@google.com, uysalorhan@google.com, gsennton@google.com, mattsziklay@google.com, mdehaini@google.com
diff --git a/libs/WindowManager/Shell/multivalentScreenshotTests/Android.bp b/libs/WindowManager/Shell/multivalentScreenshotTests/Android.bp
index 1871203..b6db6d9 100644
--- a/libs/WindowManager/Shell/multivalentScreenshotTests/Android.bp
+++ b/libs/WindowManager/Shell/multivalentScreenshotTests/Android.bp
@@ -72,8 +72,8 @@
         "platform-screenshot-diff-core",
     ],
     libs: [
-        "android.test.base",
-        "android.test.runner",
+        "android.test.base.stubs.system",
+        "android.test.runner.stubs.system",
     ],
     jni_libs: [
         "libdexmakerjvmtiagent",
diff --git a/libs/WindowManager/Shell/res/drawable/desktop_mode_ic_handle_menu_manage_windows.xml b/libs/WindowManager/Shell/res/drawable/desktop_mode_ic_handle_menu_manage_windows.xml
new file mode 100644
index 0000000..7d912a2
--- /dev/null
+++ b/libs/WindowManager/Shell/res/drawable/desktop_mode_ic_handle_menu_manage_windows.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android" android:width="24dp" android:height="24dp" android:viewportWidth="960" android:viewportHeight="960" android:tint="?attr/colorControlNormal">
+    <path android:fillColor="@android:color/black" android:pathData="M160,880Q127,880 103.5,856.5Q80,833 80,800L80,440Q80,407 103.5,383.5Q127,360 160,360L240,360L240,160Q240,127 263.5,103.5Q287,80 320,80L800,80Q833,80 856.5,103.5Q880,127 880,160L880,520Q880,553 856.5,576.5Q833,600 800,600L720,600L720,800Q720,833 696.5,856.5Q673,880 640,880L160,880ZM160,800L640,800Q640,800 640,800Q640,800 640,800L640,520L160,520L160,800Q160,800 160,800Q160,800 160,800ZM720,520L800,520Q800,520 800,520Q800,520 800,520L800,240L320,240L320,360L640,360Q673,360 696.5,383.5Q720,407 720,440L720,520Z"/>
+</vector>
diff --git a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml
index eea3de8..64f71c7 100644
--- a/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml
+++ b/libs/WindowManager/Shell/res/layout/desktop_mode_window_decor_handle_menu.xml
@@ -147,6 +147,14 @@
             android:drawableStart="@drawable/desktop_mode_ic_handle_menu_new_window"
             android:drawableTint="?androidprv:attr/materialColorOnSurface"
             style="@style/DesktopModeHandleMenuActionButton" />
+
+        <Button
+            android:id="@+id/manage_windows_button"
+            android:contentDescription="@string/manage_windows_text"
+            android:text="@string/manage_windows_text"
+            android:drawableStart="@drawable/desktop_mode_ic_handle_menu_manage_windows"
+            android:drawableTint="?androidprv:attr/materialColorOnSurface"
+            style="@style/DesktopModeHandleMenuActionButton" />
     </LinearLayout>
 
     <LinearLayout
diff --git a/libs/WindowManager/Shell/res/values/dimen.xml b/libs/WindowManager/Shell/res/values/dimen.xml
index 755e0d5..c76c470 100644
--- a/libs/WindowManager/Shell/res/values/dimen.xml
+++ b/libs/WindowManager/Shell/res/values/dimen.xml
@@ -504,9 +504,9 @@
     <!-- The width of the handle menu in desktop mode.  -->
     <dimen name="desktop_mode_handle_menu_width">216dp</dimen>
 
-    <!-- The maximum height of the handle menu in desktop mode. Four pills (52dp each) plus 2dp
-        spacing between them plus 4dp top padding. -->
-    <dimen name="desktop_mode_handle_menu_height">270dp</dimen>
+    <!-- The maximum height of the handle menu in desktop mode. Three pills at 52dp each,
+         additional actions pill 156dp, plus 2dp spacing between them plus 4dp top padding. -->
+    <dimen name="desktop_mode_handle_menu_height">322dp</dimen>
 
     <!-- The elevation set on the handle menu pills. -->
     <dimen name="desktop_mode_handle_menu_pill_elevation">1dp</dimen>
@@ -520,6 +520,9 @@
     <!-- The maximum height of the handle menu's "New Window" button in desktop mode. -->
     <dimen name="desktop_mode_handle_menu_new_window_height">52dp</dimen>
 
+    <!-- The maximum height of the handle menu's "Manage Windows" button in desktop mode. -->
+    <dimen name="desktop_mode_handle_menu_manage_windows_height">52dp</dimen>
+
     <!-- The maximum height of the handle menu's "Screenshot" button in desktop mode. -->
     <dimen name="desktop_mode_handle_menu_screenshot_height">52dp</dimen>
 
diff --git a/libs/WindowManager/Shell/res/values/strings.xml b/libs/WindowManager/Shell/res/values/strings.xml
index a353db7..a6da421 100644
--- a/libs/WindowManager/Shell/res/values/strings.xml
+++ b/libs/WindowManager/Shell/res/values/strings.xml
@@ -294,6 +294,8 @@
     <string name="open_in_browser_text">Open in browser</string>
     <!-- Accessibility text for the handle menu new window button [CHAR LIMIT=NONE] -->
     <string name="new_window_text">New Window</string>
+    <!-- Accessibility text for the handle menu new window button [CHAR LIMIT=NONE] -->
+    <string name="manage_windows_text">Manage Windows</string>
     <!-- Accessibility text for the handle menu close button [CHAR LIMIT=NONE] -->
     <string name="close_text">Close</string>
     <!-- Accessibility text for the handle menu close menu button [CHAR LIMIT=NONE] -->
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/ManageWindowsViewContainer.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/ManageWindowsViewContainer.kt
new file mode 100644
index 0000000..79becb0
--- /dev/null
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/ManageWindowsViewContainer.kt
@@ -0,0 +1,173 @@
+/*
+ * 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.shared.desktopmode
+import android.annotation.ColorInt
+import android.content.Context
+import android.graphics.Bitmap
+import android.graphics.drawable.ShapeDrawable
+import android.graphics.drawable.shapes.RoundRectShape
+import android.util.TypedValue
+import android.view.MotionEvent.ACTION_OUTSIDE
+import android.view.SurfaceView
+import android.view.ViewGroup.MarginLayoutParams
+import android.widget.LinearLayout
+import android.window.TaskSnapshot
+
+/**
+ * View for the All Windows menu option, used by both Desktop Windowing and Taskbar.
+ * The menu displays icons of all open instances of an app. Clicking the icon should launch
+ * the instance, which will be performed by the child class.
+ */
+abstract class ManageWindowsViewContainer(
+    val context: Context,
+    @ColorInt private val menuBackgroundColor: Int
+) {
+    lateinit var menuView: ManageWindowsView
+
+    /** Creates the base menu view and fills it with icon views. */
+    fun show(snapshotList: List<Pair<Int, TaskSnapshot>>,
+             onIconClickListener: ((Int) -> Unit),
+             onOutsideClickListener: (() -> Unit)): ManageWindowsView {
+        menuView = ManageWindowsView(context, menuBackgroundColor).apply {
+            this.onOutsideClickListener = onOutsideClickListener
+            this.onIconClickListener = onIconClickListener
+            this.generateIconViews(snapshotList)
+        }
+        addToContainer(menuView)
+        return menuView
+    }
+
+    /** Adds the menu view to the container responsible for displaying it. */
+    abstract fun addToContainer(menuView: ManageWindowsView)
+
+    /** Dispose of the menu, perform needed cleanup. */
+    abstract fun close()
+
+    companion object {
+        const val MANAGE_WINDOWS_MINIMUM_INSTANCES = 2
+    }
+
+    class ManageWindowsView(
+        private val context: Context,
+        menuBackgroundColor: Int
+    ) {
+        val rootView: LinearLayout = LinearLayout(context)
+        var menuHeight = 0
+        var menuWidth = 0
+        var onIconClickListener: ((Int) -> Unit)? = null
+        var onOutsideClickListener: (() -> Unit)? = null
+
+        init {
+            rootView.orientation = LinearLayout.VERTICAL
+            val menuBackground = ShapeDrawable()
+            val menuRadius = getDimensionPixelSize(MENU_RADIUS_DP)
+            menuBackground.shape = RoundRectShape(
+                FloatArray(8) { menuRadius },
+                null,
+                null
+            )
+            menuBackground.paint.color = menuBackgroundColor
+            rootView.background = menuBackground
+            rootView.elevation = getDimensionPixelSize(MENU_ELEVATION_DP)
+            rootView.setOnTouchListener { _, event ->
+                if (event.actionMasked == ACTION_OUTSIDE) {
+                    onOutsideClickListener?.invoke()
+                }
+                return@setOnTouchListener true
+            }
+        }
+
+        private fun getDimensionPixelSize(sizeDp: Float): Float {
+            return TypedValue.applyDimension(TypedValue.COMPLEX_UNIT_DIP,
+                sizeDp, context.resources.displayMetrics)
+        }
+
+        fun generateIconViews(
+            snapshotList: List<Pair<Int, TaskSnapshot>>
+        ) {
+            menuWidth = 0
+            menuHeight = 0
+            rootView.removeAllViews()
+            val instanceIconHeight = getDimensionPixelSize(ICON_HEIGHT_DP)
+            val instanceIconWidth = getDimensionPixelSize(ICON_WIDTH_DP)
+            val iconRadius = getDimensionPixelSize(ICON_RADIUS_DP)
+            val iconMargin = getDimensionPixelSize(ICON_MARGIN_DP)
+            var rowLayout: LinearLayout? = null
+            // Add each icon to the menu, adding a new row when needed.
+            for ((iconCount, taskInfoSnapshotPair) in snapshotList.withIndex()) {
+                val taskId = taskInfoSnapshotPair.first
+                val snapshot = taskInfoSnapshotPair.second
+                // Once a row is filled, make a new row and increase the menu height.
+                if (iconCount % MENU_MAX_ICONS_PER_ROW == 0) {
+                    rowLayout = LinearLayout(context)
+                    rowLayout.orientation = LinearLayout.HORIZONTAL
+                    rootView.addView(rowLayout)
+                    menuHeight += (instanceIconHeight + iconMargin).toInt()
+                }
+                val snapshotBitmap = Bitmap.wrapHardwareBuffer(
+                    snapshot.hardwareBuffer,
+                    snapshot.colorSpace
+                )
+                val scaledSnapshotBitmap = snapshotBitmap?.let {
+                    Bitmap.createScaledBitmap(
+                        it, instanceIconWidth.toInt(), instanceIconHeight.toInt(), true /* filter */
+                    )
+                }
+                val appSnapshotButton = SurfaceView(context)
+                appSnapshotButton.cornerRadius = iconRadius
+                appSnapshotButton.setZOrderOnTop(true)
+                appSnapshotButton.setOnClickListener {
+                    onIconClickListener?.invoke(taskId)
+                }
+                val lp = MarginLayoutParams(
+                    instanceIconWidth.toInt(), instanceIconHeight.toInt()
+                )
+                lp.apply {
+                    marginStart = iconMargin.toInt()
+                    topMargin = iconMargin.toInt()
+                }
+                appSnapshotButton.layoutParams = lp
+                // If we haven't already reached one full row, increment width.
+                if (iconCount < MENU_MAX_ICONS_PER_ROW) {
+                    menuWidth += (instanceIconWidth + iconMargin).toInt()
+                }
+                rowLayout?.addView(appSnapshotButton)
+                appSnapshotButton.requestLayout()
+                rowLayout?.post {
+                    appSnapshotButton.holder.surface
+                        .attachAndQueueBufferWithColorSpace(
+                            scaledSnapshotBitmap?.hardwareBuffer,
+                            scaledSnapshotBitmap?.colorSpace
+                        )
+                }
+            }
+            // Add margin again for the right/bottom of the menu.
+            menuWidth += iconMargin.toInt()
+            menuHeight += iconMargin.toInt()
+        }
+
+        companion object {
+            private const val MENU_RADIUS_DP = 26f
+            private const val ICON_WIDTH_DP = 204f
+            private const val ICON_HEIGHT_DP = 127.5f
+            private const val ICON_RADIUS_DP = 16f
+            private const val ICON_MARGIN_DP = 16f
+            private const val MENU_ELEVATION_DP = 1f
+            private const val MENU_MAX_ICONS_PER_ROW = 3
+        }
+    }
+}
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 1d16980..7e96253 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
@@ -1077,45 +1077,47 @@
             request.triggerTask != null
     }
 
+    /** Open an existing instance of an app. */
+    fun openInstance(
+        callingTask: RunningTaskInfo,
+        requestedTaskId: Int
+    ) {
+        val wct = WindowContainerTransaction()
+        val options = createNewWindowOptions(callingTask)
+        if (options.launchWindowingMode == WINDOWING_MODE_FREEFORM) {
+            wct.startTask(requestedTaskId, options.toBundle())
+            transitions.startTransition(TRANSIT_OPEN, wct, null)
+        } else {
+            val splitPosition = splitScreenController.determineNewInstancePosition(callingTask)
+            splitScreenController.startTask(requestedTaskId, splitPosition,
+                options.toBundle(), null /* hideTaskToken */)
+        }
+    }
+
+    /** Create an Intent to open a new window of a task. */
     fun openNewWindow(
-        taskInfo: RunningTaskInfo
+        callingTaskInfo: RunningTaskInfo
     ) {
         // TODO(b/337915660): Add a transition handler for these; animations
         //  need updates in some cases.
-        val newTaskWindowingMode = when {
-            taskInfo.isFreeform -> {
-                WINDOWING_MODE_FREEFORM
-            }
-            taskInfo.isFullscreen || taskInfo.isMultiWindow -> {
-                WINDOWING_MODE_MULTI_WINDOW
-            }
-            else -> {
-                error("Invalid windowing mode: ${taskInfo.windowingMode}")
-            }
-        }
-
-        val baseActivity = taskInfo.baseActivity ?: return
+        val baseActivity = callingTaskInfo.baseActivity ?: return
         val fillIn: Intent = context.packageManager
             .getLaunchIntentForPackage(
                 baseActivity.packageName
             ) ?: return
         fillIn
             .setFlags(Intent.FLAG_ACTIVITY_NEW_TASK or Intent.FLAG_ACTIVITY_MULTIPLE_TASK)
-        val options =
-            ActivityOptions.makeBasic().apply {
-                launchWindowingMode = newTaskWindowingMode
-                pendingIntentBackgroundActivityStartMode =
-                    ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS
-            }
         val launchIntent = PendingIntent.getActivity(
             context,
             /* requestCode= */ 0,
             fillIn,
             PendingIntent.FLAG_IMMUTABLE
         )
-        when (newTaskWindowingMode) {
+        val options = createNewWindowOptions(callingTaskInfo)
+        when (options.launchWindowingMode) {
             WINDOWING_MODE_MULTI_WINDOW -> {
-                val splitPosition = splitScreenController.determineNewInstancePosition(taskInfo)
+                val splitPosition = splitScreenController
+                    .determineNewInstancePosition(callingTaskInfo)
                 splitScreenController.startIntent(
                     launchIntent, context.userId, fillIn, splitPosition,
                     options.toBundle(), null /* hideTaskToken */
@@ -1130,6 +1132,25 @@
         }
     }
 
+    private fun createNewWindowOptions(callingTask: RunningTaskInfo): ActivityOptions {
+        val newTaskWindowingMode = when {
+            callingTask.isFreeform -> {
+                WINDOWING_MODE_FREEFORM
+            }
+            callingTask.isFullscreen || callingTask.isMultiWindow -> {
+                WINDOWING_MODE_MULTI_WINDOW
+            }
+            else -> {
+                error("Invalid windowing mode: ${callingTask.windowingMode}")
+            }
+        }
+        return ActivityOptions.makeBasic().apply {
+            launchWindowingMode = newTaskWindowingMode
+            pendingIntentBackgroundActivityStartMode =
+                ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS
+        }
+    }
+
     /**
      * Handles the case where a freeform task is launched from recents.
      *
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
index d72ec90..dfc5ab3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
@@ -629,15 +629,20 @@
         finishTransaction: SurfaceControl.Transaction?
     ) {
         val state = transitionState ?: return
-        if (aborted && state.startTransitionToken == transition) {
+        if (!aborted) {
+            return
+        }
+        if (state.startTransitionToken == transition) {
             ProtoLog.v(
                 ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                 "DragToDesktop: onTransitionConsumed() start transition aborted"
             )
             state.startAborted = true
-            // Cancel CUJ interaction if the transition is aborted.
+            // The start-transition (DRAG_HOLD) is aborted, cancel its jank interaction.
             interactionJankMonitor.cancel(CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_HOLD)
         } else if (state.cancelTransitionToken != transition) {
+            // This transition being aborted is neither the start, nor the cancel transition, so
+            // it must be the finish transition (DRAG_RELEASE); cancel its jank interaction.
             interactionJankMonitor.cancel(CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_RELEASE)
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java
index cf02fb5..22e8dc1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java
@@ -26,7 +26,6 @@
 import static android.view.WindowManager.LayoutParams.FLAG_HARDWARE_ACCELERATED;
 import static android.view.WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE;
 import static android.view.WindowManager.LayoutParams.LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
-import static android.view.WindowManager.LayoutParams.MATCH_PARENT;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_INTERCEPT_GLOBAL_DRAG_AND_DROP;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_NO_MOVE_ANIMATION;
 import static android.view.WindowManager.LayoutParams.SYSTEM_FLAG_SHOW_FOR_ALL_USERS;
@@ -247,9 +246,8 @@
                 R.layout.global_drop_target, null);
         rootView.setOnDragListener(this);
         rootView.setVisibility(View.INVISIBLE);
-        DragLayout dragLayout = new DragLayout(context, mSplitScreen, mIconProvider);
-        rootView.addView(dragLayout,
-                new FrameLayout.LayoutParams(MATCH_PARENT, MATCH_PARENT));
+        DragLayoutProvider dragLayout = new DragLayout(context, mSplitScreen, mIconProvider);
+        dragLayout.addDraggingView(rootView);
         try {
             wm.addView(rootView, layoutParams);
             addDisplayDropTarget(displayId, context, wm, rootView, dragLayout);
@@ -261,7 +259,7 @@
 
     @VisibleForTesting
     void addDisplayDropTarget(int displayId, Context context, WindowManager wm,
-            FrameLayout rootView, DragLayout dragLayout) {
+            FrameLayout rootView, DragLayoutProvider dragLayout) {
         mDisplayDropTargets.put(displayId,
                 new PerDisplay(displayId, context, wm, rootView, dragLayout));
     }
@@ -564,7 +562,7 @@
         final Context context;
         final WindowManager wm;
         final FrameLayout rootView;
-        final DragLayout dragLayout;
+        final DragLayoutProvider dragLayout;
         // Tracks whether the window has fully drawn since it was last made visible
         boolean hasDrawn;
 
@@ -575,7 +573,7 @@
         // The active drag session
         DragSession dragSession;
 
-        PerDisplay(int dispId, Context c, WindowManager w, FrameLayout rv, DragLayout dl) {
+        PerDisplay(int dispId, Context c, WindowManager w, FrameLayout rv, DragLayoutProvider dl) {
             displayId = dispId;
             context = c;
             wm = w;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
index 3fecbe7..dfa2437 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
@@ -23,10 +23,10 @@
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
 import static android.view.ViewTreeObserver.InternalInsetsInfo.TOUCHABLE_INSETS_REGION;
 
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_BOTTOM;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_LEFT;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_RIGHT;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_TOP;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_BOTTOM;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_LEFT;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_RIGHT;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_TOP;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 
@@ -47,6 +47,7 @@
 import android.graphics.drawable.Drawable;
 import android.view.DragEvent;
 import android.view.SurfaceControl;
+import android.view.ViewGroup;
 import android.view.ViewTreeObserver;
 import android.view.WindowInsets;
 import android.view.WindowInsets.Type;
@@ -66,13 +67,13 @@
 import com.android.wm.shell.splitscreen.SplitScreenController;
 
 import java.io.PrintWriter;
-import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Coordinates the visible drop targets for the current drag within a single display.
  */
 public class DragLayout extends LinearLayout
-        implements ViewTreeObserver.OnComputeInternalInsetsListener {
+        implements ViewTreeObserver.OnComputeInternalInsetsListener, DragLayoutProvider {
 
     // While dragging the status bar is hidden.
     private static final int HIDE_STATUS_BAR_FLAGS = StatusBarManager.DISABLE_NOTIFICATION_ICONS
@@ -80,7 +81,7 @@
             | StatusBarManager.DISABLE_CLOCK
             | StatusBarManager.DISABLE_SYSTEM_INFO;
 
-    private final DragAndDropPolicy mPolicy;
+    private final DropTarget mPolicy;
     private final SplitScreenController mSplitScreenController;
     private final IconProvider mIconProvider;
     private final StatusBarManager mStatusBarManager;
@@ -91,7 +92,7 @@
     // Whether the device is currently in left/right split mode
     private boolean mIsLeftRightSplit;
 
-    private DragAndDropPolicy.Target mCurrentTarget = null;
+    private SplitDragPolicy.Target mCurrentTarget = null;
     private DropZoneView mDropZoneView1;
     private DropZoneView mDropZoneView2;
 
@@ -113,7 +114,7 @@
         super(context);
         mSplitScreenController = splitScreenController;
         mIconProvider = iconProvider;
-        mPolicy = new DragAndDropPolicy(context, splitScreenController);
+        mPolicy = new SplitDragPolicy(context, splitScreenController);
         mStatusBarManager = context.getSystemService(StatusBarManager.class);
         mLastConfiguration.setTo(context.getResources().getConfiguration());
 
@@ -387,6 +388,13 @@
         recomputeDropTargets();
     }
 
+    @NonNull
+    @Override
+    public void addDraggingView(ViewGroup rootView) {
+        // TODO(b/349828130) We need to separate out view + logic here
+        rootView.addView(this, new ViewGroup.LayoutParams(MATCH_PARENT, MATCH_PARENT));
+    }
+
     /**
      * Recalculates the drop targets based on the current policy.
      */
@@ -394,9 +402,9 @@
         if (!mIsShowing) {
             return;
         }
-        final ArrayList<DragAndDropPolicy.Target> targets = mPolicy.getTargets(mInsets);
+        final List<SplitDragPolicy.Target> targets = mPolicy.getTargets(mInsets);
         for (int i = 0; i < targets.size(); i++) {
-            final DragAndDropPolicy.Target target = targets.get(i);
+            final SplitDragPolicy.Target target = targets.get(i);
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DRAG_AND_DROP, "Add target: %s", target);
             // Inset the draw region by a little bit
             target.drawRegion.inset(mDisplayMargin, mDisplayMargin);
@@ -419,7 +427,7 @@
         }
         // Find containing region, if the same as mCurrentRegion, then skip, otherwise, animate the
         // visibility of the current region
-        DragAndDropPolicy.Target target = mPolicy.getTargetAtLocation(x, y);
+        SplitDragPolicy.Target target = mPolicy.getTargetAtLocation(x, y);
         if (mCurrentTarget != target) {
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DRAG_AND_DROP, "Current target: %s", target);
             if (target == null) {
@@ -493,7 +501,7 @@
         mHasDropped = true;
 
         // Process the drop
-        mPolicy.handleDrop(mCurrentTarget, hideTaskToken);
+        mPolicy.onDropped(mCurrentTarget, hideTaskToken);
 
         // Start animating the drop UI out with the drag surface
         hide(event, dropCompleteCallback);
@@ -576,7 +584,7 @@
         }
     }
 
-    private void animateHighlight(DragAndDropPolicy.Target target) {
+    private void animateHighlight(SplitDragPolicy.Target target) {
         if (target.type == TYPE_SPLIT_LEFT || target.type == TYPE_SPLIT_TOP) {
             mDropZoneView1.setShowingHighlight(true);
             mDropZoneView2.setShowingHighlight(false);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayoutProvider.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayoutProvider.kt
new file mode 100644
index 0000000..3d40824
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayoutProvider.kt
@@ -0,0 +1,80 @@
+/*
+ * 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.draganddrop
+
+import android.content.res.Configuration
+import android.view.DragEvent
+import android.view.SurfaceControl
+import android.view.ViewGroup
+import android.window.WindowContainerToken
+import com.android.internal.logging.InstanceId
+import java.io.PrintWriter
+
+/** Interface to be implemented by any controllers providing a layout for DragAndDrop in Shell */
+interface DragLayoutProvider {
+    /**
+     * Updates the drag layout based on the given drag session.
+     */
+    fun updateSession(session: DragSession)
+    /**
+     * Called when a new drag is started.
+     */
+    fun prepare(session: DragSession, loggerSessionId: InstanceId?)
+
+    /**
+     * Shows the drag layout.
+     */
+    fun show()
+
+    /**
+     * Updates the visible drop target as the user drags.
+     */
+    fun update(event: DragEvent?)
+
+    /**
+     * Hides the drag layout and animates out the visible drop targets.
+     */
+    fun hide(event: DragEvent?, hideCompleteCallback: Runnable?)
+
+    /**
+     * Whether target has already been dropped or not
+     */
+    fun hasDropped(): Boolean
+
+    /**
+     * Handles the drop onto a target and animates out the visible drop targets.
+     */
+    fun drop(
+        event: DragEvent?, dragSurface: SurfaceControl,
+        hideTaskToken: WindowContainerToken?, dropCompleteCallback: Runnable?
+    ): Boolean
+
+    /**
+     * Dumps information about this drag layout.
+     */
+    fun dump(pw: PrintWriter, prefix: String?)
+
+    /**
+     * @return a View which will be added to the global root view for drag and drop
+     */
+    fun addDraggingView(viewGroup: ViewGroup)
+
+    /**
+     * Called when the configuration changes.
+     */
+    fun onConfigChanged(newConfig: Configuration?)
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DropTarget.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DropTarget.kt
new file mode 100644
index 0000000..122a105
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DropTarget.kt
@@ -0,0 +1,56 @@
+/*
+ * 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.draganddrop
+
+import android.graphics.Insets
+import android.window.WindowContainerToken
+import com.android.internal.logging.InstanceId
+
+/**
+ * Interface to be implemented by classes which want to provide drop targets
+ * for DragAndDrop in Shell
+ */
+interface DropTarget {
+    // TODO(b/349828130) Delete after flexible split launches
+    /**
+     * Called at the start of a Drag, before input events are processed.
+     */
+    fun start(dragSession: DragSession, logSessionId: InstanceId)
+    /**
+     * @return [SplitDragPolicy.Target] corresponding to the given coords in display bounds.
+     */
+    fun getTargetAtLocation(x: Int, y: Int) : SplitDragPolicy.Target
+    /**
+     * @return total number of drop targets for the current drag session.
+     */
+    fun getNumTargets() : Int
+    // TODO(b/349828130)
+
+    /**
+     * @return [List<SplitDragPolicy.Target>] to show for the current drag session.
+     */
+    fun getTargets(insets: Insets) : List<SplitDragPolicy.Target>
+    /**
+     * Called when user is hovering Drag object over the given Target
+     */
+    fun onHoveringOver(target: SplitDragPolicy.Target) {}
+    /**
+     * Called when the user has dropped the provided target (need not be the same target as
+     * [onHoveringOver])
+     */
+    fun onDropped(target: SplitDragPolicy.Target, hideTaskToken: WindowContainerToken)
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropPolicy.java b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/SplitDragPolicy.java
similarity index 93%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropPolicy.java
rename to libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/SplitDragPolicy.java
index 6fec0c1..2a19d65 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropPolicy.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/SplitDragPolicy.java
@@ -32,11 +32,11 @@
 import static android.content.Intent.FLAG_ACTIVITY_MULTIPLE_TASK;
 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
 
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_FULLSCREEN;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_BOTTOM;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_LEFT;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_RIGHT;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_TOP;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_FULLSCREEN;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_BOTTOM;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_LEFT;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_RIGHT;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_TOP;
 import static com.android.wm.shell.shared.draganddrop.DragAndDropConstants.EXTRA_DISALLOW_HIT_REGION;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
@@ -80,9 +80,9 @@
 /**
  * The policy for handling drag and drop operations to shell.
  */
-public class DragAndDropPolicy {
+public class SplitDragPolicy implements DropTarget {
 
-    private static final String TAG = DragAndDropPolicy.class.getSimpleName();
+    private static final String TAG = SplitDragPolicy.class.getSimpleName();
 
     private final Context mContext;
     // Used only for launching a fullscreen task (or as a fallback if there is no split starter)
@@ -90,18 +90,18 @@
     // Used for launching tasks into splitscreen
     private final Starter mSplitscreenStarter;
     private final SplitScreenController mSplitScreen;
-    private final ArrayList<DragAndDropPolicy.Target> mTargets = new ArrayList<>();
+    private final ArrayList<SplitDragPolicy.Target> mTargets = new ArrayList<>();
     private final RectF mDisallowHitRegion = new RectF();
 
     private InstanceId mLoggerSessionId;
     private DragSession mSession;
 
-    public DragAndDropPolicy(Context context, SplitScreenController splitScreen) {
+    public SplitDragPolicy(Context context, SplitScreenController splitScreen) {
         this(context, splitScreen, new DefaultStarter(context));
     }
 
     @VisibleForTesting
-    DragAndDropPolicy(Context context, SplitScreenController splitScreen,
+    SplitDragPolicy(Context context, SplitScreenController splitScreen,
             Starter fullscreenStarter) {
         mContext = context;
         mSplitScreen = splitScreen;
@@ -112,7 +112,7 @@
     /**
      * Starts a new drag session with the given initial drag data.
      */
-    void start(DragSession session, InstanceId loggerSessionId) {
+    public void start(DragSession session, InstanceId loggerSessionId) {
         mLoggerSessionId = loggerSessionId;
         mSession = session;
         RectF disallowHitRegion = mSession.appData != null
@@ -128,7 +128,8 @@
     /**
      * Returns the number of targets.
      */
-    int getNumTargets() {
+    @Override
+    public int getNumTargets() {
         return mTargets.size();
     }
 
@@ -136,7 +137,8 @@
      * Returns the target's regions based on the current state of the device and display.
      */
     @NonNull
-    ArrayList<Target> getTargets(Insets insets) {
+    @Override
+    public ArrayList<Target> getTargets(@NonNull Insets insets) {
         mTargets.clear();
         if (mSession == null) {
             // Return early if this isn't an app drag
@@ -222,12 +224,12 @@
      * Returns the target at the given position based on the targets previously calculated.
      */
     @Nullable
-    Target getTargetAtLocation(int x, int y) {
+    public Target getTargetAtLocation(int x, int y) {
         if (mDisallowHitRegion.contains(x, y)) {
             return null;
         }
         for (int i = mTargets.size() - 1; i >= 0; i--) {
-            DragAndDropPolicy.Target t = mTargets.get(i);
+            SplitDragPolicy.Target t = mTargets.get(i);
             if (t.hitRegion.contains(x, y)) {
                 return t;
             }
@@ -241,7 +243,7 @@
      * container transaction if possible.
      */
     @VisibleForTesting
-    void handleDrop(Target target, @Nullable WindowContainerToken hideTaskToken) {
+    public void onDropped(Target target, @Nullable WindowContainerToken hideTaskToken) {
         if (target == null || !mTargets.contains(target)) {
             return;
         }
@@ -419,8 +421,9 @@
 
     /**
      * Represents a drop target.
+     * TODO(b/349828130): Move this into {@link DropTarget}
      */
-    static class Target {
+    public static class Target {
         static final int TYPE_FULLSCREEN = 0;
         static final int TYPE_SPLIT_LEFT = 1;
         static final int TYPE_SPLIT_TOP = 2;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
index ab222c9..b4e0329 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
@@ -2033,7 +2033,10 @@
         tx.apply();
     }
 
-    private void cancelCurrentAnimator() {
+    /**
+     * Cancels the currently running animator if there is one and removes an overlay if present.
+     */
+    public void cancelCurrentAnimator() {
         final PipAnimationController.PipTransitionAnimator<?> animator =
                 mPipAnimationController.getCurrentAnimator();
         // remove any overlays if present
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
index 755e958..deb7691 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipController.java
@@ -324,7 +324,7 @@
                         return;
                     }
                     onDisplayChanged(mDisplayController.getDisplayLayout(displayId),
-                            false /* saveRestoreSnapFraction */);
+                            true /* saveRestoreSnapFraction */);
                 }
 
                 @Override
@@ -652,9 +652,11 @@
                             // there's a keyguard present
                             return;
                         }
-                        onDisplayChangedUncheck(mDisplayController
-                                        .getDisplayLayout(mPipDisplayLayoutState.getDisplayId()),
-                                false /* saveRestoreSnapFraction */);
+                        mMainExecutor.executeDelayed(() -> {
+                            onDisplayChangedUncheck(mDisplayController.getDisplayLayout(
+                                    mPipDisplayLayoutState.getDisplayId()),
+                                    false /* saveRestoreSnapFraction */);
+                        }, PIP_KEEP_CLEAR_AREAS_DELAY);
                     }
                 });
 
@@ -800,7 +802,7 @@
             }
         };
 
-        if (mPipTaskOrganizer.isInPip() && saveRestoreSnapFraction) {
+        if (mPipTransitionState.hasEnteredPip() && saveRestoreSnapFraction) {
             mMenuController.attachPipMenuView();
             // Calculate the snap fraction of the current stack along the old movement bounds
             final PipSnapAlgorithm pipSnapAlgorithm = mPipBoundsAlgorithm.getSnapAlgorithm();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 793e2aa..87b661d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -92,7 +92,7 @@
 import com.android.wm.shell.common.split.SplitScreenUtils;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.draganddrop.DragAndDropController;
-import com.android.wm.shell.draganddrop.DragAndDropPolicy;
+import com.android.wm.shell.draganddrop.SplitDragPolicy;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.shared.TransactionPool;
@@ -121,7 +121,7 @@
  * @see StageCoordinator
  */
 // TODO(b/198577848): Implement split screen flicker test to consolidate CUJ of split screen.
-public class SplitScreenController implements DragAndDropPolicy.Starter,
+public class SplitScreenController implements SplitDragPolicy.Starter,
         RemoteCallable<SplitScreenController>, KeyguardChangeListener {
     private static final String TAG = SplitScreenController.class.getSimpleName();
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHandleManageWindowsMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHandleManageWindowsMenu.kt
new file mode 100644
index 0000000..13a805a
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHandleManageWindowsMenu.kt
@@ -0,0 +1,109 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.windowdecor
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.content.Context
+import android.graphics.Point
+import android.graphics.Rect
+import android.view.WindowManager
+import android.window.TaskSnapshot
+import androidx.compose.ui.graphics.toArgb
+import com.android.wm.shell.shared.desktopmode.ManageWindowsViewContainer
+import com.android.wm.shell.shared.split.SplitScreenConstants
+import com.android.wm.shell.splitscreen.SplitScreenController
+import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalSystemViewContainer
+import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalViewContainer
+import com.android.wm.shell.windowdecor.common.DecorThemeUtil
+import com.android.wm.shell.windowdecor.extension.isFullscreen
+import com.android.wm.shell.windowdecor.extension.isMultiWindow
+
+/**
+ * Implementation of [ManageWindowsViewContainer] meant to be used in desktop header and app
+ * handle.
+ */
+class DesktopHandleManageWindowsMenu(
+    private val callerTaskInfo: RunningTaskInfo,
+    private val splitScreenController: SplitScreenController,
+    private val captionX: Int,
+    private val captionWidth: Int,
+    private val windowManagerWrapper: WindowManagerWrapper,
+    context: Context,
+    snapshotList: List<Pair<Int, TaskSnapshot>>,
+    onIconClickListener: ((Int) -> Unit),
+    onOutsideClickListener: (() -> Unit)
+) : ManageWindowsViewContainer(
+    context,
+    DecorThemeUtil(context).getColorScheme(callerTaskInfo).background.toArgb()
+) {
+    private var menuViewContainer: AdditionalViewContainer? = null
+
+    init {
+        show(snapshotList, onIconClickListener, onOutsideClickListener)
+    }
+
+    override fun close() {
+        menuViewContainer?.releaseView()
+    }
+
+    private fun calculateMenuPosition(): Point {
+        val position = Point()
+        val nonFreeformX = (captionX + (captionWidth / 2) - (menuView.menuWidth / 2))
+        when {
+            callerTaskInfo.isFreeform -> {
+                val taskBounds = callerTaskInfo.getConfiguration().windowConfiguration.bounds
+                position.set(taskBounds.left, taskBounds.top)
+            }
+            callerTaskInfo.isFullscreen -> {
+                position.set(nonFreeformX, 0)
+            }
+            callerTaskInfo.isMultiWindow -> {
+                val splitPosition = splitScreenController.getSplitPosition(callerTaskInfo.taskId)
+                val leftOrTopStageBounds = Rect()
+                val rightOrBottomStageBounds = Rect()
+                splitScreenController.getStageBounds(leftOrTopStageBounds, rightOrBottomStageBounds)
+                // TODO(b/343561161): This needs to be calculated differently if the task is in
+                //  top/bottom split.
+                when (splitPosition) {
+                    SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT -> {
+                        position.set(leftOrTopStageBounds.width() + nonFreeformX, /* y= */ 0)
+                    }
+
+                    SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT -> {
+                        position.set(nonFreeformX, /* y= */ 0)
+                    }
+                }
+            }
+        }
+        return position
+    }
+
+    override fun addToContainer(menuView: ManageWindowsView) {
+        val menuPosition = calculateMenuPosition()
+        menuViewContainer = AdditionalSystemViewContainer(
+            windowManagerWrapper,
+            callerTaskInfo.taskId,
+            menuPosition.x,
+            menuPosition.y,
+            menuView.menuWidth,
+            menuView.menuHeight,
+            WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE or
+                    WindowManager.LayoutParams.FLAG_WATCH_OUTSIDE_TOUCH,
+            menuView.rootView
+        )
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenu.kt
new file mode 100644
index 0000000..05391a8
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopHeaderManageWindowsMenu.kt
@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.windowdecor
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.content.Context
+import android.graphics.PixelFormat
+import android.graphics.Point
+import android.view.SurfaceControl
+import android.view.SurfaceControlViewHost
+import android.view.WindowManager
+import android.view.WindowlessWindowManager
+import android.window.TaskConstants
+import android.window.TaskSnapshot
+import androidx.compose.ui.graphics.toArgb
+import com.android.wm.shell.RootTaskDisplayAreaOrganizer
+import com.android.wm.shell.common.DisplayController
+import com.android.wm.shell.shared.desktopmode.ManageWindowsViewContainer
+import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalViewContainer
+import com.android.wm.shell.windowdecor.additionalviewcontainer.AdditionalViewHostViewContainer
+import com.android.wm.shell.windowdecor.common.DecorThemeUtil
+import java.util.function.Supplier
+
+/**
+ * Implementation of [ManageWindowsViewContainer] meant to be used in desktop header and app
+ * handle.
+ */
+class DesktopHeaderManageWindowsMenu(
+    private val callerTaskInfo: RunningTaskInfo,
+    private val displayController: DisplayController,
+    private val rootTdaOrganizer: RootTaskDisplayAreaOrganizer,
+    context: Context,
+    private val surfaceControlBuilderSupplier: Supplier<SurfaceControl.Builder>,
+    private val surfaceControlTransactionSupplier: Supplier<SurfaceControl.Transaction>,
+    snapshotList: List<Pair<Int, TaskSnapshot>>,
+    onIconClickListener: ((Int) -> Unit),
+    onOutsideClickListener: (() -> Unit)
+) : ManageWindowsViewContainer(
+    context,
+    DecorThemeUtil(context).getColorScheme(callerTaskInfo).background.toArgb()
+) {
+    private var menuViewContainer: AdditionalViewContainer? = null
+
+    init {
+        show(snapshotList, onIconClickListener, onOutsideClickListener)
+    }
+
+    override fun close() {
+        menuViewContainer?.releaseView()
+    }
+
+    override fun addToContainer(menuView: ManageWindowsView) {
+        val taskBounds = callerTaskInfo.getConfiguration().windowConfiguration.bounds
+        val menuPosition = Point(taskBounds.left, taskBounds.top)
+        val builder = surfaceControlBuilderSupplier.get()
+        rootTdaOrganizer.attachToDisplayArea(callerTaskInfo.displayId, builder)
+        val leash = builder
+            .setName("Manage Windows Menu")
+            .setContainerLayer()
+            .build()
+        val lp = WindowManager.LayoutParams(
+            menuView.menuWidth, menuView.menuHeight,
+            WindowManager.LayoutParams.TYPE_APPLICATION,
+            WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE
+                    or WindowManager.LayoutParams.FLAG_WATCH_OUTSIDE_TOUCH
+                    or WindowManager.LayoutParams.FLAG_SPLIT_TOUCH,
+            PixelFormat.TRANSPARENT
+        )
+        val windowManager = WindowlessWindowManager(
+            callerTaskInfo.configuration,
+            leash,
+            null // HostInputToken
+        )
+        val viewHost = SurfaceControlViewHost(
+            context,
+            displayController.getDisplay(callerTaskInfo.displayId), windowManager,
+            "MaximizeMenu"
+        )
+        menuView.let { viewHost.setView(it.rootView, lp) }
+        val t = surfaceControlTransactionSupplier.get()
+        t.setLayer(leash, TaskConstants.TASK_CHILD_LAYER_FLOATING_MENU)
+            .setPosition(leash, menuPosition.x.toFloat(), menuPosition.y.toFloat())
+            .show(leash)
+        t.apply()
+        menuViewContainer = AdditionalViewHostViewContainer(
+            leash, viewHost,
+            surfaceControlTransactionSupplier
+        )
+    }
+}
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 7692bd7..2519ce4 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
@@ -39,6 +39,7 @@
 import static com.android.wm.shell.desktopmode.DesktopModeVisualIndicator.IndicatorType.TO_SPLIT_LEFT_INDICATOR;
 import static com.android.wm.shell.desktopmode.DesktopModeVisualIndicator.IndicatorType.TO_SPLIT_RIGHT_INDICATOR;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE;
+import static com.android.wm.shell.shared.desktopmode.ManageWindowsViewContainer.MANAGE_WINDOWS_MINIMUM_INSTANCES;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
@@ -48,6 +49,8 @@
 import android.app.ActivityManager;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.ActivityTaskManager;
+import android.app.IActivityManager;
+import android.app.IActivityTaskManager;
 import android.content.Context;
 import android.content.Intent;
 import android.graphics.Point;
@@ -77,6 +80,7 @@
 import android.view.SurfaceControl.Transaction;
 import android.view.View;
 import android.widget.Toast;
+import android.window.TaskSnapshot;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -125,9 +129,12 @@
 import com.android.wm.shell.windowdecor.viewholder.AppHeaderViewHolder;
 import com.android.wm.shell.windowdecor.viewhost.WindowDecorViewHostSupplier;
 
+import kotlin.Pair;
 import kotlin.Unit;
 
 import java.io.PrintWriter;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
 import java.util.function.Supplier;
 
@@ -586,6 +593,61 @@
         mDesktopTasksController.openNewWindow(decoration.mTaskInfo);
     }
 
+    private void onManageWindows(DesktopModeWindowDecoration decoration) {
+        if (decoration == null) {
+            return;
+        }
+        decoration.closeHandleMenu();
+        decoration.createManageWindowsMenu(getTaskSnapshots(decoration.mTaskInfo),
+                /* onIconClickListener= */(Integer requestedTaskId) -> {
+                    decoration.closeManageWindowsMenu();
+                    mDesktopTasksController.openInstance(decoration.mTaskInfo, requestedTaskId);
+                    return Unit.INSTANCE;
+                });
+    }
+
+    private ArrayList<Pair<Integer, TaskSnapshot>> getTaskSnapshots(
+            @NonNull RunningTaskInfo callerTaskInfo
+    ) {
+        final ArrayList<Pair<Integer, TaskSnapshot>> snapshotList = new ArrayList<>();
+        final IActivityManager activityManager = ActivityManager.getService();
+        final IActivityTaskManager activityTaskManagerService = ActivityTaskManager.getService();
+        final List<ActivityManager.RecentTaskInfo> recentTasks;
+        try {
+            recentTasks = mActivityTaskManager.getRecentTasks(
+                    Integer.MAX_VALUE,
+                    ActivityManager.RECENT_WITH_EXCLUDED,
+                    activityManager.getCurrentUser().id);
+        } catch (RemoteException e) {
+            throw new RuntimeException(e);
+        }
+        final String callerPackageName = callerTaskInfo.baseActivity.getPackageName();
+        for (ActivityManager.RecentTaskInfo info : recentTasks) {
+            if (info.taskId == callerTaskInfo.taskId || info.baseActivity == null) continue;
+            final String infoPackageName = info.baseActivity.getPackageName();
+            if (!infoPackageName.equals(callerPackageName)) {
+                continue;
+            }
+            if (info.baseActivity != null) {
+                if (callerPackageName.equals(infoPackageName)) {
+                    // TODO(b/337903443): Fix this returning null for freeform tasks.
+                    try {
+                        TaskSnapshot screenshot = activityTaskManagerService
+                                .getTaskSnapshot(info.taskId, false);
+                        if (screenshot == null) {
+                            screenshot = activityTaskManagerService
+                                    .takeTaskSnapshot(info.taskId, false);
+                        }
+                        snapshotList.add(new Pair(info.taskId, screenshot));
+                    } catch (RemoteException e) {
+                        throw new RuntimeException(e);
+                    }
+                }
+            }
+        }
+        return snapshotList;
+    }
+
     private class DesktopModeTouchEventListener extends GestureDetector.SimpleOnGestureListener
             implements View.OnClickListener, View.OnTouchListener, View.OnLongClickListener,
             View.OnGenericMotionListener, DragDetector.MotionEventHandler {
@@ -642,7 +704,8 @@
             } else if (id == R.id.caption_handle || id == R.id.open_menu_button) {
                 if (!decoration.isHandleMenuActive()) {
                     moveTaskToFront(decoration.mTaskInfo);
-                    decoration.createHandleMenu();
+                    decoration.createHandleMenu(checkNumberOfOtherInstances(decoration.mTaskInfo)
+                                    >= MANAGE_WINDOWS_MINIMUM_INSTANCES);
                 }
             } else if (id == R.id.maximize_window) {
                 // TODO(b/346441962): move click detection logic into the decor's
@@ -1350,6 +1413,10 @@
             onNewWindow(taskInfo.taskId);
             return Unit.INSTANCE;
         });
+        windowDecoration.setManageWindowsClickListener(() -> {
+            onManageWindows(windowDecoration);
+            return Unit.INSTANCE;
+        });
         windowDecoration.setCaptionListeners(
                 touchEventListener, touchEventListener, touchEventListener, touchEventListener);
         windowDecoration.setExclusionRegionListener(mExclusionRegionListener);
@@ -1441,6 +1508,29 @@
         }
     }
 
+    /**
+     * Gets the number of instances of a task running, not including the specified task itself.
+     */
+    private int checkNumberOfOtherInstances(@NonNull RunningTaskInfo info) {
+        // TODO(b/336289597): Rather than returning number of instances, return a list of valid
+        //  instances, then refer to the list's size and reuse the list for Manage Windows menu.
+        final IActivityTaskManager activityTaskManager = ActivityTaskManager.getService();
+        final IActivityManager activityManager = ActivityManager.getService();
+        try {
+            return activityTaskManager.getRecentTasks(Integer.MAX_VALUE,
+                    ActivityManager.RECENT_WITH_EXCLUDED,
+                    activityManager.getCurrentUserId()).getList().stream().filter(
+                            recentTaskInfo -> (recentTaskInfo.taskId != info.taskId
+                                    && recentTaskInfo.baseActivity != null
+                                    && recentTaskInfo.baseActivity.getPackageName()
+                                    .equals(info.baseActivity.getPackageName())
+                            )
+            ).toList().size();
+        } catch (RemoteException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     static class InputMonitorFactory {
         InputMonitor create(InputManager inputManager, int displayId) {
             return inputManager.monitorGestureInput("caption-touch", displayId);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index 1409d30..5d16d97 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -64,6 +64,7 @@
 import android.view.ViewConfiguration;
 import android.view.WindowManager;
 import android.widget.ImageButton;
+import android.window.TaskSnapshot;
 import android.window.WindowContainerTransaction;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -86,6 +87,7 @@
 import com.android.wm.shell.shared.desktopmode.DesktopModeFlags;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource;
+import com.android.wm.shell.shared.desktopmode.ManageWindowsViewContainer;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.windowdecor.extension.TaskInfoKt;
 import com.android.wm.shell.windowdecor.viewholder.AppHandleViewHolder;
@@ -93,9 +95,12 @@
 import com.android.wm.shell.windowdecor.viewholder.WindowDecorationViewHolder;
 import com.android.wm.shell.windowdecor.viewhost.WindowDecorViewHostSupplier;
 
+import kotlin.Pair;
 import kotlin.Unit;
 import kotlin.jvm.functions.Function0;
+import kotlin.jvm.functions.Function1;
 
+import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
@@ -131,6 +136,7 @@
     private Function0<Unit> mOnToFullscreenClickListener;
     private Function0<Unit> mOnToSplitscreenClickListener;
     private Function0<Unit> mOnNewWindowClickListener;
+    private Function0<Unit> mOnManageWindowsClickListener;
     private DragPositioningCallback mDragPositioningCallback;
     private DragResizeInputListener mDragResizeListener;
     private DragDetector mDragDetector;
@@ -140,6 +146,8 @@
 
     private final Point mPositionInParent = new Point();
     private HandleMenu mHandleMenu;
+    private boolean mMinimumInstancesFound;
+    private ManageWindowsViewContainer mManageWindowsMenu;
 
     private MaximizeMenu mMaximizeMenu;
 
@@ -285,6 +293,14 @@
         mOnNewWindowClickListener = listener;
     }
 
+    /**
+     * Registers a listener to be called when the decoration's manage windows action is
+     * triggered.
+     */
+    void setManageWindowsClickListener(Function0<Unit> listener) {
+        mOnManageWindowsClickListener = listener;
+    }
+
     void setCaptionListeners(
             View.OnClickListener onCaptionButtonClickListener,
             View.OnTouchListener onCaptionTouchListener,
@@ -941,9 +957,10 @@
     /**
      * Updates app info and creates and displays handle menu window.
      */
-    void createHandleMenu() {
+    void createHandleMenu(boolean minimumInstancesFound) {
         // Requests assist content. When content is received, calls {@link #onAssistContentReceived}
         // which sets app info and creates the handle menu.
+        mMinimumInstancesFound = minimumInstancesFound;
         mAssistContentRequester.requestAssistContent(
                 mTaskInfo.taskId, this::onAssistContentReceived);
     }
@@ -956,8 +973,10 @@
         mWebUri = assistContent == null ? null : assistContent.getWebUri();
         loadAppInfoIfNeeded();
         updateGenericLink();
-
-        // Create and display handle menu
+        final boolean supportsMultiInstance = mMultiInstanceHelper
+                .supportsMultiInstanceSplit(mTaskInfo.baseActivity);
+        final boolean shouldShowManageWindowsButton = supportsMultiInstance
+                && mMinimumInstancesFound;
         mHandleMenu = mHandleMenuFactory.create(
                 this,
                 mWindowManagerWrapper,
@@ -966,9 +985,8 @@
                 mAppName,
                 mSplitScreenController,
                 DesktopModeStatus.canEnterDesktopMode(mContext),
-                Flags.enableDesktopWindowingMultiInstanceFeatures()
-                        && mMultiInstanceHelper
-                        .supportsMultiInstanceSplit(mTaskInfo.baseActivity),
+                supportsMultiInstance,
+                shouldShowManageWindowsButton,
                 getBrowserLink(),
                 mResult.mCaptionWidth,
                 mResult.mCaptionHeight,
@@ -984,6 +1002,7 @@
                 /* onToFullscreenClickListener= */ mOnToFullscreenClickListener,
                 /* onToSplitScreenClickListener= */ mOnToSplitscreenClickListener,
                 /* onNewWindowClickListener= */ mOnNewWindowClickListener,
+                /* onManageWindowsClickListener= */ mOnManageWindowsClickListener,
                 /* openInBrowserClickListener= */ (uri) -> {
                     mOpenInBrowserClickListener.accept(uri);
                     onCapturedLinkExpired();
@@ -998,6 +1017,47 @@
                     return Unit.INSTANCE;
                 }
         );
+        mMinimumInstancesFound = false;
+    }
+
+    void createManageWindowsMenu(@NonNull List<Pair<Integer, TaskSnapshot>> snapshotList,
+            @NonNull Function1<Integer, Unit> onIconClickListener
+    ) {
+        if (mTaskInfo.isFreeform()) {
+            mManageWindowsMenu = new DesktopHeaderManageWindowsMenu(
+                    mTaskInfo,
+                    mDisplayController,
+                    mRootTaskDisplayAreaOrganizer,
+                    mContext,
+                    mSurfaceControlBuilderSupplier,
+                    mSurfaceControlTransactionSupplier,
+                    snapshotList,
+                    onIconClickListener,
+                    /* onOutsideClickListener= */ () -> {
+                        closeManageWindowsMenu();
+                        return Unit.INSTANCE;
+                    }
+                    );
+        } else {
+            mManageWindowsMenu = new DesktopHandleManageWindowsMenu(
+                    mTaskInfo,
+                    mSplitScreenController,
+                    getCaptionX(),
+                    mResult.mCaptionWidth,
+                    mWindowManagerWrapper,
+                    mContext,
+                    snapshotList,
+                    onIconClickListener,
+                    /* onOutsideClickListener= */ () -> {
+                        closeManageWindowsMenu();
+                        return Unit.INSTANCE;
+                    }
+                    );
+        }
+    }
+
+    void closeManageWindowsMenu() {
+        mManageWindowsMenu.close();
     }
 
     private void updateGenericLink() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt
index 748046e..3d00a44 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/HandleMenu.kt
@@ -71,6 +71,7 @@
     private val splitScreenController: SplitScreenController,
     private val shouldShowWindowingPill: Boolean,
     private val shouldShowNewWindowButton: Boolean,
+    private val shouldShowManageWindowsButton: Boolean,
     private val openInBrowserLink: Uri?,
     private val captionWidth: Int,
     private val captionHeight: Int,
@@ -119,6 +120,7 @@
         onToFullscreenClickListener: () -> Unit,
         onToSplitScreenClickListener: () -> Unit,
         onNewWindowClickListener: () -> Unit,
+        onManageWindowsClickListener: () -> Unit,
         openInBrowserClickListener: (Uri) -> Unit,
         onCloseMenuClickListener: () -> Unit,
         onOutsideTouchListener: () -> Unit,
@@ -133,6 +135,7 @@
             onToFullscreenClickListener = onToFullscreenClickListener,
             onToSplitScreenClickListener = onToSplitScreenClickListener,
             onNewWindowClickListener = onNewWindowClickListener,
+            onManageWindowsClickListener = onManageWindowsClickListener,
             openInBrowserClickListener = openInBrowserClickListener,
             onCloseMenuClickListener = onCloseMenuClickListener,
             onOutsideTouchListener = onOutsideTouchListener,
@@ -150,6 +153,7 @@
         onToFullscreenClickListener: () -> Unit,
         onToSplitScreenClickListener: () -> Unit,
         onNewWindowClickListener: () -> Unit,
+        onManageWindowsClickListener: () -> Unit,
         openInBrowserClickListener: (Uri) -> Unit,
         onCloseMenuClickListener: () -> Unit,
         onOutsideTouchListener: () -> Unit
@@ -160,13 +164,15 @@
             captionHeight = captionHeight,
             shouldShowWindowingPill = shouldShowWindowingPill,
             shouldShowBrowserPill = shouldShowBrowserPill,
-            shouldShowNewWindowButton = shouldShowNewWindowButton
+            shouldShowNewWindowButton = shouldShowNewWindowButton,
+            shouldShowManageWindowsButton = shouldShowManageWindowsButton
         ).apply {
             bind(taskInfo, appIconBitmap, appName)
             this.onToDesktopClickListener = onToDesktopClickListener
             this.onToFullscreenClickListener = onToFullscreenClickListener
             this.onToSplitScreenClickListener = onToSplitScreenClickListener
             this.onNewWindowClickListener = onNewWindowClickListener
+            this.onManageWindowsClickListener = onManageWindowsClickListener
             this.onOpenInBrowserClickListener = {
                 openInBrowserClickListener.invoke(openInBrowserLink!!)
             }
@@ -372,7 +378,13 @@
                 R.dimen.desktop_mode_handle_menu_new_window_height
             )
         }
-        if (!SHOULD_SHOW_SCREENSHOT_BUTTON && !shouldShowNewWindowButton) {
+        if (!shouldShowManageWindowsButton) {
+            menuHeight -= loadDimensionPixelSize(
+                R.dimen.desktop_mode_handle_menu_manage_windows_height
+            )
+        }
+        if (!SHOULD_SHOW_SCREENSHOT_BUTTON && !shouldShowNewWindowButton
+            && !shouldShowManageWindowsButton) {
             menuHeight -= pillTopMargin
         }
         if (!shouldShowBrowserPill) {
@@ -405,7 +417,8 @@
         captionHeight: Int,
         private val shouldShowWindowingPill: Boolean,
         private val shouldShowBrowserPill: Boolean,
-        private val shouldShowNewWindowButton: Boolean
+        private val shouldShowNewWindowButton: Boolean,
+        private val shouldShowManageWindowsButton: Boolean
     ) {
         val rootView = LayoutInflater.from(context)
             .inflate(R.layout.desktop_mode_window_decor_handle_menu, null /* root */) as View
@@ -430,6 +443,8 @@
         private val moreActionsPill = rootView.requireViewById<View>(R.id.more_actions_pill)
         private val screenshotBtn = moreActionsPill.requireViewById<Button>(R.id.screenshot_button)
         private val newWindowBtn = moreActionsPill.requireViewById<Button>(R.id.new_window_button)
+        private val manageWindowBtn = moreActionsPill
+            .requireViewById<Button>(R.id.manage_windows_button)
 
         // Open in Browser Pill.
         private val openInBrowserPill = rootView.requireViewById<View>(R.id.open_in_browser_pill)
@@ -446,6 +461,7 @@
         var onToFullscreenClickListener: (() -> Unit)? = null
         var onToSplitScreenClickListener: (() -> Unit)? = null
         var onNewWindowClickListener: (() -> Unit)? = null
+        var onManageWindowsClickListener: (() -> Unit)? = null
         var onOpenInBrowserClickListener: (() -> Unit)? = null
         var onCloseMenuClickListener: (() -> Unit)? = null
         var onOutsideTouchListener: (() -> Unit)? = null
@@ -457,6 +473,7 @@
             browserBtn.setOnClickListener { onOpenInBrowserClickListener?.invoke() }
             collapseMenuButton.setOnClickListener { onCloseMenuClickListener?.invoke() }
             newWindowBtn.setOnClickListener { onNewWindowClickListener?.invoke() }
+            manageWindowBtn.setOnClickListener { onManageWindowsClickListener?.invoke() }
 
             rootView.setOnTouchListener { _, event ->
                 if (event.actionMasked == ACTION_OUTSIDE) {
@@ -587,6 +604,7 @@
         private fun bindMoreActionsPill(style: MenuStyle) {
             moreActionsPill.apply {
                 isGone = !shouldShowNewWindowButton && !SHOULD_SHOW_SCREENSHOT_BUTTON
+                        && !shouldShowManageWindowsButton
             }
             screenshotBtn.apply {
                 isGone = !SHOULD_SHOW_SCREENSHOT_BUTTON
@@ -603,6 +621,13 @@
                 setTextColor(style.textColor)
                 compoundDrawableTintList = ColorStateList.valueOf(style.textColor)
             }
+            manageWindowBtn.apply {
+                isGone = !shouldShowManageWindowsButton
+                background.colorFilter =
+                    BlendModeColorFilter(style.backgroundColor, BlendMode.MULTIPLY)
+                setTextColor(style.textColor)
+                compoundDrawableTintList = ColorStateList.valueOf(style.textColor)
+            }
         }
 
         private fun bindOpenInBrowserPill(style: MenuStyle) {
@@ -643,6 +668,7 @@
         splitScreenController: SplitScreenController,
         shouldShowWindowingPill: Boolean,
         shouldShowNewWindowButton: Boolean,
+        shouldShowManageWindowsButton: Boolean,
         openInBrowserLink: Uri?,
         captionWidth: Int,
         captionHeight: Int,
@@ -661,6 +687,7 @@
         splitScreenController: SplitScreenController,
         shouldShowWindowingPill: Boolean,
         shouldShowNewWindowButton: Boolean,
+        shouldShowManageWindowsButton: Boolean,
         openInBrowserLink: Uri?,
         captionWidth: Int,
         captionHeight: Int,
@@ -675,6 +702,7 @@
             splitScreenController,
             shouldShowWindowingPill,
             shouldShowNewWindowButton,
+            shouldShowManageWindowsButton,
             openInBrowserLink,
             captionWidth,
             captionHeight,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
index 5b02837..497d0e5 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
@@ -18,6 +18,8 @@
 import android.window.WindowContainerTransaction
 import androidx.test.filters.SmallTest
 import com.android.dx.mockito.inline.extended.ExtendedMockito
+import com.android.internal.jank.Cuj.CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_HOLD
+import com.android.internal.jank.Cuj.CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_RELEASE
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.ShellTestCase
@@ -448,6 +450,42 @@
         )
     }
 
+    @Test
+    fun startDragToDesktop_aborted_logsDragHoldCancelled() {
+        val transition = startDragToDesktopTransition(defaultHandler, createTask(), dragAnimator)
+
+        defaultHandler.onTransitionConsumed(transition, aborted = true, mock())
+
+        verify(mockInteractionJankMonitor).cancel(eq(CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_HOLD))
+        verify(mockInteractionJankMonitor, times(0)).cancel(
+            eq(CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_RELEASE))
+    }
+
+    @Test
+    fun mergeEndDragToDesktop_aborted_logsDragReleaseCancelled() {
+        val task = createTask()
+        val startTransition = startDrag(defaultHandler, task)
+        val endTransition = mock<IBinder>()
+        defaultHandler.onTaskResizeAnimationListener = mock()
+        defaultHandler.mergeAnimation(
+            transition = endTransition,
+            info = createTransitionInfo(
+                type = TRANSIT_DESKTOP_MODE_END_DRAG_TO_DESKTOP,
+                draggedTask = task
+            ),
+            t = mock<SurfaceControl.Transaction>(),
+            mergeTarget = startTransition,
+            finishCallback = mock<Transitions.TransitionFinishCallback>()
+        )
+
+        defaultHandler.onTransitionConsumed(endTransition, aborted = true, mock())
+
+        verify(mockInteractionJankMonitor)
+            .cancel(eq(CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_RELEASE))
+        verify(mockInteractionJankMonitor, times(0))
+            .cancel(eq(CUJ_DESKTOP_MODE_ENTER_APP_HANDLE_DRAG_HOLD))
+    }
+
     private fun startDrag(
         handler: DragToDesktopTransitionHandler,
         task: RunningTaskInfo = createTask(),
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/DragAndDropPolicyTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/SplitDragPolicyTest.java
similarity index 92%
rename from libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/DragAndDropPolicyTest.java
rename to libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/SplitDragPolicyTest.java
index 645b296..46b60499 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/DragAndDropPolicyTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/SplitDragPolicyTest.java
@@ -30,11 +30,11 @@
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_FULLSCREEN;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_BOTTOM;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_LEFT;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_RIGHT;
-import static com.android.wm.shell.draganddrop.DragAndDropPolicy.Target.TYPE_SPLIT_TOP;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_FULLSCREEN;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_BOTTOM;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_LEFT;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_RIGHT;
+import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_TOP;
 
 import static junit.framework.Assert.assertTrue;
 import static junit.framework.Assert.fail;
@@ -72,7 +72,7 @@
 import com.android.internal.logging.InstanceId;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.DisplayLayout;
-import com.android.wm.shell.draganddrop.DragAndDropPolicy.Target;
+import com.android.wm.shell.draganddrop.SplitDragPolicy.Target;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 
 import org.junit.After;
@@ -92,7 +92,7 @@
  */
 @SmallTest
 @RunWith(AndroidJUnit4.class)
-public class DragAndDropPolicyTest extends ShellTestCase {
+public class SplitDragPolicyTest extends ShellTestCase {
 
     @Mock
     private Context mContext;
@@ -104,7 +104,7 @@
     @Mock
     private SplitScreenController mSplitScreenStarter;
     @Mock
-    private DragAndDropPolicy.Starter mFullscreenStarter;
+    private SplitDragPolicy.Starter mFullscreenStarter;
 
     @Mock
     private InstanceId mLoggerSessionId;
@@ -112,7 +112,7 @@
     private DisplayLayout mLandscapeDisplayLayout;
     private DisplayLayout mPortraitDisplayLayout;
     private Insets mInsets;
-    private DragAndDropPolicy mPolicy;
+    private SplitDragPolicy mPolicy;
 
     private ClipData mActivityClipData;
     private PendingIntent mLaunchableIntentPendingIntent;
@@ -150,7 +150,7 @@
         mPortraitDisplayLayout = new DisplayLayout(info2, res, false, false);
         mInsets = Insets.of(0, 0, 0, 0);
 
-        mPolicy = spy(new DragAndDropPolicy(mContext, mSplitScreenStarter, mFullscreenStarter));
+        mPolicy = spy(new SplitDragPolicy(mContext, mSplitScreenStarter, mFullscreenStarter));
         mActivityClipData = createAppClipData(MIMETYPE_APPLICATION_ACTIVITY);
         mLaunchableIntentPendingIntent = mock(PendingIntent.class);
         when(mLaunchableIntentPendingIntent.getCreatorUserHandle())
@@ -289,7 +289,7 @@
         ArrayList<Target> targets = assertExactTargetTypes(
                 mPolicy.getTargets(mInsets), TYPE_FULLSCREEN);
 
-        mPolicy.handleDrop(filterTargetByType(targets, TYPE_FULLSCREEN), null /* hideTaskToken */);
+        mPolicy.onDropped(filterTargetByType(targets, TYPE_FULLSCREEN), null /* hideTaskToken */);
         verify(mFullscreenStarter).startIntent(any(), anyInt(), any(),
                 eq(SPLIT_POSITION_UNDEFINED), any(), any());
     }
@@ -304,12 +304,12 @@
         ArrayList<Target> targets = assertExactTargetTypes(
                 mPolicy.getTargets(mInsets), TYPE_SPLIT_LEFT, TYPE_SPLIT_RIGHT);
 
-        mPolicy.handleDrop(filterTargetByType(targets, TYPE_SPLIT_LEFT), null /* hideTaskToken */);
+        mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_LEFT), null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
                 eq(SPLIT_POSITION_TOP_OR_LEFT), any(), any());
         reset(mSplitScreenStarter);
 
-        mPolicy.handleDrop(filterTargetByType(targets, TYPE_SPLIT_RIGHT), null /* hideTaskToken */);
+        mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_RIGHT), null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
                 eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), any(), any());
     }
@@ -324,12 +324,12 @@
         ArrayList<Target> targets = assertExactTargetTypes(
                 mPolicy.getTargets(mInsets), TYPE_SPLIT_TOP, TYPE_SPLIT_BOTTOM);
 
-        mPolicy.handleDrop(filterTargetByType(targets, TYPE_SPLIT_TOP), null /* hideTaskToken */);
+        mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_TOP), null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
                 eq(SPLIT_POSITION_TOP_OR_LEFT), any(), any());
         reset(mSplitScreenStarter);
 
-        mPolicy.handleDrop(filterTargetByType(targets, TYPE_SPLIT_BOTTOM),
+        mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_BOTTOM),
                 null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
                 eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), any(), any());
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
index 6ddb678..f3944d5 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipControllerTest.java
@@ -256,7 +256,7 @@
         when(mMockPipDisplayLayoutState.getDisplayLayout()).thenReturn(mMockDisplayLayout1);
         when(mMockDisplayController.getDisplayLayout(displayId)).thenReturn(mMockDisplayLayout2);
 
-        when(mMockPipTaskOrganizer.isInPip()).thenReturn(true);
+        when(mMockPipTransitionState.hasEnteredPip()).thenReturn(true);
         mPipController.mDisplaysChangedListener.onDisplayConfigurationChanged(
                 displayId, new Configuration());
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
index 7b68ddf..1741fe4 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorationTests.java
@@ -228,6 +228,8 @@
         mTestableContext = new TestableContext(mContext);
         mTestableContext.ensureTestableResources();
         mContext.setMockPackageManager(mMockPackageManager);
+        when(mMockMultiInstanceHelper.supportsMultiInstanceSplit(any()))
+                .thenReturn(false);
         when(mMockPackageManager.getApplicationLabel(any())).thenReturn("applicationLabel");
         final ActivityInfo activityInfo = new ActivityInfo();
         activityInfo.applicationInfo = new ApplicationInfo();
@@ -235,8 +237,8 @@
         final Display defaultDisplay = mock(Display.class);
         doReturn(defaultDisplay).when(mMockDisplayController).getDisplay(Display.DEFAULT_DISPLAY);
         doReturn(mInsetsState).when(mMockDisplayController).getInsetsState(anyInt());
-        when(mMockHandleMenuFactory.create(any(), any(), anyInt(), any(), any(),
-                any(), anyBoolean(), anyBoolean(), any(), anyInt(), anyInt(), anyInt()))
+        when(mMockHandleMenuFactory.create(any(), any(), anyInt(), any(), any(), any(),
+                anyBoolean(), anyBoolean(), anyBoolean(), any(), anyInt(), anyInt(), anyInt()))
                 .thenReturn(mMockHandleMenu);
         when(mMockMultiInstanceHelper.supportsMultiInstanceSplit(any())).thenReturn(false);
         when(mMockWindowDecorViewHostSupplier.acquire(any(), eq(defaultDisplay)))
@@ -744,6 +746,7 @@
                 any(),
                 any(),
                 any(),
+                any(),
                 openInBrowserCaptor.capture(),
                 any(),
                 any()
@@ -771,6 +774,7 @@
                 any(),
                 any(),
                 any(),
+                any(),
                 openInBrowserCaptor.capture(),
                 any(),
                 any()
@@ -821,6 +825,7 @@
                 any(),
                 any(),
                 any(),
+                any(),
                 closeClickListener.capture(),
                 any()
         );
@@ -832,8 +837,10 @@
     }
 
     private void verifyHandleMenuCreated(@Nullable Uri uri) {
+
         verify(mMockHandleMenuFactory).create(any(), any(), anyInt(), any(), any(),
-                any(), anyBoolean(), anyBoolean(), eq(uri), anyInt(), anyInt(), anyInt());
+                any(), anyBoolean(), anyBoolean(), anyBoolean(), eq(uri), anyInt(),
+                anyInt(), anyInt());
     }
 
     private void createMaximizeMenu(DesktopModeWindowDecoration decoration, MaximizeMenu menu) {
@@ -932,7 +939,7 @@
     }
 
     private void createHandleMenu(@NonNull DesktopModeWindowDecoration decor) {
-        decor.createHandleMenu();
+        decor.createHandleMenu(false);
         // Call DesktopModeWindowDecoration#onAssistContentReceived because decor waits to receive
         // {@link AssistContent} before creating the menu
         decor.onAssistContentReceived(mAssistContent);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt
index 100abbb..a845231 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/HandleMenuTest.kt
@@ -236,11 +236,11 @@
         val handleMenu = HandleMenu(mockDesktopWindowDecoration,
             WindowManagerWrapper(mockWindowManager),
             layoutId, appIcon, appName, splitScreenController, shouldShowWindowingPill = true,
-            shouldShowNewWindowButton = true,
+            shouldShowNewWindowButton = true, shouldShowManageWindowsButton = false,
             null /* openInBrowserLink */, captionWidth = HANDLE_WIDTH, captionHeight = 50,
             captionX = captionX
         )
-        handleMenu.show(mock(), mock(), mock(), mock(), mock(), mock(), mock())
+        handleMenu.show(mock(), mock(), mock(), mock(), mock(), mock(), mock(), mock())
         return handleMenu
     }
 
diff --git a/libs/hwui/apex/android_bitmap.cpp b/libs/hwui/apex/android_bitmap.cpp
index c80a9b4..000f109 100644
--- a/libs/hwui/apex/android_bitmap.cpp
+++ b/libs/hwui/apex/android_bitmap.cpp
@@ -14,21 +14,21 @@
  * limitations under the License.
  */
 
-#include <log/log.h>
-
-#include "android/graphics/bitmap.h"
-#include "TypeCast.h"
-#include "GraphicsJNI.h"
-
+#include <Gainmap.h>
 #include <GraphicsJNI.h>
-#include <hwui/Bitmap.h>
 #include <SkBitmap.h>
 #include <SkColorSpace.h>
 #include <SkImageInfo.h>
 #include <SkRefCnt.h>
 #include <SkStream.h>
+#include <hwui/Bitmap.h>
+#include <log/log.h>
 #include <utils/Color.h>
 
+#include "GraphicsJNI.h"
+#include "TypeCast.h"
+#include "android/graphics/bitmap.h"
+
 using namespace android;
 
 ABitmap* ABitmap_acquireBitmapFromJava(JNIEnv* env, jobject bitmapObj) {
@@ -215,6 +215,14 @@
 int ABitmap_compress(const AndroidBitmapInfo* info, ADataSpace dataSpace, const void* pixels,
                      AndroidBitmapCompressFormat inFormat, int32_t quality, void* userContext,
                      AndroidBitmap_CompressWriteFunc fn) {
+    return ABitmap_compressWithGainmap(info, dataSpace, pixels, nullptr, -1.f, inFormat, quality,
+                                       userContext, fn);
+}
+
+int ABitmap_compressWithGainmap(const AndroidBitmapInfo* info, ADataSpace dataSpace,
+                                const void* pixels, const void* gainmapPixels, float hdrSdrRatio,
+                                AndroidBitmapCompressFormat inFormat, int32_t quality,
+                                void* userContext, AndroidBitmap_CompressWriteFunc fn) {
     Bitmap::JavaCompressFormat format;
     switch (inFormat) {
         case ANDROID_BITMAP_COMPRESS_FORMAT_JPEG:
@@ -275,7 +283,7 @@
         // besides ADATASPACE_UNKNOWN as an error?
         cs = nullptr;
     } else {
-        cs = uirenderer::DataSpaceToColorSpace((android_dataspace) dataSpace);
+        cs = uirenderer::DataSpaceToColorSpace((android_dataspace)dataSpace);
         // DataSpaceToColorSpace treats UNKNOWN as SRGB, but compress forces the
         // client to specify SRGB if that is what they want.
         if (!cs || dataSpace == ADATASPACE_UNKNOWN) {
@@ -292,16 +300,70 @@
 
     auto imageInfo =
             SkImageInfo::Make(info->width, info->height, colorType, alphaType, std::move(cs));
-    SkBitmap bitmap;
-    // We are not going to modify the pixels, but installPixels expects them to
-    // not be const, since for all it knows we might want to draw to the SkBitmap.
-    if (!bitmap.installPixels(imageInfo, const_cast<void*>(pixels), info->stride)) {
-        return ANDROID_BITMAP_RESULT_BAD_PARAMETER;
+
+    // Validate the image info
+    {
+        SkBitmap tempBitmap;
+        if (!tempBitmap.installPixels(imageInfo, const_cast<void*>(pixels), info->stride)) {
+            return ANDROID_BITMAP_RESULT_BAD_PARAMETER;
+        }
+    }
+
+    SkPixelRef pixelRef =
+            SkPixelRef(info->width, info->height, const_cast<void*>(pixels), info->stride);
+
+    auto bitmap = Bitmap::createFrom(imageInfo, pixelRef);
+
+    if (gainmapPixels) {
+        auto gainmap = sp<uirenderer::Gainmap>::make();
+        gainmap->info.fGainmapRatioMin = {
+                1.f,
+                1.f,
+                1.f,
+                1.f,
+        };
+        gainmap->info.fGainmapRatioMax = {
+                hdrSdrRatio,
+                hdrSdrRatio,
+                hdrSdrRatio,
+                1.f,
+        };
+        gainmap->info.fGainmapGamma = {
+                1.f,
+                1.f,
+                1.f,
+                1.f,
+        };
+
+        static constexpr auto kDefaultEpsilon = 1.f / 64.f;
+        gainmap->info.fEpsilonSdr = {
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                1.f,
+        };
+        gainmap->info.fEpsilonHdr = {
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                1.f,
+        };
+        gainmap->info.fDisplayRatioSdr = 1.f;
+        gainmap->info.fDisplayRatioHdr = hdrSdrRatio;
+
+        SkPixelRef gainmapPixelRef = SkPixelRef(info->width, info->height,
+                                                const_cast<void*>(gainmapPixels), info->stride);
+        auto gainmapBitmap = Bitmap::createFrom(imageInfo, gainmapPixelRef);
+        gainmap->bitmap = std::move(gainmapBitmap);
+        bitmap->setGainmap(std::move(gainmap));
     }
 
     CompressWriter stream(userContext, fn);
-    return Bitmap::compress(bitmap, format, quality, &stream) ? ANDROID_BITMAP_RESULT_SUCCESS
-                                                              : ANDROID_BITMAP_RESULT_JNI_EXCEPTION;
+
+    return bitmap->compress(format, quality, &stream)
+
+                   ? ANDROID_BITMAP_RESULT_SUCCESS
+                   : ANDROID_BITMAP_RESULT_JNI_EXCEPTION;
 }
 
 AHardwareBuffer* ABitmap_getHardwareBuffer(ABitmap* bitmapHandle) {
diff --git a/libs/hwui/apex/include/android/graphics/bitmap.h b/libs/hwui/apex/include/android/graphics/bitmap.h
index 8c4b439d..6f65e9e 100644
--- a/libs/hwui/apex/include/android/graphics/bitmap.h
+++ b/libs/hwui/apex/include/android/graphics/bitmap.h
@@ -65,6 +65,13 @@
 ANDROID_API int ABitmap_compress(const AndroidBitmapInfo* info, ADataSpace dataSpace, const void* pixels,
                      AndroidBitmapCompressFormat format, int32_t quality, void* userContext,
                      AndroidBitmap_CompressWriteFunc);
+// If gainmapPixels is null, then no gainmap is encoded, and hdrSdrRatio is
+// unused
+ANDROID_API int ABitmap_compressWithGainmap(const AndroidBitmapInfo* info, ADataSpace dataSpace,
+                                            const void* pixels, const void* gainmapPixels,
+                                            float hdrSdrRatio, AndroidBitmapCompressFormat format,
+                                            int32_t quality, void* userContext,
+                                            AndroidBitmap_CompressWriteFunc);
 /**
  *  Retrieve the native object associated with a HARDWARE Bitmap.
  *
diff --git a/libs/hwui/libhwui.map.txt b/libs/hwui/libhwui.map.txt
index d03ceb4..2414299 100644
--- a/libs/hwui/libhwui.map.txt
+++ b/libs/hwui/libhwui.map.txt
@@ -13,6 +13,7 @@
     ABitmapConfig_getFormatFromConfig;
     ABitmapConfig_getConfigFromFormat;
     ABitmap_compress;
+    ABitmap_compressWithGainmap;
     ABitmap_getHardwareBuffer;
     ACanvas_isSupportedPixelFormat;
     ACanvas_getNativeHandleFromJava;
diff --git a/packages/SettingsLib/Android.bp b/packages/SettingsLib/Android.bp
index b997c35..0cb85d8 100644
--- a/packages/SettingsLib/Android.bp
+++ b/packages/SettingsLib/Android.bp
@@ -42,8 +42,6 @@
         "SettingsLibIllustrationPreference",
         "SettingsLibLayoutPreference",
         "SettingsLibMainSwitchPreference",
-        "SettingsLibMetadata",
-        "SettingsLibPreference",
         "SettingsLibProfileSelector",
         "SettingsLibProgressBar",
         "SettingsLibRestrictedLockUtils",
diff --git a/packages/SettingsLib/src/com/android/settingslib/core/lifecycle/ObservablePreferenceFragment.java b/packages/SettingsLib/src/com/android/settingslib/core/lifecycle/ObservablePreferenceFragment.java
deleted file mode 100644
index 7994924..0000000
--- a/packages/SettingsLib/src/com/android/settingslib/core/lifecycle/ObservablePreferenceFragment.java
+++ /dev/null
@@ -1,136 +0,0 @@
-/*
- * Copyright (C) 2016 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.core.lifecycle;
-
-
-import static androidx.lifecycle.Lifecycle.Event.ON_CREATE;
-import static androidx.lifecycle.Lifecycle.Event.ON_DESTROY;
-import static androidx.lifecycle.Lifecycle.Event.ON_PAUSE;
-import static androidx.lifecycle.Lifecycle.Event.ON_RESUME;
-import static androidx.lifecycle.Lifecycle.Event.ON_START;
-import static androidx.lifecycle.Lifecycle.Event.ON_STOP;
-
-import android.annotation.CallSuper;
-import android.content.Context;
-import android.os.Bundle;
-import android.view.Menu;
-import android.view.MenuInflater;
-import android.view.MenuItem;
-
-import androidx.lifecycle.LifecycleOwner;
-import androidx.preference.PreferenceScreen;
-
-import com.android.settingslib.preference.PreferenceFragment;
-
-/**
- * Preference fragment that has hooks to observe fragment lifecycle events.
- */
-public abstract class ObservablePreferenceFragment extends PreferenceFragment
-        implements LifecycleOwner {
-
-    private final Lifecycle mLifecycle = new Lifecycle(this);
-
-    public Lifecycle getSettingsLifecycle() {
-        return mLifecycle;
-    }
-
-    @CallSuper
-    @Override
-    public void onAttach(Context context) {
-        super.onAttach(context);
-        mLifecycle.onAttach(context);
-    }
-
-    @CallSuper
-    @Override
-    public void onCreate(Bundle savedInstanceState) {
-        mLifecycle.onCreate(savedInstanceState);
-        mLifecycle.handleLifecycleEvent(ON_CREATE);
-        super.onCreate(savedInstanceState);
-    }
-
-    @Override
-    public void setPreferenceScreen(PreferenceScreen preferenceScreen) {
-        mLifecycle.setPreferenceScreen(preferenceScreen);
-        super.setPreferenceScreen(preferenceScreen);
-    }
-
-    @CallSuper
-    @Override
-    public void onSaveInstanceState(Bundle outState) {
-        super.onSaveInstanceState(outState);
-        mLifecycle.onSaveInstanceState(outState);
-    }
-
-    @CallSuper
-    @Override
-    public void onStart() {
-        mLifecycle.handleLifecycleEvent(ON_START);
-        super.onStart();
-    }
-
-    @CallSuper
-    @Override
-    public void onResume() {
-        mLifecycle.handleLifecycleEvent(ON_RESUME);
-        super.onResume();
-    }
-
-    @CallSuper
-    @Override
-    public void onPause() {
-        mLifecycle.handleLifecycleEvent(ON_PAUSE);
-        super.onPause();
-    }
-
-    @CallSuper
-    @Override
-    public void onStop() {
-        mLifecycle.handleLifecycleEvent(ON_STOP);
-        super.onStop();
-    }
-
-    @CallSuper
-    @Override
-    public void onDestroy() {
-        mLifecycle.handleLifecycleEvent(ON_DESTROY);
-        super.onDestroy();
-    }
-
-    @CallSuper
-    @Override
-    public void onCreateOptionsMenu(final Menu menu, final MenuInflater inflater) {
-        mLifecycle.onCreateOptionsMenu(menu, inflater);
-        super.onCreateOptionsMenu(menu, inflater);
-    }
-
-    @CallSuper
-    @Override
-    public void onPrepareOptionsMenu(final Menu menu) {
-        mLifecycle.onPrepareOptionsMenu(menu);
-        super.onPrepareOptionsMenu(menu);
-    }
-
-    @CallSuper
-    @Override
-    public boolean onOptionsItemSelected(final MenuItem menuItem) {
-        boolean lifecycleHandled = mLifecycle.onOptionsItemSelected(menuItem);
-        if (!lifecycleHandled) {
-            return super.onOptionsItemSelected(menuItem);
-        }
-        return lifecycleHandled;
-    }
-}
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
index 7b2a284..3cc111f 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
@@ -140,7 +140,7 @@
 
     private static Status computeStatus(@NonNull ZenModeConfig.ZenRule zenRuleExtraData) {
         if (zenRuleExtraData.enabled) {
-            if (zenRuleExtraData.isAutomaticActive()) {
+            if (zenRuleExtraData.isActive()) {
                 return Status.ENABLED_AND_ACTIVE;
             } else {
                 return Status.ENABLED;
@@ -241,10 +241,6 @@
                         formattedTime);
             }
         }
-        // TODO: b/333527800 - For TYPE_SCHEDULE_TIME rules we could do the same; however
-        //   according to the snoozing discussions the mode may or may not end at the scheduled
-        //   time if manually activated. When we resolve that point, we could calculate end time
-        //   for these modes as well.
 
         return getTriggerDescription();
     }
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index be4e9a1..f59eab0 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -934,9 +934,9 @@
         "androidx.compose.runtime_runtime",
     ],
     libs: [
-        "android.test.runner",
-        "android.test.base",
-        "android.test.mock",
+        "android.test.runner.stubs.system",
+        "android.test.base.stubs.system",
+        "android.test.mock.stubs.system",
         "truth",
     ],
 
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeSceneModule.kt
deleted file mode 100644
index c58df35..0000000
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeSceneModule.kt
+++ /dev/null
@@ -1,29 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.scene
-
-import com.android.systemui.notifications.ui.composable.NotificationsShadeScene
-import com.android.systemui.scene.ui.composable.Scene
-import dagger.Binds
-import dagger.Module
-import dagger.multibindings.IntoSet
-
-@Module
-interface NotificationsShadeSceneModule {
-
-    @Binds @IntoSet fun notificationsShade(scene: NotificationsShadeScene): Scene
-}
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeSceneModule.kt
deleted file mode 100644
index 5bb6ae4..0000000
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeSceneModule.kt
+++ /dev/null
@@ -1,29 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.scene
-
-import com.android.systemui.qs.ui.composable.QuickSettingsShadeScene
-import com.android.systemui.scene.ui.composable.Scene
-import dagger.Binds
-import dagger.Module
-import dagger.multibindings.IntoSet
-
-@Module
-interface QuickSettingsShadeSceneModule {
-
-    @Binds @IntoSet fun quickSettingsShade(scene: QuickSettingsShadeScene): Scene
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
index a4dc8fc..557257d 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalContainer.kt
@@ -72,7 +72,7 @@
     override fun matches(key: ElementKey, content: ContentKey) = true
 }
 
-private object TransitionDuration {
+object TransitionDuration {
     const val BETWEEN_HUB_AND_EDIT_MODE_MS = 1000
     const val EDIT_MODE_TO_HUB_CONTENT_MS = 167
     const val EDIT_MODE_TO_HUB_GRID_DELAY_MS = 167
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
index efe0f2e..f4d1242 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
@@ -108,6 +108,8 @@
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.drawBehind
+import androidx.compose.ui.focus.FocusRequester
+import androidx.compose.ui.focus.focusRequester
 import androidx.compose.ui.geometry.CornerRadius
 import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.geometry.Size
@@ -179,6 +181,7 @@
 import com.android.systemui.communal.widgets.WidgetConfigurator
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.phone.SystemUIDialogFactory
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
 
 @OptIn(ExperimentalMaterial3Api::class)
@@ -1155,6 +1158,15 @@
     val selectedIndex =
         selectedKey?.let { key -> contentListState.list.indexOfFirst { it.key == key } }
 
+    val interactionSource = remember { MutableInteractionSource() }
+    val focusRequester = remember { FocusRequester() }
+    if (viewModel.isEditMode && selected) {
+        LaunchedEffect(Unit) {
+            delay(TransitionDuration.BETWEEN_HUB_AND_EDIT_MODE_MS.toLong())
+            focusRequester.requestFocus()
+        }
+    }
+
     val isSelected = selectedKey == model.key
 
     val selectableModifier =
@@ -1162,7 +1174,7 @@
             Modifier.selectable(
                 selected = isSelected,
                 onClick = { viewModel.setSelectedKey(model.key) },
-                interactionSource = remember { MutableInteractionSource() },
+                interactionSource = interactionSource,
                 indication = null,
             )
         } else {
@@ -1172,6 +1184,8 @@
     Box(
         modifier =
             modifier
+                .focusRequester(focusRequester)
+                .focusable(interactionSource = interactionSource)
                 .then(selectableModifier)
                 .thenIf(!viewModel.isEditMode && !model.inQuietMode) {
                     Modifier.pointerInput(Unit) {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeScene.kt
deleted file mode 100644
index 1f4cd04..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeScene.kt
+++ /dev/null
@@ -1,116 +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.notifications.ui.composable
-
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.fillMaxWidth
-import androidx.compose.foundation.layout.padding
-import androidx.compose.runtime.Composable
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.unit.dp
-import com.android.compose.animation.scene.SceneScope
-import com.android.compose.animation.scene.UserAction
-import com.android.compose.animation.scene.UserActionResult
-import com.android.systemui.battery.BatteryMeterViewController
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.rememberViewModel
-import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeUserActionsViewModel
-import com.android.systemui.scene.session.ui.composable.SaveableSession
-import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.Scene
-import com.android.systemui.shade.shared.model.ShadeMode
-import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
-import com.android.systemui.shade.ui.composable.OverlayShade
-import com.android.systemui.shade.ui.viewmodel.ShadeHeaderViewModel
-import com.android.systemui.statusbar.notification.stack.ui.view.NotificationScrollView
-import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
-import com.android.systemui.statusbar.phone.ui.StatusBarIconController
-import com.android.systemui.statusbar.phone.ui.TintedIconManager
-import dagger.Lazy
-import javax.inject.Inject
-import kotlinx.coroutines.flow.Flow
-
-@SysUISingleton
-class NotificationsShadeScene
-@Inject
-constructor(
-    private val actionsViewModelFactory: NotificationsShadeUserActionsViewModel.Factory,
-    private val shadeHeaderViewModelFactory: ShadeHeaderViewModel.Factory,
-    private val notificationsPlaceholderViewModelFactory: NotificationsPlaceholderViewModel.Factory,
-    private val tintedIconManagerFactory: TintedIconManager.Factory,
-    private val batteryMeterViewControllerFactory: BatteryMeterViewController.Factory,
-    private val statusBarIconController: StatusBarIconController,
-    private val shadeSession: SaveableSession,
-    private val stackScrollView: Lazy<NotificationScrollView>,
-) : ExclusiveActivatable(), Scene {
-
-    override val key = Scenes.NotificationsShade
-
-    private val actionsViewModel: NotificationsShadeUserActionsViewModel by lazy {
-        actionsViewModelFactory.create()
-    }
-
-    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
-
-    override suspend fun onActivated(): Nothing {
-        actionsViewModel.activate()
-    }
-
-    @Composable
-    override fun SceneScope.Content(
-        modifier: Modifier,
-    ) {
-        val notificationsPlaceholderViewModel =
-            rememberViewModel("NotificationsShadeScene") {
-                notificationsPlaceholderViewModelFactory.create()
-            }
-
-        OverlayShade(
-            modifier = modifier,
-            onScrimClicked = {},
-        ) {
-            Column {
-                ExpandedShadeHeader(
-                    viewModelFactory = shadeHeaderViewModelFactory,
-                    createTintedIconManager = tintedIconManagerFactory::create,
-                    createBatteryMeterViewController = batteryMeterViewControllerFactory::create,
-                    statusBarIconController = statusBarIconController,
-                    modifier = Modifier.padding(horizontal = 16.dp),
-                )
-
-                NotificationScrollingStack(
-                    shadeSession = shadeSession,
-                    stackScrollView = stackScrollView.get(),
-                    viewModel = notificationsPlaceholderViewModel,
-                    maxScrimTop = { 0f },
-                    shouldPunchHoleBehindScrim = false,
-                    shouldFillMaxSize = false,
-                    shouldReserveSpaceForNavBar = false,
-                    shadeMode = ShadeMode.Dual,
-                    modifier = Modifier.fillMaxWidth(),
-                )
-
-                // Communicates the bottom position of the drawable area within the shade to NSSL.
-                NotificationStackCutoffGuideline(
-                    stackScrollView = stackScrollView.get(),
-                    viewModel = notificationsPlaceholderViewModel,
-                )
-            }
-        }
-    }
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeScene.kt
deleted file mode 100644
index e27c7e2..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeScene.kt
+++ /dev/null
@@ -1,92 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.qs.ui.composable
-
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.padding
-import androidx.compose.runtime.Composable
-import androidx.compose.ui.Modifier
-import com.android.compose.animation.scene.SceneScope
-import com.android.compose.animation.scene.UserAction
-import com.android.compose.animation.scene.UserActionResult
-import com.android.systemui.battery.BatteryMeterViewController
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.rememberViewModel
-import com.android.systemui.qs.ui.viewmodel.QuickSettingsShadeSceneContentViewModel
-import com.android.systemui.qs.ui.viewmodel.QuickSettingsShadeUserActionsViewModel
-import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.Scene
-import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
-import com.android.systemui.shade.ui.composable.OverlayShade
-import com.android.systemui.shade.ui.viewmodel.ShadeHeaderViewModel
-import com.android.systemui.statusbar.phone.ui.StatusBarIconController
-import com.android.systemui.statusbar.phone.ui.TintedIconManager
-import javax.inject.Inject
-import kotlinx.coroutines.flow.Flow
-
-@SysUISingleton
-class QuickSettingsShadeScene
-@Inject
-constructor(
-    private val actionsViewModelFactory: QuickSettingsShadeUserActionsViewModel.Factory,
-    private val contentViewModelFactory: QuickSettingsShadeSceneContentViewModel.Factory,
-    private val shadeHeaderViewModelFactory: ShadeHeaderViewModel.Factory,
-    private val tintedIconManagerFactory: TintedIconManager.Factory,
-    private val batteryMeterViewControllerFactory: BatteryMeterViewController.Factory,
-    private val statusBarIconController: StatusBarIconController,
-) : ExclusiveActivatable(), Scene {
-
-    override val key = Scenes.QuickSettingsShade
-
-    private val actionsViewModel: QuickSettingsShadeUserActionsViewModel by lazy {
-        actionsViewModelFactory.create()
-    }
-
-    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
-
-    override suspend fun onActivated(): Nothing {
-        actionsViewModel.activate()
-    }
-
-    @Composable
-    override fun SceneScope.Content(
-        modifier: Modifier,
-    ) {
-        val viewModel =
-            rememberViewModel("QuickSettingsShadeScene") { contentViewModelFactory.create() }
-
-        OverlayShade(
-            modifier = modifier,
-            onScrimClicked = {},
-        ) {
-            Column {
-                ExpandedShadeHeader(
-                    viewModelFactory = shadeHeaderViewModelFactory,
-                    createTintedIconManager = tintedIconManagerFactory::create,
-                    createBatteryMeterViewController = batteryMeterViewControllerFactory::create,
-                    statusBarIconController = statusBarIconController,
-                    modifier = Modifier.padding(QuickSettingsShade.Dimensions.Padding),
-                )
-
-                ShadeBody(
-                    viewModel = viewModel.quickSettingsContainerViewModel,
-                )
-            }
-        }
-    }
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
index f660808..f64d0ed 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
@@ -6,6 +6,7 @@
 import com.android.compose.animation.scene.transitions
 import com.android.systemui.bouncer.ui.composable.Bouncer
 import com.android.systemui.notifications.ui.composable.Notifications
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.TransitionKeys.SlightlyFasterShadeCollapse
 import com.android.systemui.scene.shared.model.TransitionKeys.ToSplitShade
@@ -20,7 +21,11 @@
 import com.android.systemui.scene.ui.composable.transitions.lockscreenToQuickSettingsTransition
 import com.android.systemui.scene.ui.composable.transitions.lockscreenToShadeTransition
 import com.android.systemui.scene.ui.composable.transitions.lockscreenToSplitShadeTransition
+import com.android.systemui.scene.ui.composable.transitions.notificationsShadeToQuickSettingsShadeTransition
 import com.android.systemui.scene.ui.composable.transitions.shadeToQuickSettingsTransition
+import com.android.systemui.scene.ui.composable.transitions.toNotificationsShadeTransition
+import com.android.systemui.scene.ui.composable.transitions.toQuickSettingsShadeTransition
+import com.android.systemui.shade.ui.composable.OverlayShade
 import com.android.systemui.shade.ui.composable.Shade
 
 /**
@@ -74,6 +79,14 @@
     from(Scenes.Lockscreen, to = Scenes.Gone) { lockscreenToGoneTransition() }
     from(Scenes.Shade, to = Scenes.QuickSettings) { shadeToQuickSettingsTransition() }
 
+    // Overlay transitions
+
+    to(Overlays.NotificationsShade) { toNotificationsShadeTransition() }
+    to(Overlays.QuickSettingsShade) { toQuickSettingsShadeTransition() }
+    from(Overlays.NotificationsShade, Overlays.QuickSettingsShade) {
+        notificationsShadeToQuickSettingsShadeTransition()
+    }
+
     // Scene overscroll
 
     overscrollDisabled(Scenes.Gone, Orientation.Vertical)
@@ -91,4 +104,10 @@
             y = Shade.Dimensions.ScrimOverscrollLimit,
         )
     }
+    overscroll(Overlays.NotificationsShade, Orientation.Vertical) {
+        translate(OverlayShade.Elements.Panel, y = OverlayShade.Dimensions.OverscrollLimit)
+    }
+    overscroll(Overlays.QuickSettingsShade, Orientation.Vertical) {
+        translate(OverlayShade.Elements.Panel, y = OverlayShade.Dimensions.OverscrollLimit)
+    }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToQuickSettingsShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToQuickSettingsShadeTransition.kt
deleted file mode 100644
index 8a03e29..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToQuickSettingsShadeTransition.kt
+++ /dev/null
@@ -1,27 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.scene.ui.composable.transitions
-
-import com.android.compose.animation.scene.Edge
-import com.android.compose.animation.scene.TransitionBuilder
-
-fun TransitionBuilder.goneToQuickSettingsShadeTransition(
-    edge: Edge = Edge.Top,
-    durationScale: Double = 1.0,
-) {
-    toQuickSettingsShadeTransition(edge, durationScale)
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToNotificationsShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToNotificationsShadeTransition.kt
deleted file mode 100644
index 02664c1..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToNotificationsShadeTransition.kt
+++ /dev/null
@@ -1,25 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.scene.ui.composable.transitions
-
-import com.android.compose.animation.scene.TransitionBuilder
-
-fun TransitionBuilder.lockscreenToNotificationsShadeTransition(
-    durationScale: Double = 1.0,
-) {
-    toNotificationsShadeTransition(durationScale = durationScale)
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToQuickSettingsShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToQuickSettingsShadeTransition.kt
deleted file mode 100644
index 19aa3a7..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToQuickSettingsShadeTransition.kt
+++ /dev/null
@@ -1,26 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.scene.ui.composable.transitions
-
-import com.android.compose.animation.scene.Edge
-import com.android.compose.animation.scene.TransitionBuilder
-
-fun TransitionBuilder.lockscreenToQuickSettingsShadeTransition(
-    durationScale: Double = 1.0,
-) {
-    toQuickSettingsShadeTransition(Edge.Top, durationScale)
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromNotificationsShadeToQuickSettingsShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromNotificationsShadeToQuickSettingsShadeTransition.kt
new file mode 100644
index 0000000..24f285e
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromNotificationsShadeToQuickSettingsShadeTransition.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.scene.ui.composable.transitions
+
+import androidx.compose.animation.core.Spring
+import androidx.compose.animation.core.spring
+import androidx.compose.animation.core.tween
+import com.android.compose.animation.scene.TransitionBuilder
+import com.android.systemui.shade.ui.composable.Shade
+import kotlin.time.Duration.Companion.milliseconds
+
+fun TransitionBuilder.notificationsShadeToQuickSettingsShadeTransition(
+    durationScale: Double = 1.0
+) {
+    spec = tween(durationMillis = (DefaultDuration * durationScale).inWholeMilliseconds.toInt())
+    swipeSpec =
+        spring(
+            stiffness = Spring.StiffnessMediumLow,
+            visibilityThreshold = Shade.Dimensions.ScrimVisibilityThreshold,
+        )
+}
+
+private val DefaultDuration = 300.milliseconds
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToQuickSettingsShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToQuickSettingsShadeTransition.kt
index 9d13647..55fa6ad 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToQuickSettingsShadeTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToQuickSettingsShadeTransition.kt
@@ -26,10 +26,7 @@
 import com.android.systemui.shade.ui.composable.Shade
 import kotlin.time.Duration.Companion.milliseconds
 
-fun TransitionBuilder.toQuickSettingsShadeTransition(
-    edge: Edge = Edge.Top,
-    durationScale: Double = 1.0,
-) {
+fun TransitionBuilder.toQuickSettingsShadeTransition(durationScale: Double = 1.0) {
     spec = tween(durationMillis = (DefaultDuration * durationScale).inWholeMilliseconds.toInt())
     swipeSpec =
         spring(
@@ -38,7 +35,7 @@
         )
     distance = UserActionDistance { fromSceneSize, _ -> fromSceneSize.height.toFloat() * 2 / 3f }
 
-    translate(OverlayShade.Elements.Panel, edge)
+    translate(OverlayShade.Elements.Panel, Edge.Top)
 
     fractionRange(end = .5f) { fade(OverlayShade.Elements.Scrim) }
 }
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 df22264..8a59e20 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
@@ -16,6 +16,7 @@
 
 package com.android.systemui.shade.ui.composable
 
+import android.view.HapticFeedbackConstants
 import android.view.ViewGroup
 import androidx.compose.animation.core.animateDpAsState
 import androidx.compose.animation.core.animateFloatAsState
@@ -60,6 +61,7 @@
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalDensity
 import androidx.compose.ui.platform.LocalLifecycleOwner
+import androidx.compose.ui.platform.LocalView
 import androidx.compose.ui.res.colorResource
 import androidx.compose.ui.res.dimensionResource
 import androidx.compose.ui.unit.dp
@@ -178,9 +180,7 @@
     override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     @Composable
-    override fun SceneScope.Content(
-        modifier: Modifier,
-    ) =
+    override fun SceneScope.Content(modifier: Modifier) =
         ShadeScene(
             notificationStackScrollView.get(),
             viewModel =
@@ -224,6 +224,13 @@
     modifier: Modifier = Modifier,
     shadeSession: SaveableSession,
 ) {
+    val view = LocalView.current
+    LaunchedEffect(Unit) {
+        if (layoutState.currentTransition?.fromContent == Scenes.Gone) {
+            view.performHapticFeedback(HapticFeedbackConstants.GESTURE_START)
+        }
+    }
+
     val shadeMode by viewModel.shadeMode.collectAsStateWithLifecycle()
     when (shadeMode) {
         is ShadeMode.Single ->
@@ -282,7 +289,7 @@
         animateSceneFloatAsState(
             value = 1f,
             key = QuickSettings.SharedValues.TilesSquishiness,
-            canOverflow = false
+            canOverflow = false,
         )
     val isEmptySpaceClickable by viewModel.isEmptySpaceClickable.collectAsStateWithLifecycle()
     val isMediaVisible by viewModel.isMediaVisible.collectAsStateWithLifecycle()
@@ -320,7 +327,7 @@
                     } else {
                         cutoutInsets
                     }
-                }
+                },
             )
         }
 
@@ -337,7 +344,7 @@
             modifier =
                 Modifier.fillMaxSize()
                     .element(Shade.Elements.BackgroundScrim)
-                    .background(colorResource(R.color.shade_scrim_background_dark)),
+                    .background(colorResource(R.color.shade_scrim_background_dark))
         )
         Layout(
             modifier =
@@ -398,13 +405,13 @@
                     .pointerInteropFilter { true }
                     .verticalNestedScrollToScene(
                         topBehavior = NestedScrollBehavior.EdgeAlways,
-                        isExternalOverscrollGesture = { false }
+                        isExternalOverscrollGesture = { false },
                     )
         ) {
             NotificationStackCutoffGuideline(
                 stackScrollView = notificationStackScrollView,
                 viewModel = notificationsPlaceholderViewModel,
-                modifier = Modifier.align(Alignment.TopCenter)
+                modifier = Modifier.align(Alignment.TopCenter),
             )
         }
     }
@@ -440,24 +447,16 @@
             canOverflow = false,
         )
     val unfoldTranslationXForStartSide by
-        viewModel
-            .unfoldTranslationX(
-                isOnStartSide = true,
-            )
-            .collectAsStateWithLifecycle(0f)
+        viewModel.unfoldTranslationX(isOnStartSide = true).collectAsStateWithLifecycle(0f)
     val unfoldTranslationXForEndSide by
-        viewModel
-            .unfoldTranslationX(
-                isOnStartSide = false,
-            )
-            .collectAsStateWithLifecycle(0f)
+        viewModel.unfoldTranslationX(isOnStartSide = false).collectAsStateWithLifecycle(0f)
 
     val navBarBottomHeight = WindowInsets.navigationBars.asPaddingValues().calculateBottomPadding()
     val bottomPadding by
         animateDpAsState(
             targetValue = if (isCustomizing) 0.dp else navBarBottomHeight,
             animationSpec = tween(customizingAnimationDuration),
-            label = "animateQSSceneBottomPaddingAsState"
+            label = "animateQSSceneBottomPaddingAsState",
         )
     val density = LocalDensity.current
     LaunchedEffect(navBarBottomHeight, density) {
@@ -516,9 +515,7 @@
                     )
         )
 
-        Column(
-            modifier = Modifier.fillMaxSize(),
-        ) {
+        Column(modifier = Modifier.fillMaxSize()) {
             CollapsedShadeHeader(
                 viewModelFactory = viewModel.shadeHeaderViewModelFactory,
                 createTintedIconManager = createTintedIconManager,
@@ -526,9 +523,7 @@
                 statusBarIconController = statusBarIconController,
                 modifier =
                     Modifier.then(brightnessMirrorShowingModifier)
-                        .padding(
-                            horizontal = { unfoldTranslationXForStartSide.roundToInt() },
-                        )
+                        .padding(horizontal = { unfoldTranslationXForStartSide.roundToInt() }),
             )
 
             Row(modifier = Modifier.fillMaxWidth().weight(1f)) {
@@ -536,14 +531,14 @@
                     modifier =
                         Modifier.element(Shade.Elements.SplitShadeStartColumn)
                             .weight(1f)
-                            .graphicsLayer { translationX = unfoldTranslationXForStartSide },
+                            .graphicsLayer { translationX = unfoldTranslationXForStartSide }
                 ) {
                     BrightnessMirror(
                         viewModel = brightnessMirrorViewModel,
                         qsSceneAdapter = viewModel.qsSceneAdapter,
                         // Need to use the offset measured from the container as the header
                         // has to be accounted for
-                        measureFromContainer = true
+                        measureFromContainer = true,
                     )
                     Column(
                         verticalArrangement = Arrangement.Top,
@@ -557,7 +552,7 @@
                                     .thenIf(!isCustomizerShowing) {
                                         Modifier.verticalScroll(
                                                 quickSettingsScrollState,
-                                                enabled = isScrollable
+                                                enabled = isScrollable,
                                             )
                                             .clipScrollableContainer(Orientation.Horizontal)
                                     }
@@ -619,16 +614,16 @@
                             .padding(
                                 end =
                                     dimensionResource(R.dimen.notification_panel_margin_horizontal),
-                                bottom = navBarBottomHeight
+                                bottom = navBarBottomHeight,
                             )
-                            .then(brightnessMirrorShowingModifier)
+                            .then(brightnessMirrorShowingModifier),
                 )
             }
         }
         NotificationStackCutoffGuideline(
             stackScrollView = notificationStackScrollView,
             viewModel = notificationsPlaceholderViewModel,
-            modifier = Modifier.align(Alignment.BottomCenter).navigationBarsPadding()
+            modifier = Modifier.align(Alignment.BottomCenter).navigationBarsPadding(),
         )
     }
 }
@@ -652,6 +647,6 @@
                 null
             } else {
                 { mediaOffsetProvider.offset }
-            }
+            },
     )
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
index 75ae414..b96e40f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/domain/interactor/CommunalInteractorTest.kt
@@ -244,10 +244,7 @@
 
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
-            userTracker.set(
-                userInfos = userInfos,
-                selectedUserIndex = 0,
-            )
+            userTracker.set(userInfos = userInfos, selectedUserIndex = 0)
             runCurrent()
 
             // Widgets available.
@@ -267,21 +264,14 @@
     fun smartspaceDynamicSizing_oneCard_fullSize() =
         testSmartspaceDynamicSizing(
             totalTargets = 1,
-            expectedSizes =
-                listOf(
-                    CommunalContentSize.FULL,
-                )
+            expectedSizes = listOf(CommunalContentSize.FULL),
         )
 
     @Test
     fun smartspace_dynamicSizing_twoCards_halfSize() =
         testSmartspaceDynamicSizing(
             totalTargets = 2,
-            expectedSizes =
-                listOf(
-                    CommunalContentSize.HALF,
-                    CommunalContentSize.HALF,
-                )
+            expectedSizes = listOf(CommunalContentSize.HALF, CommunalContentSize.HALF),
         )
 
     @Test
@@ -293,34 +283,34 @@
                     CommunalContentSize.THIRD,
                     CommunalContentSize.THIRD,
                     CommunalContentSize.THIRD,
-                )
+                ),
         )
 
     @Test
-    fun smartspace_dynamicSizing_fourCards_oneFullAndThreeThirdSize() =
+    fun smartspace_dynamicSizing_fourCards_threeThirdSizeAndOneFullSize() =
         testSmartspaceDynamicSizing(
             totalTargets = 4,
             expectedSizes =
                 listOf(
+                    CommunalContentSize.THIRD,
+                    CommunalContentSize.THIRD,
+                    CommunalContentSize.THIRD,
                     CommunalContentSize.FULL,
-                    CommunalContentSize.THIRD,
-                    CommunalContentSize.THIRD,
-                    CommunalContentSize.THIRD,
-                )
+                ),
         )
 
     @Test
-    fun smartspace_dynamicSizing_fiveCards_twoHalfAndThreeThirdSize() =
+    fun smartspace_dynamicSizing_fiveCards_threeThirdAndTwoHalfSize() =
         testSmartspaceDynamicSizing(
             totalTargets = 5,
             expectedSizes =
                 listOf(
+                    CommunalContentSize.THIRD,
+                    CommunalContentSize.THIRD,
+                    CommunalContentSize.THIRD,
                     CommunalContentSize.HALF,
                     CommunalContentSize.HALF,
-                    CommunalContentSize.THIRD,
-                    CommunalContentSize.THIRD,
-                    CommunalContentSize.THIRD,
-                )
+                ),
         )
 
     @Test
@@ -335,7 +325,7 @@
                     CommunalContentSize.THIRD,
                     CommunalContentSize.THIRD,
                     CommunalContentSize.THIRD,
-                )
+                ),
         )
 
     private fun testSmartspaceDynamicSizing(
@@ -355,7 +345,7 @@
 
             smartspaceRepository.setTimers(targets)
 
-            val smartspaceContent by collectLastValue(underTest.ongoingContent)
+            val smartspaceContent by collectLastValue(underTest.ongoingContent(false))
             assertThat(smartspaceContent?.size).isEqualTo(totalTargets)
             for (index in 0 until totalTargets) {
                 assertThat(smartspaceContent?.get(index)?.size).isEqualTo(expectedSizes[index])
@@ -371,7 +361,7 @@
             // Media is playing.
             mediaRepository.mediaActive()
 
-            val umoContent by collectLastValue(underTest.ongoingContent)
+            val umoContent by collectLastValue(underTest.ongoingContent(true))
 
             assertThat(umoContent?.size).isEqualTo(1)
             assertThat(umoContent?.get(0)).isInstanceOf(CommunalContentModel.Umo::class.java)
@@ -379,6 +369,19 @@
         }
 
     @Test
+    fun umo_mediaPlaying_mediaHostNotVisible_hidesUmo() =
+        testScope.runTest {
+            // Tutorial completed.
+            tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
+
+            // Media is playing.
+            mediaRepository.mediaActive()
+
+            val umoContent by collectLastValue(underTest.ongoingContent(false))
+            assertThat(umoContent?.size).isEqualTo(0)
+        }
+
+    @Test
     fun ongoing_shouldOrderAndSizeByTimestamp() =
         testScope.runTest {
             // Keyguard showing, and tutorial completed.
@@ -401,19 +404,19 @@
             val timer3 = smartspaceTimer("timer3", timestamp = 4L)
             smartspaceRepository.setTimers(listOf(timer1, timer2, timer3))
 
-            val ongoingContent by collectLastValue(underTest.ongoingContent)
+            val ongoingContent by collectLastValue(underTest.ongoingContent(true))
             assertThat(ongoingContent?.size).isEqualTo(4)
             assertThat(ongoingContent?.get(0)?.key)
                 .isEqualTo(CommunalContentModel.KEY.smartspace("timer3"))
-            assertThat(ongoingContent?.get(0)?.size).isEqualTo(CommunalContentSize.FULL)
+            assertThat(ongoingContent?.get(0)?.size).isEqualTo(CommunalContentSize.HALF)
             assertThat(ongoingContent?.get(1)?.key)
                 .isEqualTo(CommunalContentModel.KEY.smartspace("timer2"))
-            assertThat(ongoingContent?.get(1)?.size).isEqualTo(CommunalContentSize.THIRD)
+            assertThat(ongoingContent?.get(1)?.size).isEqualTo(CommunalContentSize.HALF)
             assertThat(ongoingContent?.get(2)?.key).isEqualTo(CommunalContentModel.KEY.umo())
-            assertThat(ongoingContent?.get(2)?.size).isEqualTo(CommunalContentSize.THIRD)
+            assertThat(ongoingContent?.get(2)?.size).isEqualTo(CommunalContentSize.HALF)
             assertThat(ongoingContent?.get(3)?.key)
                 .isEqualTo(CommunalContentModel.KEY.smartspace("timer1"))
-            assertThat(ongoingContent?.get(3)?.size).isEqualTo(CommunalContentSize.THIRD)
+            assertThat(ongoingContent?.get(3)?.size).isEqualTo(CommunalContentSize.HALF)
         }
 
     @Test
@@ -435,10 +438,7 @@
         testScope.runTest {
             // Set to main user, so we can dismiss the tile for the main user.
             val user = userRepository.asMainUser()
-            userTracker.set(
-                userInfos = listOf(user),
-                selectedUserIndex = 0,
-            )
+            userTracker.set(userInfos = listOf(user), selectedUserIndex = 0)
             runCurrent()
 
             tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
@@ -816,10 +816,7 @@
             // Only main user exists.
             val userInfos = listOf(MAIN_USER_INFO)
             userRepository.setUserInfos(userInfos)
-            userTracker.set(
-                userInfos = userInfos,
-                selectedUserIndex = 0,
-            )
+            userTracker.set(userInfos = userInfos, selectedUserIndex = 0)
             runCurrent()
 
             val widgetContent by collectLastValue(underTest.widgetContent)
@@ -853,10 +850,7 @@
             // Work profile is set up.
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
-            userTracker.set(
-                userInfos = userInfos,
-                selectedUserIndex = 0,
-            )
+            userTracker.set(userInfos = userInfos, selectedUserIndex = 0)
             runCurrent()
 
             // When work profile is paused.
@@ -899,10 +893,7 @@
 
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
-            userTracker.set(
-                userInfos = userInfos,
-                selectedUserIndex = 0,
-            )
+            userTracker.set(userInfos = userInfos, selectedUserIndex = 0)
             userRepository.setSelectedUserInfo(MAIN_USER_INFO)
             runCurrent()
 
@@ -914,7 +905,7 @@
 
             setKeyguardFeaturesDisabled(
                 USER_INFO_WORK,
-                DevicePolicyManager.KEYGUARD_DISABLE_WIDGETS_ALL
+                DevicePolicyManager.KEYGUARD_DISABLE_WIDGETS_ALL,
             )
 
             // Widgets under work profile are filtered out. Only the regular widget remains.
@@ -932,10 +923,7 @@
 
             val userInfos = listOf(MAIN_USER_INFO, USER_INFO_WORK)
             userRepository.setUserInfos(userInfos)
-            userTracker.set(
-                userInfos = userInfos,
-                selectedUserIndex = 0,
-            )
+            userTracker.set(userInfos = userInfos, selectedUserIndex = 0)
             userRepository.setSelectedUserInfo(MAIN_USER_INFO)
             runCurrent()
 
@@ -947,7 +935,7 @@
 
             setKeyguardFeaturesDisabled(
                 USER_INFO_WORK,
-                DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_NONE
+                DevicePolicyManager.KEYGUARD_DISABLE_FEATURES_NONE,
             )
 
             // Widgets under work profile are available.
@@ -967,7 +955,7 @@
             kosmos.fakeKeyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.GLANCEABLE_HUB,
                 to = KeyguardState.OCCLUDED,
-                testScope
+                testScope,
             )
 
             assertThat(showCommunalFromOccluded).isTrue()
@@ -983,7 +971,7 @@
             kosmos.fakeKeyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.LOCKSCREEN,
                 to = KeyguardState.OCCLUDED,
-                testScope
+                testScope,
             )
 
             assertThat(showCommunalFromOccluded).isFalse()
@@ -999,7 +987,7 @@
             kosmos.fakeKeyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.GLANCEABLE_HUB,
                 to = KeyguardState.OCCLUDED,
-                testScope
+                testScope,
             )
             runCurrent()
             kosmos.setCommunalAvailable(false)
@@ -1017,13 +1005,13 @@
             kosmos.fakeKeyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.GLANCEABLE_HUB,
                 to = KeyguardState.OCCLUDED,
-                testScope
+                testScope,
             )
             runCurrent()
             kosmos.fakeKeyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.OCCLUDED,
                 to = KeyguardState.PRIMARY_BOUNCER,
-                testScope
+                testScope,
             )
 
             assertThat(showCommunalFromOccluded).isTrue()
@@ -1039,7 +1027,7 @@
             kosmos.fakeKeyguardTransitionRepository.sendTransitionSteps(
                 from = KeyguardState.DREAMING,
                 to = KeyguardState.OCCLUDED,
-                testScope
+                testScope,
             )
 
             assertThat(showCommunalFromOccluded).isTrue()
@@ -1049,7 +1037,7 @@
         return CommunalSmartspaceTimer(
             smartspaceTargetId = id,
             createdTimestampMillis = timestamp,
-            remoteViews = mock(RemoteViews::class.java)
+            remoteViews = mock(RemoteViews::class.java),
         )
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
index 780d357..09daa51 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
@@ -43,6 +43,7 @@
 import com.android.systemui.communal.domain.interactor.communalTutorialInteractor
 import com.android.systemui.communal.domain.model.CommunalContentModel
 import com.android.systemui.communal.shared.log.CommunalMetricsLogger
+import com.android.systemui.communal.shared.model.CommunalContentSize
 import com.android.systemui.communal.shared.model.CommunalScenes
 import com.android.systemui.communal.ui.viewmodel.CommunalViewModel
 import com.android.systemui.communal.ui.viewmodel.CommunalViewModel.Companion.POPUP_AUTO_HIDE_TIMEOUT_MS
@@ -150,10 +151,7 @@
         kosmos.fakeFeatureFlagsClassic.set(COMMUNAL_SERVICE_ENABLED, true)
         mSetFlagsRule.enableFlags(FLAG_COMMUNAL_HUB)
 
-        kosmos.fakeUserTracker.set(
-            userInfos = listOf(MAIN_USER_INFO),
-            selectedUserIndex = 0,
-        )
+        kosmos.fakeUserTracker.set(userInfos = listOf(MAIN_USER_INFO), selectedUserIndex = 0)
         whenever(mediaHost.visible).thenReturn(true)
 
         kosmos.powerInteractor.setAwakeForTest()
@@ -249,6 +247,87 @@
         }
 
     @Test
+    fun ongoingContent_umoAndOneTimer_sizedAppropriately() =
+        testScope.runTest {
+            // Widgets available.
+            widgetRepository.addWidget(appWidgetId = 0, rank = 30)
+            widgetRepository.addWidget(appWidgetId = 1, rank = 20)
+
+            // Smartspace available.
+            smartspaceRepository.setTimers(
+                listOf(
+                    CommunalSmartspaceTimer(
+                        smartspaceTargetId = "target",
+                        createdTimestampMillis = 0L,
+                        remoteViews = Mockito.mock(RemoteViews::class.java),
+                    )
+                )
+            )
+
+            // Media playing.
+            mediaRepository.mediaActive()
+
+            val communalContent by collectLastValue(underTest.communalContent)
+
+            // One timer, UMO, two widgets, and cta.
+            assertThat(communalContent?.size).isEqualTo(5)
+
+            val timer = communalContent?.get(0)
+            val umo = communalContent?.get(1)
+
+            assertThat(timer).isInstanceOf(CommunalContentModel.Smartspace::class.java)
+            assertThat(umo).isInstanceOf(CommunalContentModel.Umo::class.java)
+
+            assertThat(timer?.size).isEqualTo(CommunalContentSize.HALF)
+            assertThat(umo?.size).isEqualTo(CommunalContentSize.HALF)
+        }
+
+    @Test
+    fun ongoingContent_umoAndTwoTimers_sizedAppropriately() =
+        testScope.runTest {
+            // Widgets available.
+            widgetRepository.addWidget(appWidgetId = 0, rank = 30)
+            widgetRepository.addWidget(appWidgetId = 1, rank = 20)
+
+            // Smartspace available.
+            smartspaceRepository.setTimers(
+                listOf(
+                    CommunalSmartspaceTimer(
+                        smartspaceTargetId = "target",
+                        createdTimestampMillis = 0L,
+                        remoteViews = Mockito.mock(RemoteViews::class.java),
+                    ),
+                    CommunalSmartspaceTimer(
+                        smartspaceTargetId = "target",
+                        createdTimestampMillis = 0L,
+                        remoteViews = Mockito.mock(RemoteViews::class.java),
+                    ),
+                )
+            )
+
+            // Media playing.
+            mediaRepository.mediaActive()
+
+            val communalContent by collectLastValue(underTest.communalContent)
+
+            // Two timers, UMO, two widgets, and cta.
+            assertThat(communalContent?.size).isEqualTo(6)
+
+            val timer1 = communalContent?.get(0)
+            val timer2 = communalContent?.get(1)
+            val umo = communalContent?.get(2)
+
+            assertThat(timer1).isInstanceOf(CommunalContentModel.Smartspace::class.java)
+            assertThat(timer2).isInstanceOf(CommunalContentModel.Smartspace::class.java)
+            assertThat(umo).isInstanceOf(CommunalContentModel.Umo::class.java)
+
+            // One full-sized timer and a half-sized timer and half-sized UMO.
+            assertThat(timer1?.size).isEqualTo(CommunalContentSize.HALF)
+            assertThat(timer2?.size).isEqualTo(CommunalContentSize.HALF)
+            assertThat(umo?.size).isEqualTo(CommunalContentSize.FULL)
+        }
+
+    @Test
     fun communalContent_mediaHostVisible_umoIncluded() =
         testScope.runTest {
             // Media playing.
@@ -497,7 +576,7 @@
                     TransitionStep(
                         from = KeyguardState.LOCKSCREEN,
                         to = KeyguardState.GLANCEABLE_HUB,
-                    )
+                    ),
             )
             // Shade not expanded.
             if (!SceneContainerFlag.isEnabled) shadeTestUtil.setLockscreenShadeExpansion(0f)
@@ -550,8 +629,8 @@
                 stateTransition =
                     TransitionStep(
                         from = KeyguardState.DREAMING,
-                        to = KeyguardState.GLANCEABLE_HUB,
-                    )
+                        to = KeyguardState.GLANCEABLE_HUB
+                    ),
             )
 
             // Then flow is not frozen
@@ -570,8 +649,8 @@
                 stateTransition =
                     TransitionStep(
                         from = KeyguardState.GLANCEABLE_HUB,
-                        to = KeyguardState.OCCLUDED,
-                    )
+                        to = KeyguardState.OCCLUDED
+                    ),
             )
 
             // Then flow is not frozen
@@ -595,7 +674,7 @@
                     TransitionStep(
                         from = KeyguardState.LOCKSCREEN,
                         to = KeyguardState.GLANCEABLE_HUB,
-                    )
+                    ),
             )
 
             // Then flow is not frozen
@@ -614,7 +693,7 @@
                         to = KeyguardState.OCCLUDED,
                         transitionState = TransitionState.STARTED,
                         value = 0f,
-                    )
+                    ),
             )
 
             // Then flow is frozen
@@ -629,7 +708,7 @@
                         to = KeyguardState.OCCLUDED,
                         transitionState = TransitionState.FINISHED,
                         value = 1f,
-                    )
+                    ),
             )
 
             // Then flow is not frozen
@@ -658,8 +737,8 @@
                 stateTransition =
                     TransitionStep(
                         from = KeyguardState.DREAMING,
-                        to = KeyguardState.GLANCEABLE_HUB,
-                    )
+                        to = KeyguardState.GLANCEABLE_HUB
+                    ),
             )
 
             // Widgets available
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorTest.kt
index 1981a2d..41cc953 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorTest.kt
@@ -25,7 +25,7 @@
 import com.android.systemui.bouncer.data.repository.fakeKeyguardBouncerRepository
 import com.android.systemui.bouncer.domain.interactor.alternateBouncerInteractor
 import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.shared.model.DismissAction
@@ -39,8 +39,6 @@
 import com.android.systemui.scene.data.repository.Transition
 import com.android.systemui.scene.data.repository.setSceneTransition
 import com.android.systemui.scene.domain.interactor.sceneInteractor
-import com.android.systemui.scene.domain.resolver.notifShadeSceneFamilyResolver
-import com.android.systemui.scene.domain.resolver.quickSettingsSceneFamilyResolver
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.testKosmos
@@ -79,12 +77,9 @@
                 dismissInteractor = dismissInteractor,
                 applicationScope = testScope.backgroundScope,
                 sceneInteractor = { kosmos.sceneInteractor },
-                deviceEntryInteractor = { kosmos.deviceEntryInteractor },
-                quickSettingsSceneFamilyResolver = { kosmos.quickSettingsSceneFamilyResolver },
-                notifShadeSceneFamilyResolver = { kosmos.notifShadeSceneFamilyResolver },
+                deviceUnlockedInteractor = { kosmos.deviceUnlockedInteractor },
                 powerInteractor = kosmos.powerInteractor,
                 alternateBouncerInteractor = kosmos.alternateBouncerInteractor,
-                keyguardInteractor = { kosmos.keyguardInteractor },
                 shadeInteractor = { kosmos.shadeInteractor },
             )
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModelTest.kt
index f6865f13..642d9a0 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModelTest.kt
@@ -21,6 +21,7 @@
 import androidx.test.filters.SmallTest
 import com.android.compose.animation.scene.Back
 import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -28,6 +29,7 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.viewmodel.SceneContainerEdge
 import com.android.systemui.shade.ui.viewmodel.notificationsShadeOverlayActionsViewModel
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
@@ -47,7 +49,7 @@
     private val underTest = kosmos.notificationsShadeOverlayActionsViewModel
 
     @Test
-    fun upTransitionSceneKey_hidesShade() =
+    fun up_hidesShade() =
         testScope.runTest {
             val actions by collectLastValue(underTest.actions)
             underTest.activateIn(this)
@@ -66,4 +68,22 @@
             assertThat((actions?.get(Back) as? UserActionResult.HideOverlay)?.overlay)
                 .isEqualTo(Overlays.NotificationsShade)
         }
+
+    @Test
+    fun downFromTopRight_switchesToQuickSettingsShade() =
+        testScope.runTest {
+            val actions by collectLastValue(underTest.actions)
+            underTest.activateIn(this)
+
+            assertThat(
+                    (actions?.get(
+                            Swipe(
+                                direction = SwipeDirection.Down,
+                                fromSource = SceneContainerEdge.TopRight,
+                            )
+                        ) as? UserActionResult.ReplaceByOverlay)
+                        ?.overlay
+                )
+                .isEqualTo(Overlays.QuickSettingsShade)
+        }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelTest.kt
index 762941d..fd1c043 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelTest.kt
@@ -21,6 +21,7 @@
 import androidx.test.filters.SmallTest
 import com.android.compose.animation.scene.Back
 import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -28,6 +29,7 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.viewmodel.SceneContainerEdge
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
@@ -46,7 +48,7 @@
     private val underTest = kosmos.quickSettingsShadeOverlayActionsViewModel
 
     @Test
-    fun upTransitionSceneKey_hidesShade() =
+    fun up_hidesShade() =
         testScope.runTest {
             val actions by collectLastValue(underTest.actions)
             underTest.activateIn(this)
@@ -57,12 +59,44 @@
         }
 
     @Test
-    fun back_hidesShade() =
+    fun back_notEditing_hidesShade() =
+        testScope.runTest {
+            val actions by collectLastValue(underTest.actions)
+            val isEditing by
+                collectLastValue(kosmos.quickSettingsContainerViewModel.editModeViewModel.isEditing)
+            underTest.activateIn(this)
+            assertThat(isEditing).isFalse()
+
+            assertThat((actions?.get(Back) as? UserActionResult.HideOverlay)?.overlay)
+                .isEqualTo(Overlays.QuickSettingsShade)
+        }
+
+    @Test
+    fun back_whileEditing_doesNotHideShade() =
         testScope.runTest {
             val actions by collectLastValue(underTest.actions)
             underTest.activateIn(this)
 
-            assertThat((actions?.get(Back) as? UserActionResult.HideOverlay)?.overlay)
-                .isEqualTo(Overlays.QuickSettingsShade)
+            kosmos.quickSettingsContainerViewModel.editModeViewModel.startEditing()
+
+            assertThat(actions?.get(Back)).isNull()
+        }
+
+    @Test
+    fun downFromTopLeft_switchesToNotificationsShade() =
+        testScope.runTest {
+            val actions by collectLastValue(underTest.actions)
+            underTest.activateIn(this)
+
+            assertThat(
+                    (actions?.get(
+                            Swipe(
+                                direction = SwipeDirection.Down,
+                                fromSource = SceneContainerEdge.TopLeft,
+                            )
+                        ) as? UserActionResult.ReplaceByOverlay)
+                        ?.overlay
+                )
+                .isEqualTo(Overlays.NotificationsShade)
         }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
index 6986cf8e..62b6391 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
@@ -37,7 +37,6 @@
 import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.lifecycle.activateIn
-import com.android.systemui.qs.ui.adapter.FakeQSSceneAdapter
 import com.android.systemui.res.R
 import com.android.systemui.scene.domain.interactor.sceneBackInteractor
 import com.android.systemui.scene.domain.interactor.sceneInteractor
@@ -46,7 +45,6 @@
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.testKosmos
-import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -61,7 +59,7 @@
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
-    private val qsFlexiglassAdapter = FakeQSSceneAdapter({ mock() })
+    private val qsFlexiglassAdapter = kosmos.fakeQsSceneAdapter
 
     private val sceneInteractor = kosmos.sceneInteractor
     private val sceneBackInteractor = kosmos.sceneBackInteractor
@@ -101,10 +99,8 @@
                     mapOf(
                         Back to UserActionResult(Scenes.Shade),
                         Swipe(SwipeDirection.Up) to UserActionResult(Scenes.Shade),
-                        Swipe(
-                            fromSource = Edge.Bottom,
-                            direction = SwipeDirection.Up,
-                        ) to UserActionResult(SceneFamilies.Home)
+                        Swipe(fromSource = Edge.Bottom, direction = SwipeDirection.Up) to
+                            UserActionResult(SceneFamilies.Home),
                     )
                 )
             assertThat(homeScene).isEqualTo(Scenes.Gone)
@@ -130,10 +126,8 @@
                     mapOf(
                         Back to UserActionResult(Scenes.Lockscreen),
                         Swipe(SwipeDirection.Up) to UserActionResult(Scenes.Lockscreen),
-                        Swipe(
-                            fromSource = Edge.Bottom,
-                            direction = SwipeDirection.Up,
-                        ) to UserActionResult(SceneFamilies.Home)
+                        Swipe(fromSource = Edge.Bottom, direction = SwipeDirection.Up) to
+                            UserActionResult(SceneFamilies.Home),
                     )
                 )
             assertThat(homeScene).isEqualTo(Scenes.Lockscreen)
@@ -161,10 +155,8 @@
                     mapOf(
                         Back to UserActionResult(Scenes.Shade),
                         Swipe(SwipeDirection.Up) to UserActionResult(Scenes.Shade),
-                        Swipe(
-                            fromSource = Edge.Bottom,
-                            direction = SwipeDirection.Up,
-                        ) to UserActionResult(SceneFamilies.Home)
+                        Swipe(fromSource = Edge.Bottom, direction = SwipeDirection.Up) to
+                            UserActionResult(SceneFamilies.Home),
                     )
                 )
             assertThat(homeScene).isEqualTo(Scenes.Gone)
@@ -187,10 +179,8 @@
                     mapOf(
                         Back to UserActionResult(Scenes.Shade),
                         Swipe(SwipeDirection.Up) to UserActionResult(Scenes.Shade),
-                        Swipe(
-                            fromSource = Edge.Bottom,
-                            direction = SwipeDirection.Up,
-                        ) to UserActionResult(SceneFamilies.Home)
+                        Swipe(fromSource = Edge.Bottom, direction = SwipeDirection.Up) to
+                            UserActionResult(SceneFamilies.Home),
                     )
                 )
             assertThat(homeScene).isEqualTo(Scenes.Lockscreen)
@@ -225,10 +215,8 @@
                     mapOf(
                         Back to UserActionResult(Scenes.Shade),
                         Swipe(SwipeDirection.Up) to UserActionResult(Scenes.Shade),
-                        Swipe(
-                            fromSource = Edge.Bottom,
-                            direction = SwipeDirection.Up,
-                        ) to UserActionResult(SceneFamilies.Home)
+                        Swipe(fromSource = Edge.Bottom, direction = SwipeDirection.Up) to
+                            UserActionResult(SceneFamilies.Home),
                     )
                 )
             assertThat(homeScene).isEqualTo(Scenes.Gone)
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/PrimaryBouncerInteractor.kt b/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/PrimaryBouncerInteractor.kt
index c28bce2..6d6cd45 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/PrimaryBouncerInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/domain/interactor/PrimaryBouncerInteractor.kt
@@ -294,7 +294,9 @@
 
     /** Tell the bouncer that bouncer is requested when device is already authenticated */
     fun notifyUserRequestedBouncerWhenAlreadyAuthenticated(userId: Int) {
-        applicationScope.launch { repository.setKeyguardAuthenticatedPrimaryAuth(userId) }
+        applicationScope.launch {
+            repository.setUserRequestedBouncerWhenAlreadyAuthenticated(userId)
+        }
     }
 
     /** Tell the bouncer that keyguard is authenticated with biometrics. */
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerContainerViewModel.kt
index 73a8810..c60f932 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerContainerViewModel.kt
@@ -17,12 +17,14 @@
 package com.android.systemui.bouncer.ui.viewmodel
 
 import androidx.compose.runtime.getValue
-import com.android.keyguard.ViewMediatorCallback
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor
+import com.android.systemui.deviceentry.domain.interactor.DeviceUnlockedInteractor
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.lifecycle.Hydrator
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
+import com.android.systemui.util.kotlin.Utils.Companion.sample
+import com.android.systemui.util.kotlin.sample
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import kotlinx.coroutines.coroutineScope
@@ -34,7 +36,7 @@
     private val legacyInteractor: PrimaryBouncerInteractor,
     private val authenticationInteractor: AuthenticationInteractor,
     private val selectedUserInteractor: SelectedUserInteractor,
-    private val viewMediatorCallback: ViewMediatorCallback?,
+    private val deviceUnlockedInteractor: DeviceUnlockedInteractor,
 ) : ExclusiveActivatable() {
 
     private val hydrator = Hydrator("BouncerContainerViewModel")
@@ -45,6 +47,18 @@
     override suspend fun onActivated(): Nothing {
         coroutineScope {
             launch {
+                legacyInteractor.isShowing
+                    .sample(deviceUnlockedInteractor.deviceUnlockStatus, ::Pair)
+                    .collect { (isShowing, unlockStatus) ->
+                        if (isShowing && unlockStatus.isUnlocked) {
+                            legacyInteractor.notifyUserRequestedBouncerWhenAlreadyAuthenticated(
+                                selectedUserInteractor.getSelectedUserId()
+                            )
+                        }
+                    }
+            }
+
+            launch {
                 authenticationInteractor.onAuthenticationResult.collect { authenticationSucceeded ->
                     if (authenticationSucceeded) {
                         legacyInteractor.notifyKeyguardAuthenticatedPrimaryAuth(
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
index b570e14..a687734 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
@@ -117,7 +117,7 @@
     sceneInteractor: SceneInteractor,
     @CommunalLog logBuffer: LogBuffer,
     @CommunalTableLog tableLogBuffer: TableLogBuffer,
-    private val managedProfileController: ManagedProfileController
+    private val managedProfileController: ManagedProfileController,
 ) {
     private val logger = Logger(logBuffer, "CommunalInteractor")
 
@@ -154,7 +154,7 @@
         allOf(
                 communalSettingsInteractor.isCommunalEnabled,
                 not(keyguardInteractor.isEncryptedOrLockdown),
-                keyguardInteractor.isKeyguardShowing
+                keyguardInteractor.isKeyguardShowing,
             )
             .distinctUntilChanged()
             .onEach { available ->
@@ -342,7 +342,7 @@
     fun changeScene(
         newScene: SceneKey,
         loggingReason: String,
-        transitionKey: TransitionKey? = null
+        transitionKey: TransitionKey? = null,
     ) = communalSceneInteractor.changeScene(newScene, loggingReason, transitionKey)
 
     fun setEditModeOpen(isOpen: Boolean) {
@@ -354,9 +354,7 @@
     }
 
     /** Show the widget editor Activity. */
-    fun showWidgetEditor(
-        shouldOpenWidgetPickerOnStart: Boolean = false,
-    ) {
+    fun showWidgetEditor(shouldOpenWidgetPickerOnStart: Boolean = false) {
         communalSceneInteractor.setEditModeState(EditModeState.STARTING)
         editWidgetsActivityStarter.startActivity(shouldOpenWidgetPickerOnStart)
     }
@@ -419,7 +417,7 @@
                     IntentFilter().apply {
                         addAction(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)
                         addAction(Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE)
-                    },
+                    }
             )
             .emitOnStart()
 
@@ -450,7 +448,7 @@
                             rank = widget.rank,
                             providerInfo = widget.providerInfo,
                             appWidgetHost = appWidgetHost,
-                            inQuietMode = isQuietModeEnabled(widget.providerInfo.profile)
+                            inQuietMode = isQuietModeEnabled(widget.providerInfo.profile),
                         )
                     }
                     is CommunalWidgetContentModel.Pending -> {
@@ -468,7 +466,7 @@
     /** Filter widgets based on whether their associated profile is allowed by device policy. */
     private fun filterWidgetsAllowedByDevicePolicy(
         list: List<CommunalWidgetContentModel>,
-        disallowedByDevicePolicyUser: UserInfo?
+        disallowedByDevicePolicyUser: UserInfo?,
     ): List<CommunalWidgetContentModel> =
         if (disallowedByDevicePolicyUser == null) {
             list
@@ -507,7 +505,7 @@
      * A flow of ongoing content, including smartspace timers and umo, ordered by creation time and
      * sized dynamically.
      */
-    val ongoingContent: Flow<List<CommunalContentModel.Ongoing>> =
+    fun ongoingContent(isMediaHostVisible: Boolean): Flow<List<CommunalContentModel.Ongoing>> =
         combine(smartspaceRepository.timers, mediaRepository.mediaModel) { timers, media ->
                 val ongoingContent = mutableListOf<CommunalContentModel.Ongoing>()
 
@@ -523,22 +521,20 @@
                 )
 
                 // Add UMO
-                if (media.hasAnyMediaOrRecommendation) {
+                if (isMediaHostVisible && media.hasAnyMediaOrRecommendation) {
                     ongoingContent.add(
                         CommunalContentModel.Umo(
-                            createdTimestampMillis = media.createdTimestampMillis,
+                            createdTimestampMillis = media.createdTimestampMillis
                         )
                     )
                 }
 
-                // Order by creation time descending
+                // Order by creation time descending.
                 ongoingContent.sortByDescending { it.createdTimestampMillis }
+                // Resize the items.
+                ongoingContent.resizeItems()
 
-                // Dynamic sizing
-                ongoingContent.forEachIndexed { index, model ->
-                    model.size = dynamicContentSize(ongoingContent.size, index)
-                }
-
+                // Return the sorted and resized items.
                 ongoingContent
             }
             .flowOn(bgDispatcher)
@@ -548,7 +544,7 @@
      * stale data following user deletion.
      */
     private fun filterWidgetsByExistingUsers(
-        list: List<CommunalWidgetContentModel>,
+        list: List<CommunalWidgetContentModel>
     ): List<CommunalWidgetContentModel> {
         val currentUserIds = userTracker.userProfiles.map { it.id }.toSet()
         return list.filter { widget ->
@@ -560,6 +556,40 @@
         }
     }
 
+    // Dynamically resizes the height of items in the list of ongoing items such that they fit in
+    // columns in as compact a space as possible.
+    //
+    // Currently there are three possible sizes. When the total number is 1, size for that  content
+    // is [FULL], when the total number is 2, size for each is [HALF], and 3, size for  each is
+    // [THIRD].
+    //
+    // This algorithm also respects each item's minimum size. All items in a column will have the
+    // same size, and all items in a column will be no smaller than any item's minimum size.
+    private fun List<CommunalContentModel.Ongoing>.resizeItems() {
+        fun resizeColumn(c: List<CommunalContentModel.Ongoing>) {
+            if (c.isEmpty()) return
+            val newSize = CommunalContentSize.toSize(span = FULL.span / c.size)
+            c.forEach { item -> item.size = newSize }
+        }
+
+        val column = mutableListOf<CommunalContentModel.Ongoing>()
+        var available = FULL.span
+
+        forEach { item ->
+            if (available < item.minSize.span) {
+                resizeColumn(column)
+                column.clear()
+                available = FULL.span
+            }
+
+            column.add(item)
+            available -= item.minSize.span
+        }
+
+        // Make sure to resize the final column.
+        resizeColumn(column)
+    }
+
     companion object {
         const val TAG = "CommunalInteractor"
 
@@ -574,31 +604,6 @@
          * of -1 means that the user's chosen screen timeout will be used instead.
          */
         const val AWAKE_INTERVAL_MS = -1
-
-        /**
-         * Calculates the content size dynamically based on the total number of contents of that
-         * type.
-         *
-         * Contents with the same type are expected to fill each column evenly. Currently there are
-         * three possible sizes. When the total number is 1, size for that content is [FULL], when
-         * the total number is 2, size for each is [HALF], and 3, size for each is [THIRD].
-         *
-         * When dynamic contents fill in multiple columns, the first column follows the algorithm
-         * above, and the remaining contents are packed in [THIRD]s. For example, when the total
-         * number if 4, the first one is [FULL], filling the column, and the remaining 3 are
-         * [THIRD].
-         *
-         * @param size The total number of contents of this type.
-         * @param index The index of the current content of this type.
-         */
-        private fun dynamicContentSize(size: Int, index: Int): CommunalContentSize {
-            val remainder = size % CommunalContentSize.entries.size
-            return CommunalContentSize.toSize(
-                span =
-                    FULL.span /
-                        if (index > remainder - 1) CommunalContentSize.entries.size else remainder
-            )
-        }
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt
index 4c821d4..c2f6e85 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/model/CommunalContentModel.kt
@@ -34,12 +34,18 @@
     /** Size to be rendered in the grid. */
     val size: CommunalContentSize
 
+    /** The minimum size content can be resized to. */
+    val minSize: CommunalContentSize
+        get() = CommunalContentSize.HALF
+
     /**
      * A type of communal content is ongoing / live / ephemeral, and can be sized and ordered
      * dynamically.
      */
     sealed interface Ongoing : CommunalContentModel {
         override var size: CommunalContentSize
+        override val minSize
+            get() = CommunalContentSize.THIRD
 
         /** Timestamp in milliseconds of when the content was created. */
         val createdTimestampMillis: Long
@@ -72,7 +78,7 @@
         data class DisabledWidget(
             override val appWidgetId: Int,
             override val rank: Int,
-            val providerInfo: AppWidgetProviderInfo
+            val providerInfo: AppWidgetProviderInfo,
         ) : WidgetContent {
             override val key = KEY.disabledWidget(appWidgetId)
             override val componentName: ComponentName = providerInfo.provider
@@ -109,10 +115,7 @@
         override val size = CommunalContentSize.HALF
     }
 
-    class Tutorial(
-        id: Int,
-        override var size: CommunalContentSize,
-    ) : CommunalContentModel {
+    class Tutorial(id: Int, override var size: CommunalContentSize) : CommunalContentModel {
         override val key = KEY.tutorial(id)
     }
 
@@ -128,6 +131,7 @@
     class Umo(
         override val createdTimestampMillis: Long,
         override var size: CommunalContentSize = CommunalContentSize.HALF,
+        override var minSize: CommunalContentSize = CommunalContentSize.HALF,
     ) : Ongoing {
         override val key = KEY.umo()
     }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
index d69ba1b..53109ac 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
@@ -125,15 +125,9 @@
     private var frozenCommunalContent: List<CommunalContentModel>? = null
 
     private val ongoingContent =
-        combine(
-            isMediaHostVisible,
-            communalInteractor.ongoingContent.onEach { mediaHost.updateViewVisibility() }
-        ) { mediaVisible, ongoingContent ->
-            if (mediaVisible) {
-                ongoingContent
-            } else {
-                // Media is not visible, don't show UMO
-                ongoingContent.filterNot { it is CommunalContentModel.Umo }
+        isMediaHostVisible.flatMapLatest { isMediaHostVisible ->
+            communalInteractor.ongoingContent(isMediaHostVisible).onEach {
+                mediaHost.updateViewVisibility()
             }
         }
 
@@ -148,8 +142,7 @@
                     ongoingContent,
                     communalInteractor.widgetContent,
                     communalInteractor.ctaTileContent,
-                ) { ongoing, widgets, ctaTile,
-                    ->
+                ) { ongoing, widgets, ctaTile ->
                     ongoing + widgets + ctaTile
                 }
             }
@@ -172,10 +165,10 @@
         allOf(
                 keyguardTransitionInteractor.isFinishedIn(
                     scene = Scenes.Communal,
-                    stateWithoutSceneContainer = KeyguardState.GLANCEABLE_HUB
+                    stateWithoutSceneContainer = KeyguardState.GLANCEABLE_HUB,
                 ),
                 keyguardInteractor.isKeyguardOccluded,
-                not(keyguardInteractor.isAbleToDream)
+                not(keyguardInteractor.isAbleToDream),
             )
             .distinctUntilChanged()
             .onEach { logger.d("isCommunalContentFlowFrozen: $it") }
@@ -208,7 +201,7 @@
         combine(
                 keyguardTransitionInteractor.isFinishedIn(
                     scene = Scenes.Communal,
-                    stateWithoutSceneContainer = KeyguardState.GLANCEABLE_HUB
+                    stateWithoutSceneContainer = KeyguardState.GLANCEABLE_HUB,
                 ),
                 communalInteractor.isIdleOnCommunal,
                 shadeInteractor.isAnyFullyExpanded,
@@ -221,7 +214,7 @@
         object : View.AccessibilityDelegate() {
             override fun onInitializeAccessibilityNodeInfo(
                 host: View,
-                info: AccessibilityNodeInfo
+                info: AccessibilityNodeInfo,
             ) {
                 super.onInitializeAccessibilityNodeInfo(host, info)
                 // Hint user to long press in order to enter edit mode
@@ -230,7 +223,7 @@
                         AccessibilityNodeInfo.AccessibilityAction.ACTION_LONG_CLICK.id,
                         resources
                             .getString(R.string.accessibility_action_label_edit_widgets)
-                            .lowercase()
+                            .lowercase(),
                     )
                 )
             }
@@ -238,7 +231,7 @@
             override fun performAccessibilityAction(
                 host: View,
                 action: Int,
-                args: Bundle?
+                args: Bundle?,
             ): Boolean {
                 when (action) {
                     AccessibilityNodeInfo.AccessibilityAction.ACTION_LONG_CLICK.id -> {
@@ -271,9 +264,7 @@
         }
     }
 
-    override fun onOpenWidgetEditor(
-        shouldOpenWidgetPickerOnStart: Boolean,
-    ) {
+    override fun onOpenWidgetEditor(shouldOpenWidgetPickerOnStart: Boolean) {
         persistScrollPosition()
         communalInteractor.showWidgetEditor(shouldOpenWidgetPickerOnStart)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractor.kt
index 8495778..eb9b07a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractor.kt
@@ -21,7 +21,7 @@
 import com.android.systemui.bouncer.shared.flag.ComposeBouncerFlags
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
+import com.android.systemui.deviceentry.domain.interactor.DeviceUnlockedInteractor
 import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.keyguard.shared.model.DismissAction
 import com.android.systemui.keyguard.shared.model.KeyguardDone
@@ -29,13 +29,12 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.scene.domain.interactor.SceneInteractor
-import com.android.systemui.scene.domain.resolver.NotifShadeSceneFamilyResolver
-import com.android.systemui.scene.domain.resolver.QuickSettingsSceneFamilyResolver
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.util.kotlin.Utils.Companion.sampleFilter
 import com.android.systemui.util.kotlin.sample
+import dagger.Lazy
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -61,14 +60,11 @@
     transitionInteractor: KeyguardTransitionInteractor,
     val dismissInteractor: KeyguardDismissInteractor,
     @Application private val applicationScope: CoroutineScope,
-    sceneInteractor: dagger.Lazy<SceneInteractor>,
-    deviceEntryInteractor: dagger.Lazy<DeviceEntryInteractor>,
-    quickSettingsSceneFamilyResolver: dagger.Lazy<QuickSettingsSceneFamilyResolver>,
-    notifShadeSceneFamilyResolver: dagger.Lazy<NotifShadeSceneFamilyResolver>,
+    sceneInteractor: Lazy<SceneInteractor>,
+    deviceUnlockedInteractor: Lazy<DeviceUnlockedInteractor>,
     powerInteractor: PowerInteractor,
     alternateBouncerInteractor: AlternateBouncerInteractor,
-    keyguardInteractor: dagger.Lazy<KeyguardInteractor>,
-    shadeInteractor: dagger.Lazy<ShadeInteractor>,
+    shadeInteractor: Lazy<ShadeInteractor>,
 ) {
     val dismissAction: Flow<DismissAction> = repository.dismissAction
 
@@ -106,18 +102,18 @@
         if (SceneContainerFlag.isEnabled) {
             combine(
                     sceneInteractor.get().currentScene,
-                    deviceEntryInteractor.get().isUnlocked,
-                ) { scene, isUnlocked ->
-                    isUnlocked &&
-                        (quickSettingsSceneFamilyResolver.get().includesScene(scene) ||
-                            notifShadeSceneFamilyResolver.get().includesScene(scene))
+                    deviceUnlockedInteractor.get().deviceUnlockStatus,
+                ) { scene, unlockStatus ->
+                    unlockStatus.isUnlocked &&
+                        (scene == Scenes.QuickSettings || scene == Scenes.Shade)
                 }
                 .distinctUntilChanged()
         } else if (ComposeBouncerFlags.isOnlyComposeBouncerEnabled()) {
-            shadeInteractor.get().isAnyExpanded.sample(
-                keyguardInteractor.get().isKeyguardDismissible
-            ) { isAnyExpanded, isKeyguardDismissible ->
-                isAnyExpanded && isKeyguardDismissible
+            combine(
+                shadeInteractor.get().isAnyExpanded,
+                deviceUnlockedInteractor.get().deviceUnlockStatus,
+            ) { isAnyExpanded, deviceUnlockStatus ->
+                isAnyExpanded && deviceUnlockStatus.isUnlocked
             }
         } else {
             flow {
@@ -132,7 +128,7 @@
         merge(
                 finishedTransitionToGone,
                 isOnShadeWhileUnlocked.filter { it }.map {},
-                dismissInteractor.dismissKeyguardRequestWithImmediateDismissAction
+                dismissInteractor.dismissKeyguardRequestWithImmediateDismissAction,
             )
             .sample(dismissAction)
             .filterNot { it is DismissAction.None }
@@ -142,11 +138,11 @@
         combine(
                 transitionInteractor.isFinishedIn(
                     scene = Scenes.Gone,
-                    stateWithoutSceneContainer = GONE
+                    stateWithoutSceneContainer = GONE,
                 ),
                 transitionInteractor.isFinishedIn(
                     scene = Scenes.Bouncer,
-                    stateWithoutSceneContainer = PRIMARY_BOUNCER
+                    stateWithoutSceneContainer = PRIMARY_BOUNCER,
                 ),
                 alternateBouncerInteractor.isVisible,
                 isOnShadeWhileUnlocked,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
index 44aafab..ad1a32e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
@@ -76,7 +76,7 @@
             .filter { enabled -> !enabled }
             .sampleCombine(
                 internalTransitionInteractor.currentTransitionInfoInternal,
-                biometricSettingsRepository.isCurrentUserInLockdown
+                biometricSettingsRepository.isCurrentUserInLockdown,
             )
             .map { (_, transitionInfo, inLockdown) ->
                 // ...we hide the keyguard, if it's showing and we're not in lockdown. In that case,
@@ -91,12 +91,18 @@
          */
         scope.launch {
             if (!SceneContainerFlag.isEnabled) {
-                showKeyguardWhenReenabled
-                    .filter { shouldDismiss -> shouldDismiss }
-                    .collect {
-                        keyguardDismissTransitionInteractor.startDismissKeyguardTransition(
-                            "keyguard disabled"
-                        )
+                repository.isKeyguardEnabled
+                    .filter { enabled -> !enabled }
+                    .sampleCombine(
+                        biometricSettingsRepository.isCurrentUserInLockdown,
+                        internalTransitionInteractor.currentTransitionInfoInternal,
+                    )
+                    .collect { (_, inLockdown, currentTransitionInfo) ->
+                        if (currentTransitionInfo.to != KeyguardState.GONE && !inLockdown) {
+                            keyguardDismissTransitionInteractor.startDismissKeyguardTransition(
+                                "keyguard disabled"
+                            )
+                        }
                     }
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardStateCallbackInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardStateCallbackInteractor.kt
index 420fbd4..7fd348b 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardStateCallbackInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardStateCallbackInteractor.kt
@@ -20,6 +20,7 @@
 import android.os.RemoteException
 import com.android.internal.policy.IKeyguardStateCallback
 import com.android.systemui.CoreStartable
+import com.android.systemui.bouncer.domain.interactor.SimBouncerInteractor
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
@@ -27,13 +28,13 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
+import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
-import javax.inject.Inject
 
 /**
  * Updates KeyguardStateCallbacks provided to KeyguardService with KeyguardTransitionInteractor
@@ -50,6 +51,8 @@
     @Background private val backgroundDispatcher: CoroutineDispatcher,
     private val selectedUserInteractor: SelectedUserInteractor,
     private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
+    private val trustInteractor: TrustInteractor,
+    private val simBouncerInteractor: SimBouncerInteractor,
 ) : CoreStartable {
     private val callbacks = mutableListOf<IKeyguardStateCallback>()
 
@@ -60,21 +63,19 @@
 
         applicationScope.launch {
             combine(
-                selectedUserInteractor.selectedUser,
-                keyguardTransitionInteractor.currentKeyguardState,
-                ::Pair
-            ).collectLatest { (selectedUser, currentState) ->
-                val iterator = callbacks.iterator()
+                    selectedUserInteractor.selectedUser,
+                    keyguardTransitionInteractor.currentKeyguardState,
+                    keyguardTransitionInteractor.startedKeyguardTransitionStep,
+                    ::Triple,
+                )
+                .collectLatest { (selectedUser, _, _) ->
+                    val iterator = callbacks.iterator()
                     withContext(backgroundDispatcher) {
                         while (iterator.hasNext()) {
                             val callback = iterator.next()
                             try {
-                                 callback.onShowingStateChanged(
-                                    currentState != KeyguardState.GONE,
-                                    selectedUser
-                                )
-                                callback.onInputRestrictedStateChanged(
-                                    currentState != KeyguardState.GONE)
+                                callback.onShowingStateChanged(!isIdleInGone(), selectedUser)
+                                callback.onInputRestrictedStateChanged(!isIdleInGone())
                             } catch (e: RemoteException) {
                                 if (e is DeadObjectException) {
                                     iterator.remove()
@@ -84,10 +85,58 @@
                     }
                 }
         }
+
+        applicationScope.launch {
+            trustInteractor.isTrusted.collectLatest { isTrusted ->
+                val iterator = callbacks.iterator()
+                withContext(backgroundDispatcher) {
+                    while (iterator.hasNext()) {
+                        val callback = iterator.next()
+                        try {
+                            callback.onTrustedChanged(isTrusted)
+                        } catch (e: RemoteException) {
+                            if (e is DeadObjectException) {
+                                iterator.remove()
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        applicationScope.launch {
+            simBouncerInteractor.isAnySimSecure.collectLatest { isSimSecured ->
+                val iterator = callbacks.iterator()
+                withContext(backgroundDispatcher) {
+                    while (iterator.hasNext()) {
+                        val callback = iterator.next()
+                        try {
+                            callback.onSimSecureStateChanged(isSimSecured)
+                        } catch (e: RemoteException) {
+                            if (e is DeadObjectException) {
+                                iterator.remove()
+                            }
+                        }
+                    }
+                }
+            }
+        }
     }
 
     fun addCallback(callback: IKeyguardStateCallback) {
         KeyguardWmStateRefactor.isUnexpectedlyInLegacyMode()
         callbacks.add(callback)
+
+        // Send initial values to new callbacks.
+        callback.onShowingStateChanged(!isIdleInGone(), selectedUserInteractor.getSelectedUserId())
+        callback.onInputRestrictedStateChanged(!isIdleInGone())
+        callback.onTrustedChanged(trustInteractor.isTrusted.value)
+        callback.onSimSecureStateChanged(simBouncerInteractor.isAnySimSecure.value)
     }
-}
\ No newline at end of file
+
+    /** Whether we're in KeyguardState.GONE and haven't started a transition to another state. */
+    private fun isIdleInGone(): Boolean {
+        return keyguardTransitionInteractor.getCurrentState() == KeyguardState.GONE &&
+            keyguardTransitionInteractor.getStartedState() == KeyguardState.GONE
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
index e19b72e..c4f231d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
@@ -77,7 +77,7 @@
             MutableSharedFlow<Float>(
                     replay = 1,
                     extraBufferCapacity = 2,
-                    onBufferOverflow = BufferOverflow.DROP_OLDEST
+                    onBufferOverflow = BufferOverflow.DROP_OLDEST,
                 )
                 .also { it.tryEmit(0f) }
         }
@@ -97,8 +97,8 @@
                 SharingStarted.Eagerly,
                 WithPrev(
                     sceneInteractor.transitionState.value,
-                    sceneInteractor.transitionState.value
-                )
+                    sceneInteractor.transitionState.value,
+                ),
             )
 
     /**
@@ -156,7 +156,7 @@
                     Log.e(
                         TAG,
                         "STARTED step ($startedStep) was preceded by a RUNNING step " +
-                            "($prevStep), which should never happen. Things could go badly here."
+                            "($prevStep), which should never happen. Things could go badly here.",
                     )
                 }
             }
@@ -202,7 +202,7 @@
             transitionMap.getOrPut(mappedEdge) {
                 MutableSharedFlow(
                     extraBufferCapacity = 10,
-                    onBufferOverflow = BufferOverflow.DROP_OLDEST
+                    onBufferOverflow = BufferOverflow.DROP_OLDEST,
                 )
             }
 
@@ -262,7 +262,7 @@
             is Edge.StateToState ->
                 Edge.create(
                     from = edge.from?.mapToSceneContainerState(),
-                    to = edge.to?.mapToSceneContainerState()
+                    to = edge.to?.mapToSceneContainerState(),
                 )
             is Edge.SceneToState -> Edge.create(UNDEFINED, edge.to)
             is Edge.StateToScene -> Edge.create(edge.from, UNDEFINED)
@@ -286,9 +286,7 @@
      * The value will be `0` (or close to `0`, due to float point arithmetic) if not in this step or
      * `1` when fully in the given state.
      */
-    fun transitionValue(
-        state: KeyguardState,
-    ): Flow<Float> {
+    fun transitionValue(state: KeyguardState): Flow<Float> {
         if (SceneContainerFlag.isEnabled && state != state.mapToSceneContainerState()) {
             Log.e(TAG, "SceneContainer is enabled but a deprecated state $state is used.")
             return transitionValue(state.mapToSceneContainerScene()!!, state)
@@ -369,10 +367,9 @@
             .stateIn(scope, SharingStarted.Eagerly, OFF)
 
     val isInTransition =
-        combine(
-            isInTransitionWhere({ true }, { true }),
-            sceneInteractor.transitionState,
-        ) { isKeyguardTransitioning, sceneTransitionState ->
+        combine(isInTransitionWhere({ true }, { true }), sceneInteractor.transitionState) {
+            isKeyguardTransitioning,
+            sceneTransitionState ->
             isKeyguardTransitioning ||
                 (SceneContainerFlag.isEnabled && sceneTransitionState.isTransitioning())
         }
@@ -465,6 +462,10 @@
         return currentKeyguardState.replayCache.last()
     }
 
+    fun getStartedState(): KeyguardState {
+        return startedKeyguardTransitionStep.value.to
+    }
+
     private val finishedKeyguardState: StateFlow<KeyguardState> =
         repository.transitions
             .filter { it.transitionState == TransitionState.FINISHED }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModel.kt
index ecae079..3b266f9 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModel.kt
@@ -21,13 +21,11 @@
 import com.android.compose.animation.scene.Edge
 import com.android.compose.animation.scene.Swipe
 import com.android.compose.animation.scene.SwipeDirection
-import com.android.compose.animation.scene.TransitionKey
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
 import com.android.systemui.scene.shared.model.Overlays
-import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.TransitionKeys.ToSplitShade
 import com.android.systemui.scene.ui.viewmodel.SceneContainerEdge
@@ -71,9 +69,8 @@
 
                             addAll(
                                 when (shadeMode) {
-                                    ShadeMode.Single -> fullscreenShadeActions()
-                                    ShadeMode.Split ->
-                                        fullscreenShadeActions(transitionKey = ToSplitShade)
+                                    ShadeMode.Single -> singleShadeActions()
+                                    ShadeMode.Split -> splitShadeActions()
                                     ShadeMode.Dual -> dualShadeActions()
                                 }
                             )
@@ -84,18 +81,26 @@
             .collect { setActions(it) }
     }
 
-    private fun fullscreenShadeActions(
-        transitionKey: TransitionKey? = null
-    ): Array<Pair<UserAction, UserActionResult>> {
-        val notifShadeSceneKey = UserActionResult(SceneFamilies.NotifShade, transitionKey)
-        val qsShadeSceneKey = UserActionResult(SceneFamilies.QuickSettings, transitionKey)
+    private fun singleShadeActions(): Array<Pair<UserAction, UserActionResult>> {
         return arrayOf(
             // Swiping down, not from the edge, always goes to shade.
-            Swipe.Down to notifShadeSceneKey,
-            swipeDown(pointerCount = 2) to notifShadeSceneKey,
+            Swipe.Down to Scenes.Shade,
+            swipeDown(pointerCount = 2) to Scenes.Shade,
             // Swiping down from the top edge goes to QS.
-            swipeDownFromTop(pointerCount = 1) to qsShadeSceneKey,
-            swipeDownFromTop(pointerCount = 2) to qsShadeSceneKey,
+            swipeDownFromTop(pointerCount = 1) to Scenes.QuickSettings,
+            swipeDownFromTop(pointerCount = 2) to Scenes.QuickSettings,
+        )
+    }
+
+    private fun splitShadeActions(): Array<Pair<UserAction, UserActionResult>> {
+        val splitShadeSceneKey = UserActionResult(Scenes.Shade, ToSplitShade)
+        return arrayOf(
+            // Swiping down, not from the edge, always goes to shade.
+            Swipe.Down to splitShadeSceneKey,
+            swipeDown(pointerCount = 2) to splitShadeSceneKey,
+            // Swiping down from the top edge goes to QS.
+            swipeDownFromTop(pointerCount = 1) to splitShadeSceneKey,
+            swipeDownFromTop(pointerCount = 2) to splitShadeSceneKey,
         )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/permission/MediaProjectionPermissionActivity.java b/packages/SystemUI/src/com/android/systemui/mediaprojection/permission/MediaProjectionPermissionActivity.java
index 4251b81..d596589 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/permission/MediaProjectionPermissionActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/permission/MediaProjectionPermissionActivity.java
@@ -53,6 +53,7 @@
 import android.text.TextUtils;
 import android.util.Log;
 import android.view.Window;
+import android.view.WindowManager;
 
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
@@ -308,6 +309,9 @@
     private void setUpDialog(AlertDialog dialog) {
         SystemUIDialog.registerDismissListener(dialog);
         SystemUIDialog.applyFlags(dialog, /* showWhenLocked= */ false);
+
+        final Window w = dialog.getWindow();
+        w.setType(WindowManager.LayoutParams.TYPE_SYSTEM_ALERT);
         SystemUIDialog.setDialogSize(dialog);
 
         dialog.setOnCancelListener(this::onDialogDismissedOrCancelled);
@@ -315,7 +319,6 @@
         dialog.create();
         dialog.getButton(DialogInterface.BUTTON_POSITIVE).setFilterTouchesWhenObscured(true);
 
-        final Window w = dialog.getWindow();
         w.addSystemFlags(SYSTEM_FLAG_HIDE_NON_SYSTEM_OVERLAY_WINDOWS);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/model/SceneContainerPlugin.kt b/packages/SystemUI/src/com/android/systemui/model/SceneContainerPlugin.kt
index db5a545..0d748a1 100644
--- a/packages/SystemUI/src/com/android/systemui/model/SceneContainerPlugin.kt
+++ b/packages/SystemUI/src/com/android/systemui/model/SceneContainerPlugin.kt
@@ -17,11 +17,13 @@
 package com.android.systemui.model
 
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.scene.domain.interactor.SceneContainerOcclusionInteractor
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BOUNCER_SHOWING
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_COMMUNAL_HUB_SHOWING
@@ -57,12 +59,12 @@
 
         val transitionState = sceneInteractor.get().transitionState.value
         val idleTransitionStateOrNull = transitionState as? ObservableTransitionState.Idle
-        val currentSceneOrNull = idleTransitionStateOrNull?.currentScene
         val invisibleDueToOcclusion = occlusionInteractor.get().invisibleDueToOcclusion.value
-        return currentSceneOrNull?.let { sceneKey ->
+        return idleTransitionStateOrNull?.let { idleState ->
             EvaluatorByFlag[flag]?.invoke(
                 SceneContainerPluginState(
-                    scene = sceneKey,
+                    scene = idleState.currentScene,
+                    overlays = idleState.currentOverlays,
                     invisibleDueToOcclusion = invisibleDueToOcclusion,
                 )
             )
@@ -89,10 +91,15 @@
                             it.invisibleDueToOcclusion -> false
                             it.scene == Scenes.Lockscreen -> true
                             it.scene == Scenes.Shade -> true
+                            Overlays.NotificationsShade in it.overlays -> true
                             else -> false
                         }
                     },
-                SYSUI_STATE_QUICK_SETTINGS_EXPANDED to { it.scene == Scenes.QuickSettings },
+                SYSUI_STATE_QUICK_SETTINGS_EXPANDED to
+                    {
+                        it.scene == Scenes.QuickSettings ||
+                            Overlays.QuickSettingsShade in it.overlays
+                    },
                 SYSUI_STATE_BOUNCER_SHOWING to { it.scene == Scenes.Bouncer },
                 SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING to
                     {
@@ -106,5 +113,9 @@
             )
     }
 
-    data class SceneContainerPluginState(val scene: SceneKey, val invisibleDueToOcclusion: Boolean)
+    data class SceneContainerPluginState(
+        val scene: SceneKey,
+        val overlays: Set<OverlayKey>,
+        val invisibleDueToOcclusion: Boolean,
+    )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModel.kt
index b6868c1..63bfbd1 100644
--- a/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModel.kt
@@ -18,9 +18,13 @@
 
 import com.android.compose.animation.scene.Back
 import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
+import com.android.compose.animation.scene.UserActionResult.HideOverlay
+import com.android.compose.animation.scene.UserActionResult.ReplaceByOverlay
 import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.viewmodel.SceneContainerEdge
 import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
@@ -32,8 +36,10 @@
     override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
         setActions(
             mapOf(
-                Swipe.Up to UserActionResult.HideOverlay(Overlays.NotificationsShade),
-                Back to UserActionResult.HideOverlay(Overlays.NotificationsShade),
+                Swipe.Up to HideOverlay(Overlays.NotificationsShade),
+                Back to HideOverlay(Overlays.NotificationsShade),
+                Swipe(direction = SwipeDirection.Down, fromSource = SceneContainerEdge.TopRight) to
+                    ReplaceByOverlay(Overlays.QuickSettingsShade),
             )
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModel.kt
index 61c4c8c..31519a9 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModel.kt
@@ -18,24 +18,41 @@
 
 import com.android.compose.animation.scene.Back
 import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.SwipeDirection
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
+import com.android.compose.animation.scene.UserActionResult.HideOverlay
+import com.android.compose.animation.scene.UserActionResult.ReplaceByOverlay
 import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.viewmodel.SceneContainerEdge
 import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
+import kotlinx.coroutines.flow.map
 
 /** Models the UI state for the user actions for navigating to other scenes or overlays. */
-class QuickSettingsShadeOverlayActionsViewModel @AssistedInject constructor() :
+class QuickSettingsShadeOverlayActionsViewModel
+@AssistedInject
+constructor(private val containerViewModel: QuickSettingsContainerViewModel) :
     UserActionsViewModel() {
 
     override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
-        setActions(
-            buildMap {
-                put(Swipe.Up, UserActionResult.HideOverlay(Overlays.QuickSettingsShade))
-                put(Back, UserActionResult.HideOverlay(Overlays.QuickSettingsShade))
+        containerViewModel.editModeViewModel.isEditing
+            .map { isEditing ->
+                buildMap {
+                    put(Swipe.Up, HideOverlay(Overlays.QuickSettingsShade))
+                    // When editing, back should go back to QS from edit mode (i.e. remain in the
+                    // same overlay).
+                    if (!isEditing) {
+                        put(Back, HideOverlay(Overlays.QuickSettingsShade))
+                    }
+                    put(
+                        Swipe(SwipeDirection.Down, fromSource = SceneContainerEdge.TopLeft),
+                        ReplaceByOverlay(Overlays.NotificationsShade),
+                    )
+                }
             }
-        )
+            .collect { actions -> setActions(actions) }
     }
 
     @AssistedFactory
diff --git a/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt b/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt
index 1aa982f..e441a23 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt
@@ -22,8 +22,6 @@
 import com.android.systemui.scene.domain.SceneDomainModule
 import com.android.systemui.scene.domain.interactor.WindowRootViewVisibilityInteractor
 import com.android.systemui.scene.domain.resolver.HomeSceneFamilyResolverModule
-import com.android.systemui.scene.domain.resolver.NotifShadeSceneFamilyResolverModule
-import com.android.systemui.scene.domain.resolver.QuickSettingsSceneFamilyResolverModule
 import com.android.systemui.scene.domain.startable.KeyguardStateCallbackStartable
 import com.android.systemui.scene.domain.startable.SceneContainerStartable
 import com.android.systemui.scene.domain.startable.ScrimStartable
@@ -47,7 +45,6 @@
             EmptySceneModule::class,
             GoneSceneModule::class,
             NotificationsShadeOverlayModule::class,
-            NotificationsShadeSceneModule::class,
             NotificationsShadeSessionModule::class,
             QuickSettingsShadeOverlayModule::class,
             QuickSettingsSceneModule::class,
@@ -56,8 +53,6 @@
 
             // List SceneResolver modules for supported SceneFamilies
             HomeSceneFamilyResolverModule::class,
-            NotifShadeSceneFamilyResolverModule::class,
-            QuickSettingsSceneFamilyResolverModule::class,
         ]
 )
 interface KeyguardlessSceneContainerFrameworkModule {
diff --git a/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt b/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
index 7f0cf86..a89f752 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
@@ -22,8 +22,6 @@
 import com.android.systemui.scene.domain.SceneDomainModule
 import com.android.systemui.scene.domain.interactor.WindowRootViewVisibilityInteractor
 import com.android.systemui.scene.domain.resolver.HomeSceneFamilyResolverModule
-import com.android.systemui.scene.domain.resolver.NotifShadeSceneFamilyResolverModule
-import com.android.systemui.scene.domain.resolver.QuickSettingsSceneFamilyResolverModule
 import com.android.systemui.scene.domain.startable.KeyguardStateCallbackStartable
 import com.android.systemui.scene.domain.startable.SceneContainerStartable
 import com.android.systemui.scene.domain.startable.ScrimStartable
@@ -52,16 +50,12 @@
             QuickSettingsSceneModule::class,
             ShadeSceneModule::class,
             QuickSettingsShadeOverlayModule::class,
-            QuickSettingsShadeSceneModule::class,
             NotificationsShadeOverlayModule::class,
-            NotificationsShadeSceneModule::class,
             NotificationsShadeSessionModule::class,
             SceneDomainModule::class,
 
             // List SceneResolver modules for supported SceneFamilies
             HomeSceneFamilyResolverModule::class,
-            NotifShadeSceneFamilyResolverModule::class,
-            QuickSettingsSceneFamilyResolverModule::class,
         ]
 )
 interface SceneContainerFrameworkModule {
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/resolver/NotifShadeSceneFamilyResolver.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/resolver/NotifShadeSceneFamilyResolver.kt
deleted file mode 100644
index a313273..0000000
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/resolver/NotifShadeSceneFamilyResolver.kt
+++ /dev/null
@@ -1,71 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.scene.domain.resolver
-
-import com.android.compose.animation.scene.SceneKey
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.scene.shared.model.SceneFamilies
-import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.shade.domain.interactor.ShadeModeInteractor
-import com.android.systemui.shade.shared.model.ShadeMode
-import dagger.Binds
-import dagger.Module
-import dagger.multibindings.IntoSet
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.stateIn
-
-@SysUISingleton
-class NotifShadeSceneFamilyResolver
-@Inject
-constructor(
-    @Application applicationScope: CoroutineScope,
-    shadeModeInteractor: ShadeModeInteractor,
-) : SceneResolver {
-    override val targetFamily: SceneKey = SceneFamilies.NotifShade
-
-    override val resolvedScene: StateFlow<SceneKey> =
-        shadeModeInteractor.shadeMode
-            .map(::notifShadeScene)
-            .stateIn(
-                applicationScope,
-                started = SharingStarted.Eagerly,
-                initialValue = notifShadeScene(shadeModeInteractor.shadeMode.value),
-            )
-
-    override fun includesScene(scene: SceneKey): Boolean = scene in notifShadeScenes
-
-    private fun notifShadeScene(shadeMode: ShadeMode) =
-        when (shadeMode) {
-            is ShadeMode.Single -> Scenes.Shade
-            is ShadeMode.Dual -> Scenes.NotificationsShade
-            is ShadeMode.Split -> Scenes.Shade
-        }
-
-    companion object {
-        val notifShadeScenes = setOf(Scenes.NotificationsShade, Scenes.Shade)
-    }
-}
-
-@Module
-interface NotifShadeSceneFamilyResolverModule {
-    @Binds @IntoSet fun bindResolver(interactor: NotifShadeSceneFamilyResolver): SceneResolver
-}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/resolver/QuickSettingsSceneFamilyResolver.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/resolver/QuickSettingsSceneFamilyResolver.kt
deleted file mode 100644
index 923e712a..0000000
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/resolver/QuickSettingsSceneFamilyResolver.kt
+++ /dev/null
@@ -1,72 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.scene.domain.resolver
-
-import com.android.compose.animation.scene.SceneKey
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.scene.shared.model.SceneFamilies
-import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.shade.domain.interactor.ShadeModeInteractor
-import com.android.systemui.shade.shared.model.ShadeMode
-import dagger.Binds
-import dagger.Module
-import dagger.multibindings.IntoSet
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.stateIn
-
-@SysUISingleton
-class QuickSettingsSceneFamilyResolver
-@Inject
-constructor(
-    @Application applicationScope: CoroutineScope,
-    shadeModeInteractor: ShadeModeInteractor,
-) : SceneResolver {
-    override val targetFamily: SceneKey = SceneFamilies.QuickSettings
-
-    override val resolvedScene: StateFlow<SceneKey> =
-        shadeModeInteractor.shadeMode
-            .map(::quickSettingsScene)
-            .stateIn(
-                applicationScope,
-                started = SharingStarted.Eagerly,
-                initialValue = quickSettingsScene(shadeModeInteractor.shadeMode.value),
-            )
-
-    override fun includesScene(scene: SceneKey): Boolean = scene in quickSettingsScenes
-
-    private fun quickSettingsScene(shadeMode: ShadeMode) =
-        when (shadeMode) {
-            is ShadeMode.Single -> Scenes.QuickSettings
-            is ShadeMode.Dual -> Scenes.QuickSettingsShade
-            is ShadeMode.Split -> Scenes.Shade
-        }
-
-    companion object {
-        val quickSettingsScenes =
-            setOf(Scenes.QuickSettings, Scenes.QuickSettingsShade, Scenes.Shade)
-    }
-}
-
-@Module
-interface QuickSettingsSceneFamilyResolverModule {
-    @Binds @IntoSet fun bindResolver(interactor: QuickSettingsSceneFamilyResolver): SceneResolver
-}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
index c5e0ccc..e11ffcc 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/SceneContainerStartable.kt
@@ -107,6 +107,7 @@
  * Hooks up business logic that manipulates the state of the [SceneInteractor] for the system UI
  * scene container based on state from other systems.
  */
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class SceneContainerStartable
 @Inject
@@ -408,8 +409,7 @@
                         }
                         isOnPrimaryBouncer -> {
                             // When the device becomes unlocked in primary Bouncer,
-                            // notify dismiss succeeded and
-                            // go to previous scene or Gone.
+                            // notify dismiss succeeded and go to previous scene or Gone.
                             dismissCallbackRegistry.notifyDismissSucceeded()
                             if (
                                 previousScene.value == Scenes.Lockscreen ||
@@ -596,12 +596,12 @@
             combine(
                     sceneInteractor.transitionState
                         .mapNotNull { it as? ObservableTransitionState.Idle }
-                        .map { it.currentScene }
                         .distinctUntilChanged(),
                     occlusionInteractor.invisibleDueToOcclusion,
-                ) { sceneKey, invisibleDueToOcclusion ->
+                ) { idleState, invisibleDueToOcclusion ->
                     SceneContainerPlugin.SceneContainerPluginState(
-                        scene = sceneKey,
+                        scene = idleState.currentScene,
+                        overlays = idleState.currentOverlays,
                         invisibleDueToOcclusion = invisibleDueToOcclusion,
                     )
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt
index 115d664..82b4b1c 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt
@@ -43,22 +43,6 @@
     @JvmField val Lockscreen = SceneKey("lockscreen")
 
     /**
-     * The notifications shade scene primarily shows a scrollable list of notifications as an
-     * overlay UI.
-     *
-     * It's used only in the dual shade configuration, where there are two separate shades: one for
-     * notifications (this scene) and another for [QuickSettingsShade].
-     *
-     * It's not used in the single/accordion configuration (swipe down once to reveal the shade,
-     * swipe down again the to expand quick settings) or in the "split" shade configuration (on
-     * large screens or unfolded foldables, where notifications and quick settings are shown
-     * side-by-side in their own columns).
-     */
-    @Deprecated("The notifications shade scene has been replaced by an overlay")
-    @JvmField
-    val NotificationsShade = SceneKey("notifications_shade")
-
-    /**
      * The quick settings scene shows the quick setting tiles.
      *
      * This scene is used for single/accordion configuration (swipe down once to reveal the shade,
@@ -69,27 +53,12 @@
      * scene is used.
      *
      * For the dual shade configuration, where there are two separate shades: one for notifications
-     * and one for quick settings, [NotificationsShade] and [QuickSettingsShade] scenes are used
-     * respectively.
+     * and one for quick settings, the overlays `NotificationsShade` and `QuickSettingsShade` are
+     * used respectively.
      */
     @JvmField val QuickSettings = SceneKey("quick_settings")
 
     /**
-     * The quick settings shade scene shows the quick setting tiles as an overlay UI.
-     *
-     * It's used only in the dual shade configuration, where there are two separate shades: one for
-     * quick settings (this scene) and another for [NotificationsShade].
-     *
-     * It's not used in the single/accordion configuration (swipe down once to reveal the shade,
-     * swipe down again the to expand quick settings) or in the "split" shade configuration (on
-     * large screens or unfolded foldables, where notifications and quick settings are shown
-     * side-by-side in their own columns).
-     */
-    @Deprecated("The quick settings shade scene has been replaced by an overlay")
-    @JvmField
-    val QuickSettingsShade = SceneKey("quick_settings_shade")
-
-    /**
      * The shade is the scene that shows a scrollable list of notifications and the minimized
      * version of quick settings (AKA "quick quick settings" or "QQS").
      *
@@ -97,7 +66,8 @@
      * swipe down again the to expand quick settings) and for the "split" shade configuration (on
      * large screens or unfolded foldables, where notifications and quick settings are shown
      * side-by-side in their own columns). For the dual shade configuration, where there are two
-     * separate shades: one for notifications and one for quick settings, other scenes are used.
+     * separate shades: one for notifications and one for quick settings, the overlays
+     * `NotificationsShade` and `QuickSettingsShade` are used respectively.
      */
     @JvmField val Shade = SceneKey("shade")
 }
@@ -114,16 +84,4 @@
      * depending on whether the device is unlocked and has been entered.
      */
     @JvmField val Home = SceneKey(debugName = "scene_family_home")
-
-    /**
-     * The scene that contains the full, interactive notification shade. The specific scene it
-     * resolves to can depend on dual / split / single shade settings.
-     */
-    @JvmField val NotifShade = SceneKey(debugName = "scene_family_notif_shade")
-
-    /**
-     * The scene that contains the full QuickSettings (not to be confused with Quick-QuickSettings).
-     * The specific scene it resolves to can depend on dual / split / single shade settings.
-     */
-    @JvmField val QuickSettings = SceneKey(debugName = "scene_family_quick_settings")
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/footer/ui/viewbinder/FooterViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/footer/ui/viewbinder/FooterViewBinder.kt
index 637cadd..920541d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/footer/ui/viewbinder/FooterViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/footer/ui/viewbinder/FooterViewBinder.kt
@@ -44,7 +44,7 @@
                     viewModel,
                     clearAllNotifications,
                     launchNotificationSettings,
-                    launchNotificationHistory
+                    launchNotificationHistory,
                 )
             }
         }
@@ -55,21 +55,15 @@
         viewModel: FooterViewModel,
         clearAllNotifications: View.OnClickListener,
         launchNotificationSettings: View.OnClickListener,
-        launchNotificationHistory: View.OnClickListener
+        launchNotificationHistory: View.OnClickListener,
     ) = coroutineScope {
-        launch {
-            bindClearAllButton(
-                footer,
-                viewModel,
-                clearAllNotifications,
-            )
-        }
+        launch { bindClearAllButton(footer, viewModel, clearAllNotifications) }
         launch {
             bindManageOrHistoryButton(
                 footer,
                 viewModel,
                 launchNotificationSettings,
-                launchNotificationHistory
+                launchNotificationHistory,
             )
         }
         launch { bindMessage(footer, viewModel) }
@@ -80,8 +74,6 @@
         viewModel: FooterViewModel,
         clearAllNotifications: View.OnClickListener,
     ) = coroutineScope {
-        footer.setClearAllButtonClickListener(clearAllNotifications)
-
         launch {
             viewModel.clearAllButton.labelId.collect { textId ->
                 footer.setClearAllButtonText(textId)
@@ -96,18 +88,21 @@
 
         launch {
             viewModel.clearAllButton.isVisible.collect { isVisible ->
+                if (isVisible.value) {
+                    footer.setClearAllButtonClickListener(clearAllNotifications)
+                } else {
+                    // When the button isn't visible, it also shouldn't react to clicks. This is
+                    // necessary because when the clear all button is not visible, it's actually
+                    // just the alpha that becomes 0 so it can still be tapped.
+                    footer.setClearAllButtonClickListener(null)
+                }
+
                 if (isVisible.isAnimating) {
-                    footer.setClearAllButtonVisible(
-                        isVisible.value,
-                        /* animate = */ true,
-                    ) { _ ->
+                    footer.setClearAllButtonVisible(isVisible.value, /* animate= */ true) { _ ->
                         isVisible.stopAnimating()
                     }
                 } else {
-                    footer.setClearAllButtonVisible(
-                        isVisible.value,
-                        /* animate = */ false,
-                    )
+                    footer.setClearAllButtonVisible(isVisible.value, /* animate= */ false)
                 }
             }
         }
@@ -143,22 +138,24 @@
 
         launch {
             viewModel.manageOrHistoryButton.isVisible.collect { isVisible ->
-                // NOTE: This visibility change is never animated.
+                // NOTE: This visibility change is never animated. We also don't need to do anything
+                // special about the onClickListener here, since we're changing the visibility to
+                // GONE so it won't be clickable anyway.
                 footer.setManageOrHistoryButtonVisible(isVisible.value)
             }
         }
     }
 
-    private suspend fun bindMessage(
-        footer: FooterView,
-        viewModel: FooterViewModel,
-    ) = coroutineScope {
-        // Bind the resource IDs
-        footer.setMessageString(viewModel.message.messageId)
-        footer.setMessageIcon(viewModel.message.iconId)
+    private suspend fun bindMessage(footer: FooterView, viewModel: FooterViewModel) =
+        coroutineScope {
+            // Bind the resource IDs
+            footer.setMessageString(viewModel.message.messageId)
+            footer.setMessageIcon(viewModel.message.iconId)
 
-        launch {
-            viewModel.message.isVisible.collect { visible -> footer.setFooterLabelVisible(visible) }
+            launch {
+                viewModel.message.isVisible.collect { visible ->
+                    footer.setFooterLabelVisible(visible)
+                }
+            }
         }
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/RotationChangeProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/RotationChangeProviderTest.kt
index f5bcc21..feee0a3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/RotationChangeProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/RotationChangeProviderTest.kt
@@ -38,6 +38,7 @@
 import org.mockito.Mock
 import org.mockito.Mockito.spy
 import org.mockito.Mockito.verify
+import org.mockito.Mockito.inOrder
 import org.mockito.Mockito.verifyNoMoreInteractions
 import org.mockito.MockitoAnnotations
 
@@ -79,6 +80,26 @@
     }
 
     @Test
+    fun onRotationChanged_rotationSentMultipleWithTheSameValue_listenerReceivesUpdateOnce() {
+        sendRotationUpdate(42)
+        sendRotationUpdate(42)
+        sendRotationUpdate(42)
+
+        verify(listener).onRotationChanged(42)
+    }
+
+    @Test
+    fun onRotationChanged_rotationSentMultipleTimesWithDifferentValues_listenerReceivesUpdates() {
+        sendRotationUpdate(0)
+        sendRotationUpdate(1)
+
+        with(inOrder(listener)) {
+            verify(listener).onRotationChanged(0)
+            verify(listener).onRotationChanged(1)
+        }
+    }
+
+    @Test
     fun onRotationChanged_subscribersRemoved_noRotationChangeReceived() {
         sendRotationUpdate(42)
         verify(listener).onRotationChanged(42)
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorKosmos.kt
index d920c4f..574bbcd 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissActionInteractorKosmos.kt
@@ -17,14 +17,12 @@
 package com.android.systemui.keyguard.domain.interactor
 
 import com.android.systemui.bouncer.domain.interactor.alternateBouncerInteractor
-import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
 import com.android.systemui.keyguard.data.repository.keyguardRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.domain.interactor.sceneInteractor
-import com.android.systemui.scene.domain.resolver.notifShadeSceneFamilyResolver
-import com.android.systemui.scene.domain.resolver.quickSettingsSceneFamilyResolver
 import com.android.systemui.shade.domain.interactor.shadeInteractor
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 
@@ -37,12 +35,9 @@
             dismissInteractor = keyguardDismissInteractor,
             applicationScope = testScope.backgroundScope,
             sceneInteractor = { sceneInteractor },
-            deviceEntryInteractor = { deviceEntryInteractor },
-            quickSettingsSceneFamilyResolver = { quickSettingsSceneFamilyResolver },
-            notifShadeSceneFamilyResolver = { notifShadeSceneFamilyResolver },
+            deviceUnlockedInteractor = { deviceUnlockedInteractor },
             powerInteractor = powerInteractor,
             alternateBouncerInteractor = alternateBouncerInteractor,
-            keyguardInteractor = { keyguardInteractor },
             shadeInteractor = { shadeInteractor },
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelKosmos.kt
index 8fc40e4..6ced8c3 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelKosmos.kt
@@ -18,8 +18,12 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.qs.ui.adapter.FakeQSSceneAdapter
+import com.android.systemui.util.mockito.mock
+
+val Kosmos.fakeQsSceneAdapter: FakeQSSceneAdapter by Fixture { FakeQSSceneAdapter({ mock() }) }
 
 val Kosmos.quickSettingsShadeOverlayActionsViewModel:
     QuickSettingsShadeOverlayActionsViewModel by Fixture {
-    QuickSettingsShadeOverlayActionsViewModel()
+    QuickSettingsShadeOverlayActionsViewModel(quickSettingsContainerViewModel)
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/resolver/SceneResolverKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/resolver/SceneResolverKosmos.kt
index d17b575..8b12425 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/resolver/SceneResolverKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/resolver/SceneResolverKosmos.kt
@@ -24,16 +24,10 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.scene.shared.model.SceneFamilies
-import com.android.systemui.shade.domain.interactor.shadeModeInteractor
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 val Kosmos.sceneFamilyResolvers: Map<SceneKey, SceneResolver>
-    get() =
-        mapOf(
-            SceneFamilies.Home to homeSceneFamilyResolver,
-            SceneFamilies.NotifShade to notifShadeSceneFamilyResolver,
-            SceneFamilies.QuickSettings to quickSettingsSceneFamilyResolver,
-        )
+    get() = mapOf(SceneFamilies.Home to homeSceneFamilyResolver)
 
 val Kosmos.homeSceneFamilyResolver by
     Kosmos.Fixture {
@@ -43,19 +37,3 @@
             keyguardEnabledInteractor = keyguardEnabledInteractor,
         )
     }
-
-val Kosmos.notifShadeSceneFamilyResolver by
-    Kosmos.Fixture {
-        NotifShadeSceneFamilyResolver(
-            applicationScope = applicationCoroutineScope,
-            shadeModeInteractor = shadeModeInteractor,
-        )
-    }
-
-val Kosmos.quickSettingsSceneFamilyResolver by
-    Kosmos.Fixture {
-        QuickSettingsSceneFamilyResolver(
-            applicationScope = applicationCoroutineScope,
-            shadeModeInteractor = shadeModeInteractor,
-        )
-    }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt
index 0458f53..b1e6fe2 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/RotationChangeProvider.kt
@@ -103,7 +103,7 @@
 
                 if (displayId == display.displayId) {
                     val currentRotation = display.rotation
-                    if (lastRotation.compareAndSet(lastRotation.get(), currentRotation)) {
+                    if (lastRotation.getAndSet(currentRotation) != currentRotation) {
                         listeners.forEach { it.onRotationChanged(currentRotation) }
                     }
                 }
diff --git a/packages/services/CameraExtensionsProxy/src/com/android/cameraextensions/CameraExtensionsProxyService.java b/packages/services/CameraExtensionsProxy/src/com/android/cameraextensions/CameraExtensionsProxyService.java
index 136738f..acb74d3 100644
--- a/packages/services/CameraExtensionsProxy/src/com/android/cameraextensions/CameraExtensionsProxyService.java
+++ b/packages/services/CameraExtensionsProxy/src/com/android/cameraextensions/CameraExtensionsProxyService.java
@@ -1620,7 +1620,7 @@
                 }
                 ret.outputConfigs.add(entry);
             }
-            if (Flags.extension10Bit() && EFV_SUPPORTED) {
+            if (EFV_SUPPORTED) {
                 ret.colorSpace = sessionConfig.getColorSpace();
             } else {
                 ret.colorSpace = ColorSpaceProfiles.UNSPECIFIED;
@@ -2596,7 +2596,7 @@
 
     private static CameraOutputConfig getCameraOutputConfig(Camera2OutputConfigImpl output) {
         CameraOutputConfig ret = new CameraOutputConfig();
-        if (Flags.extension10Bit() && EFV_SUPPORTED) {
+        if (EFV_SUPPORTED) {
             ret.dynamicRangeProfile = output.getDynamicRangeProfile();
         } else {
             ret.dynamicRangeProfile = DynamicRangeProfiles.STANDARD;
diff --git a/ravenwood/Android.bp b/ravenwood/Android.bp
index d0ea474..d1a3bf9 100644
--- a/ravenwood/Android.bp
+++ b/ravenwood/Android.bp
@@ -341,8 +341,8 @@
 android_ravenwood_libgroup {
     name: "ravenwood-runtime",
     data: [
-        "framework-res",
-        "ravenwood-empty-res",
+        ":framework-res",
+        ":ravenwood-empty-res",
     ],
     libs: [
         "100-framework-minus-apex.ravenwood",
diff --git a/ravenwood/TEST_MAPPING b/ravenwood/TEST_MAPPING
index 7f9d9c2..f9b5d2c 100644
--- a/ravenwood/TEST_MAPPING
+++ b/ravenwood/TEST_MAPPING
@@ -7,6 +7,9 @@
     { "name": "RavenwoodMockitoTest_device" },
     { "name": "RavenwoodBivalentTest_device" },
 
+    { "name": "RavenwoodBivalentInstTest_nonself_inst" },
+    { "name": "RavenwoodBivalentInstTest_self_inst_device" },
+
     // The sysui tests should match vendor/unbundled_google/packages/SystemUIGoogle/TEST_MAPPING
     {
       "name": "SystemUIGoogleTests",
@@ -138,6 +141,14 @@
       "host": true
     },
     {
+      "name": "RavenwoodBivalentInstTest_nonself_inst",
+      "host": true
+    },
+    {
+      "name": "RavenwoodBivalentInstTest_self_inst",
+      "host": true
+    },
+    {
       "name": "RavenwoodBivalentTest",
       "host": true
     },
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
index 68472c1..7a6f9e3 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
@@ -108,7 +108,7 @@
         Log.v(TAG, "onBeforeInnerRunnerStart: description=" + description);
 
         // Prepare the environment before the inner runner starts.
-        RavenwoodRunnerState.forRunner(runner).enterTestClass(description);
+        runner.mState.enterTestClass(description);
     }
 
     /**
@@ -119,7 +119,7 @@
         Log.v(TAG, "onAfterInnerRunnerFinished: description=" + description);
 
         RavenwoodTestStats.getInstance().onClassFinished(description);
-        RavenwoodRunnerState.forRunner(runner).exitTestClass();
+        runner.mState.exitTestClass();
     }
 
     /**
@@ -133,10 +133,10 @@
 
         if (scope == Scope.Instance && order == Order.Outer) {
             // Start of a test method.
-            RavenwoodRunnerState.forRunner(runner).enterTestMethod(description);
+            runner.mState.enterTestMethod(description);
         }
 
-        final var classDescription = RavenwoodRunnerState.forRunner(runner).getClassDescription();
+        final var classDescription = runner.mState.getClassDescription();
 
         // Class-level annotations are checked by the runner already, so we only check
         // method-level annotations here.
@@ -160,11 +160,11 @@
             Scope scope, Order order, Throwable th) {
         Log.v(TAG, "onAfter: description=" + description + ", " + scope + ", " + order + ", " + th);
 
-        final var classDescription = RavenwoodRunnerState.forRunner(runner).getClassDescription();
+        final var classDescription = runner.mState.getClassDescription();
 
         if (scope == Scope.Instance && order == Order.Outer) {
             // End of a test method.
-            RavenwoodRunnerState.forRunner(runner).exitTestMethod();
+            runner.mState.exitTestMethod();
             RavenwoodTestStats.getInstance().onTestFinished(classDescription, description,
                     th == null ? Result.Passed : Result.Failed);
         }
@@ -214,7 +214,7 @@
             Description description, RavenwoodRule rule) throws Throwable {
         Log.v(TAG, "onRavenwoodRuleEnter: description=" + description);
 
-        RavenwoodRunnerState.forRunner(runner).enterRavenwoodRule(rule);
+        runner.mState.enterRavenwoodRule(rule);
     }
 
 
@@ -225,6 +225,6 @@
             Description description, RavenwoodRule rule) throws Throwable {
         Log.v(TAG, "onRavenwoodRuleExit: description=" + description);
 
-        RavenwoodRunnerState.forRunner(runner).exitRavenwoodRule(rule);
+        runner.mState.exitRavenwoodRule(rule);
     }
 }
\ No newline at end of file
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodConfigState.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodConfigState.java
new file mode 100644
index 0000000..3535cb2
--- /dev/null
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodConfigState.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.platform.test.ravenwood;
+
+import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_EMPTY_RESOURCES_APK;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import android.annotation.Nullable;
+import android.app.ResourcesManager;
+import android.content.res.Resources;
+import android.view.DisplayAdjustments;
+
+import java.io.File;
+import java.util.HashMap;
+
+/**
+ * Used to store various states associated with {@link RavenwoodConfig} that's inly needed
+ * in junit-impl.
+ *
+ * We don't want to put it in junit-src to avoid having to recompile all the downstream
+ * dependencies after changing this class.
+ *
+ * All members must be called from the runner's main thread.
+ */
+public class RavenwoodConfigState {
+    private static final String TAG = "RavenwoodConfigState";
+
+    private final RavenwoodConfig mConfig;
+
+    public RavenwoodConfigState(RavenwoodConfig config) {
+        mConfig = config;
+    }
+
+    /** Map from path -> resources. */
+    private final HashMap<File, Resources> mCachedResources = new HashMap<>();
+
+    /**
+     * Load {@link Resources} from an APK, with cache.
+     */
+    public Resources loadResources(@Nullable File apkPath) {
+        var cached = mCachedResources.get(apkPath);
+        if (cached != null) {
+            return cached;
+        }
+
+        var fileToLoad = apkPath != null ? apkPath : new File(RAVENWOOD_EMPTY_RESOURCES_APK);
+
+        assertTrue("File " + fileToLoad + " doesn't exist.", fileToLoad.isFile());
+
+        final String path = fileToLoad.getAbsolutePath();
+        final var emptyPaths = new String[0];
+
+        ResourcesManager.getInstance().initializeApplicationPaths(path, emptyPaths);
+
+        final var ret = ResourcesManager.getInstance().getResources(null, path,
+                emptyPaths, emptyPaths, emptyPaths,
+                emptyPaths, null, null,
+                new DisplayAdjustments().getCompatibilityInfo(),
+                RavenwoodRuntimeEnvironmentController.class.getClassLoader(), null);
+
+        assertNotNull(ret);
+
+        mCachedResources.put(apkPath, ret);
+        return ret;
+    }
+}
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodContext.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodContext.java
index 48bed79..239c806 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodContext.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodContext.java
@@ -60,6 +60,8 @@
     private final File mCacheDir;
     private final Supplier<Resources> mResourcesSupplier;
 
+    private RavenwoodContext mAppContext;
+
     @GuardedBy("mLock")
     private Resources mResources;
 
@@ -77,8 +79,8 @@
         mPackageName = packageName;
         mMainThread = mainThread;
         mResourcesSupplier = resourcesSupplier;
-        mFilesDir = createTempDir("files-dir");
-        mCacheDir = createTempDir("cache-dir");
+        mFilesDir = createTempDir(packageName + "_files-dir");
+        mCacheDir = createTempDir(packageName + "_cache-dir");
 
         // Services provided by a typical shipping device
         registerService(ClipboardManager.class,
@@ -131,34 +133,35 @@
     @Override
     public Looper getMainLooper() {
         Objects.requireNonNull(mMainThread,
-                "Test must request setProvideMainThread() via RavenwoodRule");
+                "Test must request setProvideMainThread() via RavenwoodConfig");
         return mMainThread.getLooper();
     }
 
     @Override
     public Handler getMainThreadHandler() {
         Objects.requireNonNull(mMainThread,
-                "Test must request setProvideMainThread() via RavenwoodRule");
+                "Test must request setProvideMainThread() via RavenwoodConfig");
         return mMainThread.getThreadHandler();
     }
 
     @Override
     public Executor getMainExecutor() {
         Objects.requireNonNull(mMainThread,
-                "Test must request setProvideMainThread() via RavenwoodRule");
+                "Test must request setProvideMainThread() via RavenwoodConfig");
         return mMainThread.getThreadExecutor();
     }
 
     @Override
     public String getPackageName() {
         return Objects.requireNonNull(mPackageName,
-                "Test must request setPackageName() via RavenwoodRule");
+                "Test must request setPackageName() (or setTargetPackageName())"
+                + " via RavenwoodConfig");
     }
 
     @Override
     public String getOpPackageName() {
         return Objects.requireNonNull(mPackageName,
-                "Test must request setPackageName() via RavenwoodRule");
+                "Test must request setPackageName() via RavenwoodConfig");
     }
 
     @Override
@@ -227,6 +230,15 @@
         return new File(RAVENWOOD_RESOURCE_APK).getAbsolutePath();
     }
 
+    public void setApplicationContext(RavenwoodContext appContext) {
+        mAppContext = appContext;
+    }
+
+    @Override
+    public Context getApplicationContext() {
+        return mAppContext;
+    }
+
     /**
      * Wrap the given {@link Supplier} to become memoized.
      *
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
index d73afd4..04b67c4 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
@@ -19,7 +19,6 @@
 
 import static org.junit.Assert.fail;
 
-import android.annotation.NonNull;
 import android.annotation.Nullable;
 
 import com.android.internal.annotations.GuardedBy;
@@ -35,12 +34,11 @@
 import java.util.WeakHashMap;
 
 /**
- * Used to store various states associated with the current test runner.
+ * Used to store various states associated with the current test runner that's inly needed
+ * in junit-impl.
  *
- * This class could be added to {@link RavenwoodAwareTestRunner} as a field, but we don't
- * want to put it in junit-src/ (for one, that'll cause all the downstream dependencies to be
- * rebuilt when this file is updated), so we manage it separately using a Map from each
- * {@link RavenwoodAwareTestRunner} instance to a {@link RavenwoodRunnerState}.
+ * We don't want to put it in junit-src to avoid having to recompile all the downstream
+ * dependencies after changing this class.
  *
  * All members must be called from the runner's main thread.
  */
@@ -51,23 +49,12 @@
     private static final WeakHashMap<RavenwoodAwareTestRunner, RavenwoodRunnerState> sStates =
             new WeakHashMap<>();
 
-    /**
-     * Get the instance for a given runner.
-     */
-    public static RavenwoodRunnerState forRunner(@NonNull RavenwoodAwareTestRunner runner) {
-        synchronized (sStates) {
-            var ret = sStates.get(runner);
-            if (ret == null) {
-                ret = new RavenwoodRunnerState(runner);
-                sStates.put(runner, ret);
-            }
-            return ret;
-        }
-    }
-
     private final RavenwoodAwareTestRunner mRunner;
 
-    private RavenwoodRunnerState(RavenwoodAwareTestRunner runner) {
+    /**
+     * Ctor.
+     */
+    public RavenwoodRunnerState(RavenwoodAwareTestRunner runner) {
         mRunner = runner;
     }
 
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
index 03c9001..f4756c5 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
@@ -16,12 +16,10 @@
 
 package android.platform.test.ravenwood;
 
-import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_EMPTY_RESOURCES_APK;
 import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_RESOURCE_APK;
 import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_VERBOSE_LOGGING;
 
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 import android.app.ActivityManager;
 import android.app.Instrumentation;
@@ -34,7 +32,6 @@
 import android.os.Looper;
 import android.os.ServiceManager;
 import android.util.Log;
-import android.view.DisplayAdjustments;
 
 import androidx.test.platform.app.InstrumentationRegistry;
 
@@ -163,31 +160,52 @@
             main = null;
         }
 
-        // TODO This should be integrated into LoadedApk
-        final Supplier<Resources> resourcesSupplier = () -> {
-            var resApkFile = new File(RAVENWOOD_RESOURCE_APK);
-            if (!resApkFile.isFile()) {
-                resApkFile = new File(RAVENWOOD_EMPTY_RESOURCES_APK);
-            }
-            assertTrue(resApkFile.isFile());
-            final String res = resApkFile.getAbsolutePath();
-            final var emptyPaths = new String[0];
+        final boolean isSelfInstrumenting =
+                Objects.equals(config.mTestPackageName, config.mTargetPackageName);
 
-            ResourcesManager.getInstance().initializeApplicationPaths(res, emptyPaths);
-
-            final var ret = ResourcesManager.getInstance().getResources(null, res,
-                    emptyPaths, emptyPaths, emptyPaths,
-                    emptyPaths, null, null,
-                    new DisplayAdjustments().getCompatibilityInfo(),
-                    RavenwoodRuntimeEnvironmentController.class.getClassLoader(), null);
-
-            assertNotNull(ret);
-            return ret;
+        // This will load the resources from the apk set to `resource_apk` in the build file.
+        // This is supposed to be the "target app"'s resources.
+        final Supplier<Resources> targetResourcesLoader = () -> {
+            var file = new File(RAVENWOOD_RESOURCE_APK);
+            return config.mState.loadResources(file.exists() ? file : null);
         };
+        // Set up test context's resources.
+        // If the target package name == test package name, then we use the main resources.
+        // Otherwise, we don't simulate loading resources from the test APK yet.
+        // (we need to add `test_resource_apk` to `android_ravenwood_test`)
+        final Supplier<Resources> testResourcesLoader;
+        if (isSelfInstrumenting) {
+            testResourcesLoader = targetResourcesLoader;
+        } else {
+            testResourcesLoader = () -> {
+                fail("Cannot load resources from the test context (yet)."
+                        + " Use target context's resources instead.");
+                return null; // unreachable.
+            };
+        }
 
-        config.mContext = new RavenwoodContext(config.mPackageName, main, resourcesSupplier);
+        var testContext = new RavenwoodContext(
+                config.mTestPackageName, main, testResourcesLoader);
+        var targetContext = new RavenwoodContext(
+                config.mTargetPackageName, main, targetResourcesLoader);
+
+        // Set up app context.
+        var appContext = new RavenwoodContext(
+                config.mTargetPackageName, main, targetResourcesLoader);
+        appContext.setApplicationContext(appContext);
+        if (isSelfInstrumenting) {
+            testContext.setApplicationContext(appContext);
+            targetContext.setApplicationContext(appContext);
+        } else {
+            // When instrumenting into another APK, the test context doesn't have an app context.
+            targetContext.setApplicationContext(appContext);
+        }
+        config.mTestContext = testContext;
+        config.mTargetContext = targetContext;
+
+        // Prepare other fields.
         config.mInstrumentation = new Instrumentation();
-        config.mInstrumentation.basicInit(config.mContext);
+        config.mInstrumentation.basicInit(config.mTestContext, config.mTargetContext);
         InstrumentationRegistry.registerInstance(config.mInstrumentation, Bundle.EMPTY);
 
         RavenwoodSystemServer.init(config);
@@ -224,10 +242,14 @@
 
         InstrumentationRegistry.registerInstance(null, Bundle.EMPTY);
         config.mInstrumentation = null;
-        if (config.mContext != null) {
-            ((RavenwoodContext) config.mContext).cleanUp();
+        if (config.mTestContext != null) {
+            ((RavenwoodContext) config.mTestContext).cleanUp();
         }
-        config.mContext = null;
+        if (config.mTargetContext != null) {
+            ((RavenwoodContext) config.mTargetContext).cleanUp();
+        }
+        config.mTestContext = null;
+        config.mTargetContext = null;
 
         if (config.mProvideMainThread) {
             Looper.getMainLooper().quit();
@@ -240,7 +262,9 @@
         ServiceManager.reset$ravenwood();
 
         setSystemProperties(RavenwoodSystemProperties.DEFAULT_VALUES);
-        Binder.restoreCallingIdentity(sOriginalIdentityToken);
+        if (sOriginalIdentityToken != -1) {
+            Binder.restoreCallingIdentity(sOriginalIdentityToken);
+        }
         android.os.Process.reset$ravenwood();
 
         ResourcesManager.setInstance(null); // Better structure needed.
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java
index f3a93c1..d4090e2 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java
@@ -67,7 +67,7 @@
 
         sStartedServices = new ArraySet<>();
         sTimings = new TimingsTraceAndSlog();
-        sServiceManager = new SystemServiceManager(config.mContext);
+        sServiceManager = new SystemServiceManager(config.mTestContext);
         sServiceManager.setStartInfo(false,
                 SystemClock.elapsedRealtime(),
                 SystemClock.uptimeMillis());
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
index bc944d7..cb8af0c 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
@@ -171,6 +171,12 @@
     private Description mDescription = null;
     private Throwable mExceptionInConstructor = null;
 
+    /**
+     * Stores internal states / methods associated with this runner that's only needed in
+     * junit-impl.
+     */
+    final RavenwoodRunnerState mState = new RavenwoodRunnerState(this);
+
     private Error logAndFail(String message, Throwable exception) {
         Log.e(TAG, message, exception);
         throw new AssertionError(message, exception);
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java
index 04e0bed..ea33aa6 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java
@@ -63,7 +63,10 @@
     int mUid = NOBODY_UID;
     int mPid = sNextPid.getAndIncrement();
 
-    String mPackageName;
+    String mTestPackageName;
+    String mTargetPackageName;
+
+    int mMinSdkLevel;
 
     boolean mProvideMainThread = false;
 
@@ -71,9 +74,15 @@
 
     final List<Class<?>> mServicesRequired = new ArrayList<>();
 
-    volatile Context mContext;
+    volatile Context mTestContext;
+    volatile Context mTargetContext;
     volatile Instrumentation mInstrumentation;
 
+    /**
+     * Stores internal states / methods associated with this config that's only needed in
+     * junit-impl.
+     */
+    final RavenwoodConfigState mState = new RavenwoodConfigState(this);
     private RavenwoodConfig() {
     }
 
@@ -84,6 +93,12 @@
         return RavenwoodRule.isOnRavenwood();
     }
 
+    private void setDefaults() {
+        if (mTargetPackageName == null) {
+            mTargetPackageName = mTestPackageName;
+        }
+    }
+
     public static class Builder {
         private final RavenwoodConfig mConfig = new RavenwoodConfig();
 
@@ -109,11 +124,28 @@
         }
 
         /**
-         * Configure the identity of this process to be the given package name for the duration
-         * of the test. Has no effect on non-Ravenwood environments.
+         * Configure the package name of the test, which corresponds to
+         * {@link Instrumentation#getContext()}.
          */
         public Builder setPackageName(@NonNull String packageName) {
-            mConfig.mPackageName = Objects.requireNonNull(packageName);
+            mConfig.mTestPackageName = Objects.requireNonNull(packageName);
+            return this;
+        }
+
+        /**
+         * Configure the package name of the target app, which corresponds to
+         * {@link Instrumentation#getTargetContext()}. Defaults to {@link #setPackageName}.
+         */
+        public Builder setTargetPackageName(@NonNull String packageName) {
+            mConfig.mTargetPackageName = Objects.requireNonNull(packageName);
+            return this;
+        }
+
+        /**
+         * Configure the min SDK level of the test.
+         */
+        public Builder setMinSdkLevel(int sdkLevel) {
+            mConfig.mMinSdkLevel = sdkLevel;
             return this;
         }
 
@@ -178,6 +210,7 @@
         }
 
         public RavenwoodConfig build() {
+            mConfig.setDefaults();
             return mConfig;
         }
     }
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
index 7847e7c..984106b 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
@@ -216,7 +216,7 @@
      */
     @Deprecated
     public Context getContext() {
-        return Objects.requireNonNull(mConfiguration.mContext,
+        return Objects.requireNonNull(mConfiguration.mTestContext,
                 "Context is only available during @Test execution");
     }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodConfigState.java
similarity index 69%
rename from packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
rename to ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodConfigState.java
index 48ec198..43a28ba 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
+++ b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodConfigState.java
@@ -13,13 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+package android.platform.test.ravenwood;
 
-package com.android.systemui.scene.ui.composable.transitions
-
-import com.android.compose.animation.scene.TransitionBuilder
-
-fun TransitionBuilder.goneToNotificationsShadeTransition(
-    durationScale: Double = 1.0,
-) {
-    toNotificationsShadeTransition(durationScale)
+/** Stub class. The actual implementaetion is in junit-impl-src. */
+public class RavenwoodConfigState {
+    public RavenwoodConfigState(RavenwoodConfig config) {
+    }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
similarity index 69%
copy from packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
copy to ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
index 48ec198..83cbc52 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
+++ b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRunnerState.java
@@ -13,13 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+package android.platform.test.ravenwood;
 
-package com.android.systemui.scene.ui.composable.transitions
-
-import com.android.compose.animation.scene.TransitionBuilder
-
-fun TransitionBuilder.goneToNotificationsShadeTransition(
-    durationScale: Double = 1.0,
-) {
-    toNotificationsShadeTransition(durationScale)
+/** Stub class. The actual implementaetion is in junit-impl-src. */
+public class RavenwoodRunnerState {
+    public RavenwoodRunnerState(RavenwoodAwareTestRunner runner) {
+    }
 }
diff --git a/ravenwood/tests/bivalentinst/Android.bp b/ravenwood/tests/bivalentinst/Android.bp
new file mode 100644
index 0000000..38d1b299
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/Android.bp
@@ -0,0 +1,149 @@
+package {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_ravenwood_test {
+    name: "RavenwoodBivalentInstTest_self_inst",
+
+    srcs: [
+        "test/**/*.java",
+    ],
+    exclude_srcs: [
+        "test/**/*_nonself.java",
+    ],
+
+    static_libs: [
+        "RavenwoodBivalentInstTest_self_inst_device_R",
+
+        "androidx.annotation_annotation",
+        "androidx.test.ext.junit",
+        "androidx.test.rules",
+
+        "junit",
+        "truth",
+    ],
+    // TODO(b/366246777) uncomment it and the test.
+    // resource_apk: "RavenwoodBivalentInstTest_self_inst_device",
+    auto_gen_config: true,
+}
+
+android_ravenwood_test {
+    name: "RavenwoodBivalentInstTest_nonself_inst",
+
+    srcs: [
+        "test/**/*.java",
+    ],
+    exclude_srcs: [
+        "test/**/*_self.java",
+    ],
+
+    static_libs: [
+        "RavenwoodBivalentInstTestTarget_R",
+        "RavenwoodBivalentInstTest_nonself_inst_device_R",
+
+        "androidx.annotation_annotation",
+        "androidx.test.ext.junit",
+        "androidx.test.rules",
+
+        "junit",
+        "truth",
+    ],
+    // TODO(b/366246777) uncomment it and the test.
+    // resource_apk: "RavenwoodBivalentInstTestTarget",
+    auto_gen_config: true,
+}
+
+// We have 3 R.javas from the 3 packages (2 test apks below, and 1 target APK)
+// RavenwoodBivalentInstTest needs to use all of them, but we can't add all the
+// {.aapt.srcjar}'s together because that'd cause
+// "duplicate declaration of androidx.test.core.R$string."
+// So we build them as separate libraries, and include them as static_libs.
+java_library {
+    name: "RavenwoodBivalentInstTestTarget_R",
+    srcs: [
+        ":RavenwoodBivalentInstTestTarget{.aapt.srcjar}",
+    ],
+}
+
+java_library {
+    name: "RavenwoodBivalentInstTest_self_inst_device_R",
+    srcs: [
+        ":RavenwoodBivalentInstTest_self_inst_device{.aapt.srcjar}",
+    ],
+}
+
+java_library {
+    name: "RavenwoodBivalentInstTest_nonself_inst_device_R",
+    srcs: [
+        ":RavenwoodBivalentInstTest_nonself_inst_device{.aapt.srcjar}",
+    ],
+}
+
+android_test {
+    name: "RavenwoodBivalentInstTest_self_inst_device",
+
+    srcs: [
+        "test/**/*.java",
+    ],
+    exclude_srcs: [
+        "test/**/*_nonself.java",
+    ],
+    static_libs: [
+        "junit",
+        "truth",
+
+        "androidx.annotation_annotation",
+        "androidx.test.ext.junit",
+        "androidx.test.rules",
+
+        "ravenwood-junit",
+    ],
+    test_suites: [
+        "device-tests",
+    ],
+    use_resource_processor: false,
+    manifest: "AndroidManifest-self-inst.xml",
+    test_config: "AndroidTest-self-inst.xml",
+    optimize: {
+        enabled: false,
+    },
+}
+
+android_test {
+    name: "RavenwoodBivalentInstTest_nonself_inst_device",
+
+    srcs: [
+        "test/**/*.java",
+    ],
+    exclude_srcs: [
+        "test/**/*_self.java",
+    ],
+    static_libs: [
+        "junit",
+        "truth",
+
+        "androidx.annotation_annotation",
+        "androidx.test.ext.junit",
+        "androidx.test.rules",
+
+        "ravenwood-junit",
+    ],
+    data: [
+        ":RavenwoodBivalentInstTestTarget",
+    ],
+    test_suites: [
+        "device-tests",
+    ],
+    use_resource_processor: false,
+    manifest: "AndroidManifest-nonself-inst.xml",
+    test_config: "AndroidTest-nonself-inst.xml",
+    instrumentation_for: "RavenwoodBivalentInstTestTarget",
+    optimize: {
+        enabled: false,
+    },
+}
diff --git a/ravenwood/tests/bivalentinst/AndroidManifest-nonself-inst.xml b/ravenwood/tests/bivalentinst/AndroidManifest-nonself-inst.xml
new file mode 100644
index 0000000..a5a1f17
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/AndroidManifest-nonself-inst.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+        package="com.android.ravenwood.bivalentinsttest_nonself_inst">
+
+    <application android:debuggable="true" >
+        <uses-library android:name="android.test.runner" />
+    </application>
+
+    <instrumentation
+        android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.ravenwood.bivalentinst_target_app"
+        />
+</manifest>
diff --git a/ravenwood/tests/bivalentinst/AndroidManifest-self-inst.xml b/ravenwood/tests/bivalentinst/AndroidManifest-self-inst.xml
new file mode 100644
index 0000000..3dc4c56
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/AndroidManifest-self-inst.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+        package="com.android.ravenwood.bivalentinsttest_self_inst">
+
+    <application android:debuggable="true" >
+        <uses-library android:name="android.test.runner" />
+    </application>
+
+    <instrumentation
+        android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.ravenwood.bivalentinsttest_self_inst"
+        />
+</manifest>
diff --git a/ravenwood/tests/bivalentinst/AndroidTest-nonself-inst.xml b/ravenwood/tests/bivalentinst/AndroidTest-nonself-inst.xml
new file mode 100644
index 0000000..9491c53
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/AndroidTest-nonself-inst.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration>
+    <option name="test-suite-tag" value="apct" />
+    <option name="test-suite-tag" value="apct-instrumentation" />
+
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="RavenwoodBivalentInstTestTarget.apk" />
+        <option name="test-file-name" value="RavenwoodBivalentInstTest_nonself_inst_device.apk" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.ravenwood.bivalentinsttest_nonself_inst" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+    </test>
+</configuration>
diff --git a/ravenwood/tests/bivalentinst/AndroidTest-self-inst.xml b/ravenwood/tests/bivalentinst/AndroidTest-self-inst.xml
new file mode 100644
index 0000000..3079c06
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/AndroidTest-self-inst.xml
@@ -0,0 +1,29 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration>
+    <option name="test-suite-tag" value="apct" />
+    <option name="test-suite-tag" value="apct-instrumentation" />
+
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="test-file-name" value="RavenwoodBivalentInstTest_self_inst_device.apk" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.ravenwood.bivalentinsttest_self_inst" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+    </test>
+</configuration>
diff --git a/ravenwood/tests/bivalentinst/res/values/strings.xml b/ravenwood/tests/bivalentinst/res/values/strings.xml
new file mode 100644
index 0000000..73ef650
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/res/values/strings.xml
@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+           xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string translatable="false" name="test_string_in_test">String in test APK</string>
+</resources>
diff --git a/ravenwood/tests/bivalentinst/targetapp/Android.bp b/ravenwood/tests/bivalentinst/targetapp/Android.bp
new file mode 100644
index 0000000..7528a62
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/targetapp/Android.bp
@@ -0,0 +1,20 @@
+package {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_app {
+    name: "RavenwoodBivalentInstTestTarget",
+    srcs: [
+        "src/**/*.java",
+    ],
+    sdk_version: "current",
+    optimize: {
+        enabled: false,
+    },
+    use_resource_processor: false,
+}
diff --git a/ravenwood/tests/bivalentinst/targetapp/AndroidManifest.xml b/ravenwood/tests/bivalentinst/targetapp/AndroidManifest.xml
new file mode 100644
index 0000000..0715f5d
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/targetapp/AndroidManifest.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+        package="com.android.ravenwood.bivalentinst_target_app">
+    <application>
+    </application>
+</manifest>
diff --git a/ravenwood/tests/bivalentinst/targetapp/res/values/strings.xml b/ravenwood/tests/bivalentinst/targetapp/res/values/strings.xml
new file mode 100644
index 0000000..395bc2a
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/targetapp/res/values/strings.xml
@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<resources xmlns:android="http://schemas.android.com/apk/res/android"
+           xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string translatable="false" name="test_string_in_target">Test string in target APK</string>
+</resources>
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt b/ravenwood/tests/bivalentinst/targetapp/src/com/android/ravenwoodtest/bivalentinst/Empty.java
similarity index 69%
copy from packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
copy to ravenwood/tests/bivalentinst/targetapp/src/com/android/ravenwoodtest/bivalentinst/Empty.java
index 48ec198..15e50ec 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
+++ b/ravenwood/tests/bivalentinst/targetapp/src/com/android/ravenwoodtest/bivalentinst/Empty.java
@@ -13,13 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+package com.android.ravenwoodtest.bivalentinst;
 
-package com.android.systemui.scene.ui.composable.transitions
-
-import com.android.compose.animation.scene.TransitionBuilder
-
-fun TransitionBuilder.goneToNotificationsShadeTransition(
-    durationScale: Double = 1.0,
-) {
-    toNotificationsShadeTransition(durationScale)
+/**
+ * Empty class. We need it because an instrumentation target APK must have code.
+ */
+public class Empty {
 }
diff --git a/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_nonself.java b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_nonself.java
new file mode 100644
index 0000000..9f3ca6f
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_nonself.java
@@ -0,0 +1,115 @@
+/*
+ * 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.ravenwoodtest.bivalentinst;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.Instrumentation;
+import android.content.Context;
+import android.platform.test.annotations.DisabledOnRavenwood;
+import android.platform.test.ravenwood.RavenwoodConfig;
+import android.platform.test.ravenwood.RavenwoodConfig.Config;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for the case where the instrumentation target is _not_ the test APK itself.
+ */
+@RunWith(AndroidJUnit4.class)
+public class RavenwoodInstrumentationTest_nonself {
+    private static final String TARGET_PACKAGE_NAME =
+            "com.android.ravenwood.bivalentinst_target_app";
+    private static final String TEST_PACKAGE_NAME =
+            "com.android.ravenwood.bivalentinsttest_nonself_inst";
+
+    @Config
+    public static final RavenwoodConfig sConfig = new RavenwoodConfig.Builder()
+            .setPackageName(TEST_PACKAGE_NAME)
+            .setTargetPackageName(TARGET_PACKAGE_NAME)
+            .build();
+
+    private static Instrumentation sInstrumentation;
+    private static Context sTestContext;
+    private static Context sTargetContext;
+
+    @BeforeClass
+    public static void beforeClass() {
+        sInstrumentation = InstrumentationRegistry.getInstrumentation();
+        sTestContext = sInstrumentation.getContext();
+        sTargetContext = sInstrumentation.getTargetContext();
+    }
+
+    @Test
+    public void testTestContextPackageName() {
+        assertThat(sTestContext.getPackageName()).isEqualTo(TEST_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTargetContextPackageName() {
+        assertThat(sTargetContext.getPackageName()).isEqualTo(TARGET_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTestAppContext() {
+        // Test context doesn't have an app context.
+        assertThat(sTestContext.getApplicationContext()).isNull();
+    }
+
+    @Test
+    public void testTargetAppContextPackageName() {
+        assertThat(sTargetContext.getApplicationContext().getPackageName())
+                .isEqualTo(TARGET_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTargetAppAppContextPackageName() {
+        assertThat(sTargetContext.getApplicationContext()
+                .getApplicationContext().getPackageName())
+                .isEqualTo(TARGET_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testContextSameness() {
+        assertThat(sTargetContext).isNotSameInstanceAs(sTestContext);
+
+        assertThat(sTargetContext).isNotSameInstanceAs(sTargetContext.getApplicationContext());
+
+        assertThat(sTargetContext.getApplicationContext()).isSameInstanceAs(
+                sTargetContext.getApplicationContext().getApplicationContext());
+    }
+
+    @Test
+    @DisabledOnRavenwood(reason = "b/366246777")
+    public void testTargetAppResource() {
+        assertThat(sTargetContext.getString(
+                com.android.ravenwood.bivalentinst_target_app.R.string.test_string_in_target))
+                .isEqualTo("Test string in target APK");
+    }
+
+    @Test
+    @DisabledOnRavenwood(
+            reason = "Loading resources from non-self-instrumenting test APK isn't supported yet")
+    public void testTestAppResource() {
+        assertThat(sTestContext.getString(
+                com.android.ravenwood.bivalentinsttest_nonself_inst.R.string.test_string_in_test))
+                .isEqualTo("String in test APK");
+    }
+}
diff --git a/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_self.java b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_self.java
new file mode 100644
index 0000000..fdff222
--- /dev/null
+++ b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_self.java
@@ -0,0 +1,126 @@
+/*
+ * 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.ravenwoodtest.bivalentinst;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.Instrumentation;
+import android.content.Context;
+import android.platform.test.annotations.DisabledOnRavenwood;
+import android.platform.test.ravenwood.RavenwoodConfig;
+import android.platform.test.ravenwood.RavenwoodConfig.Config;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for the case where the instrumentation target is the test APK itself.
+ */
+@RunWith(AndroidJUnit4.class)
+public class RavenwoodInstrumentationTest_self {
+
+    private static final String TARGET_PACKAGE_NAME =
+            "com.android.ravenwood.bivalentinsttest_self_inst";
+    private static final String TEST_PACKAGE_NAME =
+            "com.android.ravenwood.bivalentinsttest_self_inst";
+
+    @Config
+    public static final RavenwoodConfig sConfig = new RavenwoodConfig.Builder()
+            .setPackageName(TEST_PACKAGE_NAME)
+            .setTargetPackageName(TARGET_PACKAGE_NAME)
+            .build();
+
+
+    private static Instrumentation sInstrumentation;
+    private static Context sTestContext;
+    private static Context sTargetContext;
+
+    @BeforeClass
+    public static void beforeClass() {
+        sInstrumentation = InstrumentationRegistry.getInstrumentation();
+        sTestContext = sInstrumentation.getContext();
+        sTargetContext = sInstrumentation.getTargetContext();
+    }
+
+    @Test
+    public void testTestContextPackageName() {
+        assertThat(sTestContext.getPackageName()).isEqualTo(TEST_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTargetContextPackageName() {
+        assertThat(sTargetContext.getPackageName()).isEqualTo(TARGET_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTestAppContextPackageName() {
+        assertThat(sTestContext.getApplicationContext().getPackageName())
+                .isEqualTo(TEST_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTestAppAppContextPackageName() {
+        assertThat(sTestContext.getApplicationContext().getPackageName())
+                .isEqualTo(TEST_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTargetAppContextPackageName() {
+        assertThat(sTargetContext.getApplicationContext()
+                .getApplicationContext().getPackageName())
+                .isEqualTo(TARGET_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testTargetAppAppContextPackageName() {
+        assertThat(sTargetContext.getApplicationContext()
+                .getApplicationContext().getPackageName())
+                .isEqualTo(TARGET_PACKAGE_NAME);
+    }
+
+    @Test
+    public void testContextSameness() {
+        assertThat(sTargetContext).isNotSameInstanceAs(sTestContext);
+
+        assertThat(sTestContext).isNotSameInstanceAs(sTestContext.getApplicationContext());
+        assertThat(sTargetContext).isNotSameInstanceAs(sTargetContext.getApplicationContext());
+
+        assertThat(sTestContext.getApplicationContext()).isSameInstanceAs(
+                sTestContext.getApplicationContext().getApplicationContext());
+        assertThat(sTargetContext.getApplicationContext()).isSameInstanceAs(
+                sTargetContext.getApplicationContext().getApplicationContext());
+    }
+
+    @Test
+    @DisabledOnRavenwood(reason = "b/366246777")
+    public void testTargetAppResource() {
+        assertThat(sTargetContext.getString(
+                com.android.ravenwood.bivalentinsttest_self_inst.R.string.test_string_in_test))
+                .isEqualTo("String in test APK");
+    }
+
+    @Test
+    @DisabledOnRavenwood(reason = "b/366246777")
+    public void testTestAppResource() {
+        assertThat(sTestContext.getString(
+                com.android.ravenwood.bivalentinsttest_self_inst.R.string.test_string_in_test))
+                .isEqualTo("String in test APK");
+    }
+}
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
index a38512e..f7f9a85 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
@@ -127,7 +127,7 @@
             }
         }
 
-        stats.totalProcessTime = log.iTime("$executableName processing $inJar") {
+        stats.totalProcessTime = log.vTime("$executableName processing $inJar") {
             ZipFile(inJar).use { inZip ->
                 val inEntries = inZip.entries()
 
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt
index e8341e5..10fe0a3 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt
@@ -20,6 +20,20 @@
 import com.android.hoststubgen.SetOnce
 import com.android.hoststubgen.ensureFileExists
 import com.android.hoststubgen.log
+import java.nio.file.Paths
+import kotlin.io.path.exists
+
+/**
+ * If this file exits, we also read options from it. This is "unsafe" because it could break
+ * incremental builds, if it sets any flag that affects the output file.
+ * (however, for now, there's no such options.)
+ *
+ * For example, to enable verbose logging, do `echo '-v' > ~/.raveniezr-unsafe`
+ *
+ * (but even the content of this file changes, soong won't rerun the command, so you need to
+ * remove the output first and then do a build again.)
+ */
+private val RAVENIZER_DOTFILE = System.getenv("HOME") + "/.raveniezr-unsafe"
 
 class RavenizerOptions(
     /** Input jar file*/
@@ -35,9 +49,16 @@
     var fatalValidation: SetOnce<Boolean> = SetOnce(false),
 ) {
     companion object {
-        fun parseArgs(args: Array<String>): RavenizerOptions {
+
+        fun parseArgs(origArgs: Array<String>): RavenizerOptions {
+            val args = origArgs.toMutableList()
+            if (Paths.get(RAVENIZER_DOTFILE).exists()) {
+                log.i("Reading options from $RAVENIZER_DOTFILE")
+                args.add(0, "@$RAVENIZER_DOTFILE")
+            }
+
             val ret = RavenizerOptions()
-            val ai = ArgIterator.withAtFiles(args)
+            val ai = ArgIterator.withAtFiles(args.toTypedArray())
 
             while (true) {
                 val arg = ai.nextArgOptional()
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt
index bd9d96d..cf6d6f6 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt
@@ -184,7 +184,7 @@
             av.visit("value", ravenwoodTestRunnerType.type)
             av.visitEnd()
         }
-        log.i("Update the @RunWith: ${classInternalName.toHumanReadableClassName()}")
+        log.v("Update the @RunWith: ${classInternalName.toHumanReadableClassName()}")
     }
 
     /*
@@ -442,7 +442,7 @@
             // Don't process a class if it has a @NoRavenizer annotation.
             classes.findClass(className)?.let { cn ->
                 if (cn.findAnyAnnotation(noRavenizerAnotType.descAsSet) != null) {
-                    log.w("Class ${className.toHumanReadableClassName()} has" +
+                    log.i("Class ${className.toHumanReadableClassName()} has" +
                         " @${noRavenizerAnotType.humanReadableName}. Skipping."
                     )
                     return false
diff --git a/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
index 56f373d..0044b4b 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
@@ -75,7 +75,7 @@
      * AppSearchSession} database.
      */
     AndroidFuture<AppSearchBatchResult<String, GenericDocument>> getByDocumentId(
-            GetByDocumentIdRequest getRequest);
+            @NonNull GetByDocumentIdRequest getRequest);
 
     /**
      * Retrieves documents from the open {@link AppSearchSession} that match a given query string
diff --git a/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java b/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java
index 5f60804..d140258 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java
@@ -61,7 +61,8 @@
  */
 public class MetadataSyncAdapter {
     private static final String TAG = MetadataSyncAdapter.class.getSimpleName();
-    private final FutureAppSearchSession mFutureAppSearchSession;
+    private final FutureAppSearchSession mRuntimeMetadataSearchSession;
+    private final FutureAppSearchSession mStaticMetadataSearchSession;
     private final Executor mSyncExecutor;
     private final PackageManager mPackageManager;
 
@@ -72,10 +73,12 @@
 
     public MetadataSyncAdapter(
             @NonNull Executor syncExecutor,
-            @NonNull FutureAppSearchSession futureAppSearchSession,
+            @NonNull FutureAppSearchSession runtimeMetadataSearchSession,
+            @NonNull FutureAppSearchSession staticMetadataSearchSession,
             @NonNull PackageManager packageManager) {
         mSyncExecutor = Objects.requireNonNull(syncExecutor);
-        mFutureAppSearchSession = Objects.requireNonNull(futureAppSearchSession);
+        mRuntimeMetadataSearchSession = Objects.requireNonNull(runtimeMetadataSearchSession);
+        mStaticMetadataSearchSession = Objects.requireNonNull(staticMetadataSearchSession);
         mPackageManager = Objects.requireNonNull(packageManager);
     }
 
@@ -104,11 +107,13 @@
             throws ExecutionException, InterruptedException {
         ArrayMap<String, ArraySet<String>> staticPackageToFunctionMap =
                 getPackageToFunctionIdMap(
+                        mStaticMetadataSearchSession,
                         AppFunctionStaticMetadataHelper.STATIC_SCHEMA_TYPE,
                         AppFunctionStaticMetadataHelper.PROPERTY_FUNCTION_ID,
                         AppFunctionStaticMetadataHelper.PROPERTY_PACKAGE_NAME);
         ArrayMap<String, ArraySet<String>> runtimePackageToFunctionMap =
                 getPackageToFunctionIdMap(
+                        mRuntimeMetadataSearchSession,
                         RUNTIME_SCHEMA_TYPE,
                         AppFunctionRuntimeMetadata.PROPERTY_FUNCTION_ID,
                         AppFunctionRuntimeMetadata.PROPERTY_PACKAGE_NAME);
@@ -129,7 +134,7 @@
             RemoveByDocumentIdRequest removeByDocumentIdRequest =
                     buildRemoveRuntimeMetadataRequest(removedFunctionsDiffMap);
             AppSearchBatchResult<String, Void> removeDocumentBatchResult =
-                    mFutureAppSearchSession.remove(removeByDocumentIdRequest).get();
+                    mRuntimeMetadataSearchSession.remove(removeByDocumentIdRequest).get();
             if (!removeDocumentBatchResult.isSuccess()) {
                 throw convertFailedAppSearchResultToException(
                         removeDocumentBatchResult.getFailures().values());
@@ -140,11 +145,12 @@
             // TODO(b/357551503): only set schema on package diff
             SetSchemaRequest addSetSchemaRequest =
                     buildSetSchemaRequestForRuntimeMetadataSchemas(appRuntimeMetadataSchemas);
-            Objects.requireNonNull(mFutureAppSearchSession.setSchema(addSetSchemaRequest).get());
+            Objects.requireNonNull(
+                    mRuntimeMetadataSearchSession.setSchema(addSetSchemaRequest).get());
             PutDocumentsRequest putDocumentsRequest =
                     buildPutRuntimeMetadataRequest(addedFunctionsDiffMap);
             AppSearchBatchResult<String, Void> putDocumentBatchResult =
-                    mFutureAppSearchSession.put(putDocumentsRequest).get();
+                    mRuntimeMetadataSearchSession.put(putDocumentsRequest).get();
             if (!putDocumentBatchResult.isSuccess()) {
                 throw convertFailedAppSearchResultToException(
                         putDocumentBatchResult.getFailures().values());
@@ -326,13 +332,17 @@
      * This method returns a map of package names to a set of function ids from the AppFunction
      * metadata.
      *
-     * @param schemaType The name space of the AppFunction metadata.
+     * @param searchSession The {@link FutureAppSearchSession} to search the AppFunction metadata.
+     * @param schemaType The schema type of the AppFunction metadata.
+     * @param propertyFunctionId The property name of the function id in the AppFunction metadata.
+     * @param propertyPackageName The property name of the package name in the AppFunction metadata.
      * @return A map of package names to a set of function ids from the AppFunction metadata.
      */
     @NonNull
     @VisibleForTesting
     @WorkerThread
-    ArrayMap<String, ArraySet<String>> getPackageToFunctionIdMap(
+    static ArrayMap<String, ArraySet<String>> getPackageToFunctionIdMap(
+            @NonNull FutureAppSearchSession searchSession,
             @NonNull String schemaType,
             @NonNull String propertyFunctionId,
             @NonNull String propertyPackageName)
@@ -343,7 +353,7 @@
         ArrayMap<String, ArraySet<String>> packageToFunctionIds = new ArrayMap<>();
 
         FutureSearchResults futureSearchResults =
-                mFutureAppSearchSession
+                searchSession
                         .search(
                                 "",
                                 buildMetadataSearchSpec(
diff --git a/services/core/Android.bp b/services/core/Android.bp
index 1b5b7e8..4e36e3f 100644
--- a/services/core/Android.bp
+++ b/services/core/Android.bp
@@ -237,6 +237,7 @@
         "dreams_flags_lib",
         "aconfig_new_storage_flags_lib",
         "powerstats_flags_lib",
+        "locksettings_flags_lib",
     ],
     javac_shard_size: 50,
     javacflags: [
diff --git a/services/core/java/com/android/server/adaptiveauth/AdaptiveAuthService.java b/services/core/java/com/android/server/adaptiveauth/AdaptiveAuthService.java
index 4da5cfc..9398c7a 100644
--- a/services/core/java/com/android/server/adaptiveauth/AdaptiveAuthService.java
+++ b/services/core/java/com/android/server/adaptiveauth/AdaptiveAuthService.java
@@ -34,7 +34,6 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
-import android.os.PowerManager;
 import android.os.SystemClock;
 import android.util.Log;
 import android.util.Slog;
@@ -73,7 +72,6 @@
     private final LockSettingsInternal mLockSettings;
     private final BiometricManager mBiometricManager;
     private final KeyguardManager mKeyguardManager;
-    private final PowerManager mPowerManager;
     private final WindowManagerInternal mWindowManager;
     private final UserManagerInternal mUserManager;
     @VisibleForTesting
@@ -93,7 +91,6 @@
         mBiometricManager = Objects.requireNonNull(
                 context.getSystemService(BiometricManager.class));
         mKeyguardManager = Objects.requireNonNull(context.getSystemService(KeyguardManager.class));
-        mPowerManager = Objects.requireNonNull(context.getSystemService(PowerManager.class));
         mWindowManager = Objects.requireNonNull(
                 LocalServices.getService(WindowManagerInternal.class));
         mUserManager = Objects.requireNonNull(LocalServices.getService(UserManagerInternal.class));
@@ -290,9 +287,6 @@
                     parentUserId);
         }
 
-        // Power off the display
-        mPowerManager.goToSleep(SystemClock.uptimeMillis());
-
         // Lock the device
         mWindowManager.lockNow();
 
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 2416ab6..75e9fad 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -1833,7 +1833,7 @@
 
     @Override
     @EnforcePermission(UPDATE_DEVICE_STATS)
-    public void noteScreenState(final int state) {
+    public void noteScreenState(final int displayId, final int state, final int reason) {
         super.noteScreenState_enforcePermission();
 
         synchronized (mLock) {
@@ -1843,7 +1843,8 @@
             mHandler.post(() -> {
                 if (DBG) Slog.d(TAG, "begin noteScreenState");
                 synchronized (mStats) {
-                    mStats.noteScreenStateLocked(0, state, elapsedRealtime, uptime, currentTime);
+                    mStats.noteScreenStateLocked(
+                            displayId, state, reason, elapsedRealtime, uptime, currentTime);
                 }
                 if (DBG) Slog.d(TAG, "end noteScreenState");
             });
@@ -1853,7 +1854,7 @@
 
     @Override
     @EnforcePermission(UPDATE_DEVICE_STATS)
-    public void noteScreenBrightness(final int brightness) {
+    public void noteScreenBrightness(final int displayId, final int brightness) {
         super.noteScreenBrightness_enforcePermission();
 
         synchronized (mLock) {
@@ -1861,7 +1862,8 @@
             final long uptime = SystemClock.uptimeMillis();
             mHandler.post(() -> {
                 synchronized (mStats) {
-                    mStats.noteScreenBrightnessLocked(0, brightness, elapsedRealtime, uptime);
+                    mStats.noteScreenBrightnessLocked(
+                            displayId, brightness, elapsedRealtime, uptime);
                 }
             });
         }
diff --git a/services/core/java/com/android/server/cpu/CpuInfoReader.java b/services/core/java/com/android/server/cpu/CpuInfoReader.java
index 984ad1d..a68451a 100644
--- a/services/core/java/com/android/server/cpu/CpuInfoReader.java
+++ b/services/core/java/com/android/server/cpu/CpuInfoReader.java
@@ -40,6 +40,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -80,13 +81,14 @@
     /** package **/ @interface CpusetCategory{}
 
     // TODO(b/242722241): Protect updatable variables with a local lock.
-    private final File mCpusetDir;
     private final long mMinReadIntervalMillis;
     private final SparseIntArray mCpusetCategoriesByCpus = new SparseIntArray();
     private final SparseArray<File> mCpuFreqPolicyDirsById = new SparseArray<>();
     private final SparseArray<StaticPolicyInfo> mStaticPolicyInfoById = new SparseArray<>();
     private final SparseArray<LongSparseLongArray> mTimeInStateByPolicyId = new SparseArray<>();
+    private final AtomicBoolean mShouldReadCpusetCategories;
 
+    private File mCpusetDir;
     private File mCpuFreqDir;
     private File mProcStatFile;
     private SparseArray<CpuUsageStats> mCumulativeCpuUsageStats = new SparseArray<>();
@@ -106,10 +108,13 @@
         mCpuFreqDir = cpuFreqDir;
         mProcStatFile = procStatFile;
         mMinReadIntervalMillis = minReadIntervalMillis;
+        mShouldReadCpusetCategories = new AtomicBoolean(true);
     }
 
     /**
      * Initializes CpuInfoReader and returns a boolean to indicate whether the reader is enabled.
+     *
+     * <p>Returns {@code true} on success. Otherwise, returns {@code false}.
      */
     public boolean init() {
         if (mCpuFreqPolicyDirsById.size() > 0) {
@@ -139,8 +144,7 @@
             Slogf.e(TAG, "Missing proc stat file at %s", mProcStatFile.getAbsolutePath());
             return false;
         }
-        readCpusetCategories();
-        if (mCpusetCategoriesByCpus.size() == 0) {
+        if (!readCpusetCategories()) {
             Slogf.e(TAG, "Failed to read cpuset information from %s", mCpusetDir.getAbsolutePath());
             return false;
         }
@@ -163,10 +167,19 @@
         return true;
     }
 
+  public void stopPeriodicCpusetReading() {
+        mShouldReadCpusetCategories.set(false);
+        if (!readCpusetCategories()) {
+            Slogf.e(TAG, "Failed to read cpuset information from %s",
+                    mCpusetDir.getAbsolutePath());
+            mIsEnabled = false;
+        }
+    }
+
     /**
      * Reads CPU information from proc and sys fs files exposed by the Kernel.
      *
-     * @return SparseArray keyed by CPU core ID; {@code null} on error or when disabled.
+     * <p>Returns SparseArray keyed by CPU core ID; {@code null} on error or when disabled.
      */
     @Nullable
     public SparseArray<CpuInfo> readCpuInfos() {
@@ -183,6 +196,12 @@
         }
         mLastReadUptimeMillis = uptimeMillis;
         mLastReadCpuInfos = null;
+        if (mShouldReadCpusetCategories.get() && !readCpusetCategories()) {
+            Slogf.e(TAG, "Failed to read cpuset information from %s",
+                    mCpusetDir.getAbsolutePath());
+            mIsEnabled = false;
+            return null;
+        }
         SparseArray<CpuUsageStats> cpuUsageStatsByCpus = readLatestCpuUsageStats();
         if (cpuUsageStatsByCpus == null || cpuUsageStatsByCpus.size() == 0) {
             Slogf.e(TAG, "Failed to read latest CPU usage stats");
@@ -324,7 +343,7 @@
     /**
      * Sets the CPU frequency for testing.
      *
-     * <p>Return {@code true} on success. Otherwise, returns {@code false}.
+     * <p>Returns {@code true} on success. Otherwise, returns {@code false}.
      */
     @VisibleForTesting
     boolean setCpuFreqDir(File cpuFreqDir) {
@@ -354,7 +373,7 @@
     /**
      * Sets the proc stat file for testing.
      *
-     * <p>Return true on success. Otherwise, returns false.
+     * <p>Returns {@code true} on success. Otherwise, returns {@code false}.
      */
     @VisibleForTesting
     boolean setProcStatFile(File procStatFile) {
@@ -366,6 +385,21 @@
         return true;
     }
 
+    /**
+     * Set the cpuset directory for testing.
+     *
+     * <p>Returns {@code true} on success. Otherwise, returns {@code false}.
+     */
+    @VisibleForTesting
+    boolean setCpusetDir(File cpusetDir) {
+        if (!cpusetDir.exists() && !cpusetDir.isDirectory()) {
+            Slogf.e(TAG, "Missing or invalid cpuset directory at %s", cpusetDir.getAbsolutePath());
+            return false;
+        }
+        mCpusetDir = cpusetDir;
+        return true;
+    }
+
     private void populateCpuFreqPolicyDirsById(File[] policyDirs) {
         mCpuFreqPolicyDirsById.clear();
         for (int i = 0; i < policyDirs.length; i++) {
@@ -381,12 +415,27 @@
         }
     }
 
-    private void readCpusetCategories() {
+    /**
+     * Reads cpuset categories by CPU.
+     *
+     * <p>The cpusets are read from the cpuset category specific directories
+     * under the /dev/cpuset directory. The cpuset categories are subject to change at any point
+     * during system bootup, as determined by the init rules specified within the init.rc files.
+     * Therefore, it's necessary to read the cpuset categories each time before accessing CPU usage
+     * statistics until the system boot completes. Once the boot is complete, the latest changes to
+     * the cpuset categories will take a few seconds to propagate. Thus, on boot complete,
+     * the periodic reading is stopped with a delay of
+     * {@link CpuMonitorService#STOP_PERIODIC_CPUSET_READING_DELAY_MILLISECONDS}.
+     *
+     * <p>Returns {@code true} on success. Otherwise, returns {@code false}.
+     */
+    private boolean readCpusetCategories() {
         File[] cpusetDirs = mCpusetDir.listFiles(File::isDirectory);
         if (cpusetDirs == null) {
             Slogf.e(TAG, "Missing cpuset directories at %s", mCpusetDir.getAbsolutePath());
-            return;
+            return false;
         }
+        mCpusetCategoriesByCpus.clear();
         for (int i = 0; i < cpusetDirs.length; i++) {
             File dir = cpusetDirs[i];
             @CpusetCategory int cpusetCategory;
@@ -418,6 +467,7 @@
                 }
             }
         }
+        return mCpusetCategoriesByCpus.size() > 0;
     }
 
     private void readStaticPolicyInfo() {
diff --git a/services/core/java/com/android/server/cpu/CpuMonitorService.java b/services/core/java/com/android/server/cpu/CpuMonitorService.java
index 7ea2c1b..88ff7e4 100644
--- a/services/core/java/com/android/server/cpu/CpuMonitorService.java
+++ b/services/core/java/com/android/server/cpu/CpuMonitorService.java
@@ -22,6 +22,7 @@
 import static com.android.server.cpu.CpuAvailabilityMonitoringConfig.CPUSET_BACKGROUND;
 import static com.android.server.cpu.CpuInfoReader.FLAG_CPUSET_CATEGORY_BACKGROUND;
 import static com.android.server.cpu.CpuInfoReader.FLAG_CPUSET_CATEGORY_TOP_APP;
+import static com.android.server.SystemService.PHASE_BOOT_COMPLETED;
 
 import android.annotation.Nullable;
 import android.content.Context;
@@ -82,6 +83,15 @@
     //  frequently. Should this duration be increased as well when this happens?
     private static final long LATEST_AVAILABILITY_DURATION_MILLISECONDS =
             TimeUnit.SECONDS.toMillis(30);
+    /**
+     * Delay to stop the periodic cpuset reading after boot complete.
+     *
+     * Device specific implementations can update cpuset on boot complete. This may take
+     * a few seconds to propagate. So, wait for a few minutes before stopping the periodic cpuset
+     * reading.
+     */
+    private static final long STOP_PERIODIC_CPUSET_READING_DELAY_MILLISECONDS =
+            TimeUnit.MINUTES.toMillis(2);
 
     private final Context mContext;
     private final HandlerThread mHandlerThread;
@@ -90,6 +100,7 @@
     private final long mNormalMonitoringIntervalMillis;
     private final long mDebugMonitoringIntervalMillis;
     private final long mLatestAvailabilityDurationMillis;
+    private final long mStopPeriodicCpusetReadingDelayMillis;
     private final Object mLock = new Object();
     @GuardedBy("mLock")
     private final SparseArrayMap<CpuMonitorInternal.CpuAvailabilityCallback,
@@ -153,13 +164,15 @@
         this(context, new CpuInfoReader(), new ServiceThread(TAG,
                         Process.THREAD_PRIORITY_BACKGROUND, /* allowIo= */ true),
                 Build.IS_USERDEBUG || Build.IS_ENG, NORMAL_MONITORING_INTERVAL_MILLISECONDS,
-                DEBUG_MONITORING_INTERVAL_MILLISECONDS, LATEST_AVAILABILITY_DURATION_MILLISECONDS);
+                DEBUG_MONITORING_INTERVAL_MILLISECONDS, LATEST_AVAILABILITY_DURATION_MILLISECONDS,
+                STOP_PERIODIC_CPUSET_READING_DELAY_MILLISECONDS);
     }
 
     @VisibleForTesting
     CpuMonitorService(Context context, CpuInfoReader cpuInfoReader, HandlerThread handlerThread,
             boolean shouldDebugMonitor, long normalMonitoringIntervalMillis,
-            long debugMonitoringIntervalMillis, long latestAvailabilityDurationMillis) {
+            long debugMonitoringIntervalMillis, long latestAvailabilityDurationMillis,
+            long stopPeriodicCpusetReadingDelayMillis) {
         super(context);
         mContext = context;
         mHandlerThread = handlerThread;
@@ -167,6 +180,7 @@
         mNormalMonitoringIntervalMillis = normalMonitoringIntervalMillis;
         mDebugMonitoringIntervalMillis = debugMonitoringIntervalMillis;
         mLatestAvailabilityDurationMillis = latestAvailabilityDurationMillis;
+        mStopPeriodicCpusetReadingDelayMillis = stopPeriodicCpusetReadingDelayMillis;
         mCpuInfoReader = cpuInfoReader;
         mCpusetInfosByCpuset = new SparseArray<>(2);
         mCpusetInfosByCpuset.append(CPUSET_ALL, new CpusetInfo(CPUSET_ALL));
@@ -200,6 +214,16 @@
         }
     }
 
+    @Override
+    public void onBootPhase(int phase) {
+        if (phase != PHASE_BOOT_COMPLETED) {
+            return;
+        }
+        Slogf.i(TAG, "Stopping periodic cpuset reading on boot complete");
+        mHandler.postDelayed(() -> mCpuInfoReader.stopPeriodicCpusetReading(),
+                mStopPeriodicCpusetReadingDelayMillis);
+    }
+
     @VisibleForTesting
     long getCurrentMonitoringIntervalMillis() {
         synchronized (mLock) {
diff --git a/services/core/java/com/android/server/display/DisplayPowerController.java b/services/core/java/com/android/server/display/DisplayPowerController.java
index c3faec0..04573f4 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController.java
@@ -536,7 +536,9 @@
         mLastBrightnessEvent = new BrightnessEvent(mDisplayId);
         mTempBrightnessEvent = new BrightnessEvent(mDisplayId);
 
-        if (mDisplayId == Display.DEFAULT_DISPLAY) {
+        if (flags.isBatteryStatsEnabledForAllDisplays()) {
+            mBatteryStats = BatteryStatsService.getService();
+        } else if (mDisplayId == Display.DEFAULT_DISPLAY) {
             mBatteryStats = BatteryStatsService.getService();
         } else {
             mBatteryStats = null;
@@ -2791,8 +2793,7 @@
                 screenState, mDisplayStatsId, reason);
         if (mBatteryStats != null) {
             try {
-                // TODO(multi-display): make this multi-display
-                mBatteryStats.noteScreenState(screenState);
+                mBatteryStats.noteScreenState(mDisplayId, screenState, reason);
             } catch (RemoteException e) {
                 // same process
             }
@@ -2807,7 +2808,7 @@
                 int brightnessInt = mFlags.isBrightnessIntRangeUserPerceptionEnabled()
                         ? BrightnessSynchronizer.brightnessFloatToIntSetting(mContext, brightness)
                         : BrightnessSynchronizer.brightnessFloatToInt(brightness);
-                mBatteryStats.noteScreenBrightness(brightnessInt);
+                mBatteryStats.noteScreenBrightness(mDisplayId, brightnessInt);
             } catch (RemoteException e) {
                 // same process
             }
diff --git a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
index 69b67c8..f600e7f 100644
--- a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
+++ b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
@@ -200,6 +200,11 @@
             Flags::normalBrightnessForDozeParameter
     );
 
+    private final FlagState mEnableBatteryStatsForAllDisplays = new FlagState(
+            Flags.FLAG_ENABLE_BATTERY_STATS_FOR_ALL_DISPLAYS,
+            Flags::enableBatteryStatsForAllDisplays
+    );
+
     /**
      * @return {@code true} if 'port' is allowed in display layout configuration file.
      */
@@ -415,6 +420,14 @@
     }
 
     /**
+      * @return {@code true} if battery stats is enabled for all displays, not just the primary
+      * display.
+      */
+    public boolean isBatteryStatsEnabledForAllDisplays() {
+        return mEnableBatteryStatsForAllDisplays.isEnabled();
+    }
+
+    /**
      * dumps all flagstates
      * @param pw printWriter
      */
@@ -456,6 +469,7 @@
         pw.println(" " + mNewHdrBrightnessModifier);
         pw.println(" " + mNormalBrightnessForDozeParameter);
         pw.println(" " + mIdleScreenConfigInSubscribingLightSensor);
+        pw.println(" " + mEnableBatteryStatsForAllDisplays);
     }
 
     private static class FlagState {
diff --git a/services/core/java/com/android/server/display/feature/display_flags.aconfig b/services/core/java/com/android/server/display/feature/display_flags.aconfig
index 70230b4..9968ba5 100644
--- a/services/core/java/com/android/server/display/feature/display_flags.aconfig
+++ b/services/core/java/com/android/server/display/feature/display_flags.aconfig
@@ -349,3 +349,11 @@
       purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "enable_battery_stats_for_all_displays"
+    namespace: "display_manager"
+    description: "Flag to enable battery stats for all displays."
+    bug: "366112793"
+    is_fixed_read_only: true
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/input/debug/TouchpadVisualizationView.java b/services/core/java/com/android/server/input/debug/TouchpadVisualizationView.java
index 2ed6f44..67c3621 100644
--- a/services/core/java/com/android/server/input/debug/TouchpadVisualizationView.java
+++ b/services/core/java/com/android/server/input/debug/TouchpadVisualizationView.java
@@ -39,35 +39,43 @@
     TouchpadHardwareState mLatestHardwareState = new TouchpadHardwareState(0, 0, 0, 0,
             new TouchpadFingerState[]{});
 
-    private final Paint mOvalPaint;
+    private final Paint mOvalStrokePaint;
+    private final Paint mOvalFillPaint;
+    private final RectF mTempOvalRect = new RectF();
 
     public TouchpadVisualizationView(Context context,
             TouchpadHardwareProperties touchpadHardwareProperties) {
         super(context);
         mTouchpadHardwareProperties = touchpadHardwareProperties;
         mScaleFactor = 1;
-        mOvalPaint = new Paint();
-        mOvalPaint.setAntiAlias(true);
-        mOvalPaint.setARGB(255, 0, 0, 0);
-        mOvalPaint.setStyle(Paint.Style.STROKE);
+        mOvalStrokePaint = new Paint();
+        mOvalStrokePaint.setAntiAlias(true);
+        mOvalStrokePaint.setARGB(255, 0, 0, 0);
+        mOvalStrokePaint.setStyle(Paint.Style.STROKE);
+        mOvalFillPaint = new Paint();
+        mOvalFillPaint.setAntiAlias(true);
+        mOvalFillPaint.setARGB(255, 0, 0, 0);
     }
 
-    private final RectF mOvalRect = new RectF();
-
-    private void drawOval(Canvas canvas, float x, float y, float major, float minor, float angle,
-            Paint paint) {
+    private void drawOval(Canvas canvas, float x, float y, float major, float minor, float angle) {
         canvas.save(Canvas.MATRIX_SAVE_FLAG);
         canvas.rotate(angle, x, y);
-        mOvalRect.left = x - minor / 2;
-        mOvalRect.right = x + minor / 2;
-        mOvalRect.top = y - major / 2;
-        mOvalRect.bottom = y + major / 2;
-        canvas.drawOval(mOvalRect, paint);
+        mTempOvalRect.left = x - minor / 2;
+        mTempOvalRect.right = x + minor / 2;
+        mTempOvalRect.top = y - major / 2;
+        mTempOvalRect.bottom = y + major / 2;
+        canvas.drawOval(mTempOvalRect, mOvalStrokePaint);
+        canvas.drawOval(mTempOvalRect, mOvalFillPaint);
         canvas.restore();
     }
 
     @Override
     protected void onDraw(Canvas canvas) {
+        float maximumPressure = 0;
+        for (TouchpadFingerState touchpadFingerState : mLatestHardwareState.getFingerStates()) {
+            maximumPressure = Math.max(maximumPressure, touchpadFingerState.getPressure());
+        }
+
         for (TouchpadFingerState touchpadFingerState : mLatestHardwareState.getFingerStates()) {
             float newX = translateRange(mTouchpadHardwareProperties.getLeft(),
                     mTouchpadHardwareProperties.getRight(), 0, getWidth(),
@@ -88,7 +96,11 @@
             float newTouchMajor = touchpadFingerState.getTouchMajor() * mScaleFactor / resY;
             float newTouchMinor = touchpadFingerState.getTouchMinor() * mScaleFactor / resX;
 
-            drawOval(canvas, newX, newY, newTouchMajor, newTouchMinor, newAngle, mOvalPaint);
+            float pressureToOpacity = translateRange(0, maximumPressure, 0, 255,
+                    touchpadFingerState.getPressure());
+            mOvalFillPaint.setAlpha((int) pressureToOpacity);
+
+            drawOval(canvas, newX, newY, newTouchMajor, newTouchMinor, newAngle);
         }
     }
 
diff --git a/services/core/java/com/android/server/locksettings/Android.bp b/services/core/java/com/android/server/locksettings/Android.bp
new file mode 100644
index 0000000..53f1ac6
--- /dev/null
+++ b/services/core/java/com/android/server/locksettings/Android.bp
@@ -0,0 +1,11 @@
+aconfig_declarations {
+    name: "locksettings_flags",
+    package: "com.android.server.locksettings",
+    container: "system",
+    srcs: ["*.aconfig"],
+}
+
+java_aconfig_library {
+    name: "locksettings_flags_lib",
+    aconfig_declarations: "locksettings_flags",
+}
diff --git a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
index f44b852..820c0ef 100644
--- a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
+++ b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
@@ -273,11 +273,6 @@
                     "server_based_ror_enabled", false);
         }
 
-        public boolean waitForInternet() {
-            return DeviceConfig.getBoolean(
-                    DeviceConfig.NAMESPACE_OTA, "wait_for_internet_ror", false);
-        }
-
         public boolean isNetworkConnected() {
             final ConnectivityManager connectivityManager =
                     mContext.getSystemService(ConnectivityManager.class);
@@ -433,7 +428,7 @@
 
     /** Wrapper function to set error code serialized through handler, */
     private void setLoadEscrowDataErrorCode(@RebootEscrowErrorCode int value, Handler handler) {
-        if (mInjector.waitForInternet()) {
+        if (Flags.waitForInternetRor()) {
             mInjector.post(
                     handler,
                     () -> {
@@ -516,7 +511,7 @@
             mWakeLock.acquire(mInjector.getWakeLockTimeoutMillis());
         }
 
-        if (mInjector.waitForInternet()) {
+        if (Flags.waitForInternetRor()) {
             // Timeout to stop retrying same as the wake lock timeout.
             mInjector.postDelayed(
                     retryHandler,
@@ -553,7 +548,7 @@
             return;
         }
 
-        if (mInjector.waitForInternet()) {
+        if (Flags.waitForInternetRor()) {
             if (mRebootEscrowTimedOut) {
                 Slog.w(TAG, "Failed to load reboot escrow data within timeout");
                 compareAndSetLoadEscrowDataErrorCode(
diff --git a/services/core/java/com/android/server/locksettings/flags.aconfig b/services/core/java/com/android/server/locksettings/flags.aconfig
new file mode 100644
index 0000000..6818de9
--- /dev/null
+++ b/services/core/java/com/android/server/locksettings/flags.aconfig
@@ -0,0 +1,9 @@
+package: "com.android.server.locksettings"
+container: "system"
+
+flag {
+     name: "wait_for_internet_ror"
+     namespace: "sudo"
+     description: "Feature flag to wait for internet connectivity before calling resume on reboot server."
+     bug: "231660348"
+}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/notification/GroupHelper.java b/services/core/java/com/android/server/notification/GroupHelper.java
index 4fa7112..82e00d9 100644
--- a/services/core/java/com/android/server/notification/GroupHelper.java
+++ b/services/core/java/com/android/server/notification/GroupHelper.java
@@ -757,8 +757,12 @@
 
             // scenario 3: sparse/singleton groups
             if (Flags.notificationForceGroupSingletons()) {
-                groupSparseGroups(record, notificationList, summaryByGroupKey, sectioner,
-                    fullAggregateGroupKey);
+                try {
+                    groupSparseGroups(record, notificationList, summaryByGroupKey, sectioner,
+                        fullAggregateGroupKey);
+                } catch (Throwable e) {
+                    Slog.wtf(TAG, "Failed to group sparse groups", e);
+                }
             }
         }
     }
diff --git a/services/core/java/com/android/server/notification/ZenModeEventLogger.java b/services/core/java/com/android/server/notification/ZenModeEventLogger.java
index b03a54e..fcc5e97 100644
--- a/services/core/java/com/android/server/notification/ZenModeEventLogger.java
+++ b/services/core/java/com/android/server/notification/ZenModeEventLogger.java
@@ -419,7 +419,7 @@
 
             if (config.automaticRules != null) {
                 for (ZenModeConfig.ZenRule rule : config.automaticRules.values()) {
-                    if (rule != null && rule.isAutomaticActive()) {
+                    if (rule != null && rule.isActive()) {
                         rules.add(rule);
                     }
                 }
diff --git a/services/core/java/com/android/server/notification/ZenModeHelper.java b/services/core/java/com/android/server/notification/ZenModeHelper.java
index 2ada9ae4..e9db1b5 100644
--- a/services/core/java/com/android/server/notification/ZenModeHelper.java
+++ b/services/core/java/com/android/server/notification/ZenModeHelper.java
@@ -887,7 +887,7 @@
                 return Condition.STATE_UNKNOWN;
             }
             if (Flags.modesApi() && Flags.modesUi()) {
-                return rule.isAutomaticActive() ? STATE_TRUE : STATE_FALSE;
+                return rule.isActive() ? STATE_TRUE : STATE_FALSE;
             } else {
                 // Buggy, does not consider snoozing!
                 return rule.condition != null ? rule.condition.state : STATE_FALSE;
@@ -967,12 +967,12 @@
                 // snoozing-unsnoozing or activating-stopping.
                 if (condition.state == STATE_TRUE) {
                     rule.resetConditionOverride();
-                    if (!rule.isAutomaticActive()) {
+                    if (!rule.isActive()) {
                         rule.setConditionOverride(OVERRIDE_ACTIVATE);
                     }
                 } else if (condition.state == STATE_FALSE) {
                     rule.resetConditionOverride();
-                    if (rule.isAutomaticActive()) {
+                    if (rule.isActive()) {
                         rule.setConditionOverride(OVERRIDE_DEACTIVATE);
                     }
                 }
@@ -1609,7 +1609,7 @@
                     // User deactivation of DND means just turning off the manual DND rule.
                     // For API calls (different origin) keep old behavior of snoozing all rules.
                     for (ZenRule automaticRule : newConfig.automaticRules.values()) {
-                        if (automaticRule.isAutomaticActive()) {
+                        if (automaticRule.isActive()) {
                             automaticRule.setConditionOverride(OVERRIDE_DEACTIVATE);
                         }
                     }
@@ -1618,7 +1618,7 @@
                 if (zenMode == Global.ZEN_MODE_OFF) {
                     newConfig.manualRule = null;
                     for (ZenRule automaticRule : newConfig.automaticRules.values()) {
-                        if (automaticRule.isAutomaticActive()) {
+                        if (automaticRule.isActive()) {
                             automaticRule.setConditionOverride(OVERRIDE_DEACTIVATE);
                         }
                     }
@@ -1665,7 +1665,7 @@
                 mConfig.manualRule.dumpDebug(proto, ZenModeProto.ENABLED_ACTIVE_CONDITIONS);
             }
             for (ZenRule rule : mConfig.automaticRules.values()) {
-                if (rule.isAutomaticActive()) {
+                if (rule.isActive()) {
                     rule.dumpDebug(proto, ZenModeProto.ENABLED_ACTIVE_CONDITIONS);
                 }
             }
@@ -2020,9 +2020,9 @@
                             scheduleEnabledBroadcast(
                                     rule.getPkg(), config.user, rule.id, rule.enabled);
                         }
-                        if (original.isAutomaticActive() != rule.isAutomaticActive()) {
+                        if (original.isActive() != rule.isActive()) {
                             scheduleActivationBroadcast(
-                                    rule.getPkg(), config.user, rule.id, rule.isAutomaticActive());
+                                    rule.getPkg(), config.user, rule.id, rule.isActive());
                         }
                     }
                 }
@@ -2106,7 +2106,7 @@
             if (mConfig.isManualActive()) return mConfig.manualRule.zenMode;
             int zen = Global.ZEN_MODE_OFF;
             for (ZenRule automaticRule : mConfig.automaticRules.values()) {
-                if (automaticRule.isAutomaticActive()) {
+                if (automaticRule.isActive()) {
                     if (zenSeverity(automaticRule.zenMode) > zenSeverity(zen)) {
                         // automatic rule triggered dnd and user hasn't seen update dnd dialog
                         if (Settings.Secure.getInt(mContext.getContentResolver(),
@@ -2182,7 +2182,7 @@
             }
 
             for (ZenRule automaticRule : mConfig.automaticRules.values()) {
-                if (automaticRule.isAutomaticActive()) {
+                if (automaticRule.isActive()) {
                     // Active rules with INTERRUPTION_FILTER_ALL are not included in consolidated
                     // policy. This is relevant in case some other active rule has a more
                     // restrictive INTERRUPTION_FILTER but a more lenient ZenPolicy!
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index 680b1ac..cb8e1a0 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -5031,9 +5031,7 @@
         if (mPretendScreenOff != pretendScreenOff) {
             mPretendScreenOff = pretendScreenOff;
             final int primaryScreenState = mPerDisplayBatteryStats[0].screenState;
-            noteScreenStateLocked(0, primaryScreenState,
-                    mClock.elapsedRealtime(), mClock.uptimeMillis(),
-                    mClock.currentTimeMillis());
+            noteScreenStateLocked(0, primaryScreenState);
         }
     }
 
@@ -5554,15 +5552,29 @@
         }
     }
 
+    private static String getScreenStateTag(
+            int display, int state, @Display.StateReason int reason) {
+        return String.format(
+                "display=%d state=%s reason=%s",
+                display, Display.stateToString(state), Display.stateReasonToString(reason));
+    }
+
     @GuardedBy("this")
     public void noteScreenStateLocked(int display, int state) {
-        noteScreenStateLocked(display, state, mClock.elapsedRealtime(), mClock.uptimeMillis(),
-                mClock.currentTimeMillis());
+        noteScreenStateLocked(display, state, Display.STATE_REASON_UNKNOWN,
+                mClock.elapsedRealtime(), mClock.uptimeMillis(), mClock.currentTimeMillis());
     }
 
     @GuardedBy("this")
     public void noteScreenStateLocked(int display, int displayState,
-            long elapsedRealtimeMs, long uptimeMs, long currentTimeMs) {
+            @Display.StateReason int displayStateReason, long elapsedRealtimeMs, long uptimeMs,
+            long currentTimeMs) {
+        if (Flags.batteryStatsScreenStateEvent()) {
+            mHistory.recordEvent(
+                    elapsedRealtimeMs, uptimeMs, HistoryItem.EVENT_DISPLAY_STATE_CHANGED,
+                    getScreenStateTag(display, displayState, displayStateReason),
+                    Process.INVALID_UID);
+        }
         // Battery stats relies on there being 4 states. To accommodate this, new states beyond the
         // original 4 are mapped to one of the originals.
         if (displayState > MAX_TRACKED_SCREEN_STATE) {
diff --git a/services/core/java/com/android/server/power/stats/flags.aconfig b/services/core/java/com/android/server/power/stats/flags.aconfig
index cc0a283..05d29f5 100644
--- a/services/core/java/com/android/server/power/stats/flags.aconfig
+++ b/services/core/java/com/android/server/power/stats/flags.aconfig
@@ -68,3 +68,11 @@
     description: "Disable deprecated BatteryUsageStatsAtom pulled atom"
     bug: "324602949"
 }
+
+flag {
+    name: "battery_stats_screen_state_event"
+    namespace: "backstage_power"
+    description: "Guards the battery stats event for screen state changes."
+    bug: "364350206"
+    is_fixed_read_only: true
+}
diff --git a/services/core/java/com/android/server/uri/UriGrantsManagerInternal.java b/services/core/java/com/android/server/uri/UriGrantsManagerInternal.java
index 195e91c..49825f1 100644
--- a/services/core/java/com/android/server/uri/UriGrantsManagerInternal.java
+++ b/services/core/java/com/android/server/uri/UriGrantsManagerInternal.java
@@ -64,13 +64,36 @@
             String targetPkg, int targetUserId);
 
     /**
-     * Same as {@link #checkGrantUriPermissionFromIntent(Intent, int, String, int)}, but with an
-     * extra parameter {@code requireContentUriPermissionFromCaller}, which is the value from {@link
-     * android.R.attr#requireContentUriPermissionFromCaller} attribute.
+     * Same as {@link #checkGrantUriPermissionFromIntent(Intent, int, String, int)}, but with:
+     * - {@code requireContentUriPermissionFromCaller}, which is the value from {@link
+     *   android.R.attr#requireContentUriPermissionFromCaller} attribute.
+     * - {@code requestHashCode}, which is required to differentiate activity launches for logging
+     *   ContentOrFileUriEventReported message.
      */
     NeededUriGrants checkGrantUriPermissionFromIntent(Intent intent, int callingUid,
             String targetPkg, int targetUserId,
-            @RequiredContentUriPermission int requireContentUriPermissionFromCaller);
+            @RequiredContentUriPermission int requireContentUriPermissionFromCaller,
+            int requestHashCode);
+
+    /**
+     * Notify that an activity launch request has been completed and perform the following actions:
+     * - If the activity launch was unsuccessful, then clean up all the collected the content URIs
+     *   that were passed during that launch.
+     * - If the activity launch was successful, then log cog content URIs that were passed during
+     *   that launch. Specifically:
+     *   - The caller didn't have read permission to them.
+     *   - The activity's {@link android.R.attr#requireContentUriPermissionFromCaller} was set to
+     *     "none".
+     *
+     * <p>Note that:
+     * - The API has to be called after
+     *   {@link #checkGrantUriPermissionFromIntent(Intent, int, String, int, int, int)} was called.
+     * - The API is not idempotent, i.e. content URIs may be logged only once because the API clears
+     *   the content URIs after logging.
+     */
+    void notifyActivityLaunchRequestCompleted(int requestHashCode, boolean isSuccessfulLaunch,
+            String intentAction, int callingUid, String callingActivityName, int calleeUid,
+            String calleeActivityName, boolean isStartActivityForResult);
 
     /**
      * Extend a previously calculated set of permissions grants to the given
diff --git a/services/core/java/com/android/server/uri/UriGrantsManagerService.java b/services/core/java/com/android/server/uri/UriGrantsManagerService.java
index a581b08..3479b6c 100644
--- a/services/core/java/com/android/server/uri/UriGrantsManagerService.java
+++ b/services/core/java/com/android/server/uri/UriGrantsManagerService.java
@@ -24,6 +24,7 @@
 import static android.content.Intent.FLAG_GRANT_PERSISTABLE_URI_PERMISSION;
 import static android.content.Intent.FLAG_GRANT_PREFIX_URI_PERMISSION;
 import static android.content.pm.ActivityInfo.CONTENT_URI_PERMISSION_NONE;
+import static android.content.pm.ActivityInfo.CONTENT_URI_PERMISSION_READ;
 import static android.content.pm.ActivityInfo.CONTENT_URI_PERMISSION_READ_OR_WRITE;
 import static android.content.pm.ActivityInfo.isRequiredContentUriPermissionRead;
 import static android.content.pm.ActivityInfo.isRequiredContentUriPermissionWrite;
@@ -39,6 +40,8 @@
 import static android.os.Process.myUid;
 
 import static com.android.internal.util.XmlUtils.writeBooleanAttribute;
+import static com.android.internal.util.FrameworkStatsLog.CONTENT_OR_FILE_URI_EVENT_REPORTED;
+import static com.android.internal.util.FrameworkStatsLog.CONTENT_OR_FILE_URI_EVENT_REPORTED__EVENT_TYPE__CONTENT_URI_WITHOUT_CALLER_READ_PERMISSION;
 import static com.android.server.uri.UriGrantsManagerService.H.PERSIST_URI_GRANTS_MSG;
 
 import static org.xmlpull.v1.XmlPullParser.END_DOCUMENT;
@@ -78,6 +81,7 @@
 import android.provider.Downloads;
 import android.text.format.DateUtils;
 import android.util.ArrayMap;
+import android.util.ArraySet;
 import android.util.AtomicFile;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -86,6 +90,7 @@
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ArrayUtils;
+import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.Preconditions;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
@@ -153,6 +158,22 @@
     private final SparseArray<ArrayMap<GrantUri, UriPermission>>
             mGrantedUriPermissions = new SparseArray<>();
 
+    /**
+     * Global map of activity launches to sets of passed content URIs. Specifically:
+     * - The caller didn't have read permission to them.
+     * - The callee activity's {@link android.R.attr#requireContentUriPermissionFromCaller} was set
+     *   to "none".
+     *
+     * <p>This map is used for logging the ContentOrFileUriEventReported message.
+     *
+     * <p>The launch id is the ActivityStarter.Request#hashCode and has to be received from
+     * ActivityStarter to {@link #checkGrantUriPermissionFromIntentUnlocked(int, String, Intent,
+     * int, NeededUriGrants, int, Integer, Integer)}.
+     */
+    @GuardedBy("mLaunchToContentUrisWithoutCallerReadPermission")
+    private final SparseArray<ArraySet<Uri>> mLaunchToContentUrisWithoutCallerReadPermission =
+            new SparseArray<>();
+
     private UriGrantsManagerService() {
         this(SystemServiceManager.ensureSystemDir(), "uri-grants");
     }
@@ -613,7 +634,8 @@
     /** Like checkGrantUriPermission, but takes an Intent. */
     private NeededUriGrants checkGrantUriPermissionFromIntentUnlocked(int callingUid,
             String targetPkg, Intent intent, int mode, NeededUriGrants needed, int targetUserId,
-            @RequiredContentUriPermission Integer requireContentUriPermissionFromCaller) {
+            @RequiredContentUriPermission Integer requireContentUriPermissionFromCaller,
+            Integer requestHashCode) {
         if (DEBUG) Slog.v(TAG,
                 "Checking URI perm to data=" + (intent != null ? intent.getData() : null)
                         + " clip=" + (intent != null ? intent.getClipData() : null)
@@ -635,8 +657,9 @@
         }
 
         if (android.security.Flags.contentUriPermissionApis()) {
-            enforceRequireContentUriPermissionFromCallerOnIntentExtraStream(intent, contentUserHint,
-                    mode, callingUid, requireContentUriPermissionFromCaller);
+            enforceRequireContentUriPermissionFromCallerOnIntentExtraStreamUnlocked(intent,
+                    contentUserHint, mode, callingUid, requireContentUriPermissionFromCaller,
+                    requestHashCode);
         }
 
         Uri data = intent.getData();
@@ -660,8 +683,9 @@
         if (data != null) {
             GrantUri grantUri = GrantUri.resolve(contentUserHint, data, mode);
             if (android.security.Flags.contentUriPermissionApis()) {
-                enforceRequireContentUriPermissionFromCaller(requireContentUriPermissionFromCaller,
-                        grantUri, callingUid);
+                enforceRequireContentUriPermissionFromCallerUnlocked(
+                        requireContentUriPermissionFromCaller, grantUri, callingUid,
+                        requestHashCode);
             }
             targetUid = checkGrantUriPermissionUnlocked(callingUid, targetPkg, grantUri, mode,
                     targetUid);
@@ -678,8 +702,9 @@
                 if (uri != null) {
                     GrantUri grantUri = GrantUri.resolve(contentUserHint, uri, mode);
                     if (android.security.Flags.contentUriPermissionApis()) {
-                        enforceRequireContentUriPermissionFromCaller(
-                                requireContentUriPermissionFromCaller, grantUri, callingUid);
+                        enforceRequireContentUriPermissionFromCallerUnlocked(
+                                requireContentUriPermissionFromCaller, grantUri, callingUid,
+                                requestHashCode);
                     }
                     targetUid = checkGrantUriPermissionUnlocked(callingUid, targetPkg,
                             grantUri, mode, targetUid);
@@ -694,7 +719,7 @@
                     if (clipIntent != null) {
                         NeededUriGrants newNeeded = checkGrantUriPermissionFromIntentUnlocked(
                                 callingUid, targetPkg, clipIntent, mode, needed, targetUserId,
-                                requireContentUriPermissionFromCaller);
+                                requireContentUriPermissionFromCaller, requestHashCode);
                         if (newNeeded != null) {
                             needed = newNeeded;
                         }
@@ -706,17 +731,32 @@
         return needed;
     }
 
-    private void enforceRequireContentUriPermissionFromCaller(
+    private void enforceRequireContentUriPermissionFromCallerUnlocked(
             @RequiredContentUriPermission Integer requireContentUriPermissionFromCaller,
-            GrantUri grantUri, int uid) {
-        // Ignore if requireContentUriPermissionFromCaller hasn't been set or the URI is a
+            GrantUri grantUri, int callingUid, Integer requestHashCode) {
+        // Exit early if requireContentUriPermissionFromCaller hasn't been set or the URI is a
         // non-content URI.
         if (requireContentUriPermissionFromCaller == null
                 || requireContentUriPermissionFromCaller == CONTENT_URI_PERMISSION_NONE
                 || !ContentResolver.SCHEME_CONTENT.equals(grantUri.uri.getScheme())) {
+            tryAddingContentUriWithoutCallerReadPermissionWhenAttributeIsNoneUnlocked(
+                    requireContentUriPermissionFromCaller, grantUri, callingUid, requestHashCode);
             return;
         }
 
+        final boolean hasPermission = hasRequireContentUriPermissionFromCallerUnlocked(
+                requireContentUriPermissionFromCaller, grantUri, callingUid);
+
+        if (!hasPermission) {
+            throw new SecurityException("You can't launch this activity because you don't have the"
+                    + " required " + ActivityInfo.requiredContentUriPermissionToShortString(
+                            requireContentUriPermissionFromCaller) + " access to " + grantUri.uri);
+        }
+    }
+
+    private boolean hasRequireContentUriPermissionFromCallerUnlocked(
+            @RequiredContentUriPermission Integer requireContentUriPermissionFromCaller,
+            GrantUri grantUri, int uid) {
         final boolean readMet = !isRequiredContentUriPermissionRead(
                 requireContentUriPermissionFromCaller)
                 || checkContentUriPermissionFullUnlocked(grantUri, uid,
@@ -727,26 +767,48 @@
                 || checkContentUriPermissionFullUnlocked(grantUri, uid,
                 Intent.FLAG_GRANT_WRITE_URI_PERMISSION);
 
-        boolean hasPermission =
-                requireContentUriPermissionFromCaller == CONTENT_URI_PERMISSION_READ_OR_WRITE
-                        ? (readMet || writeMet) : (readMet && writeMet);
+        return requireContentUriPermissionFromCaller == CONTENT_URI_PERMISSION_READ_OR_WRITE
+                ? (readMet || writeMet) : (readMet && writeMet);
+    }
 
-        if (!hasPermission) {
-            throw new SecurityException("You can't launch this activity because you don't have the"
-                    + " required " + ActivityInfo.requiredContentUriPermissionToShortString(
-                            requireContentUriPermissionFromCaller) + " access to " + grantUri.uri);
+    private void tryAddingContentUriWithoutCallerReadPermissionWhenAttributeIsNoneUnlocked(
+            @RequiredContentUriPermission Integer requireContentUriPermissionFromCaller,
+            GrantUri grantUri, int callingUid, Integer requestHashCode) {
+        // We're interested in requireContentUriPermissionFromCaller that is set to
+        // CONTENT_URI_PERMISSION_NONE and content URIs. Hence, ignore if
+        // requireContentUriPermissionFromCaller is not set to CONTENT_URI_PERMISSION_NONE or the
+        // URI is a non-content URI.
+        if (requireContentUriPermissionFromCaller == null
+                || requireContentUriPermissionFromCaller != CONTENT_URI_PERMISSION_NONE
+                || !ContentResolver.SCHEME_CONTENT.equals(grantUri.uri.getScheme())
+                || requestHashCode == null) {
+            return;
+        }
+
+        if (!hasRequireContentUriPermissionFromCallerUnlocked(CONTENT_URI_PERMISSION_READ, grantUri,
+                callingUid)) {
+            synchronized (mLaunchToContentUrisWithoutCallerReadPermission) {
+                if (mLaunchToContentUrisWithoutCallerReadPermission.get(requestHashCode) == null) {
+                    mLaunchToContentUrisWithoutCallerReadPermission
+                            .put(requestHashCode, new ArraySet<>());
+                }
+                mLaunchToContentUrisWithoutCallerReadPermission.get(requestHashCode)
+                        .add(grantUri.uri);
+            }
         }
     }
 
-    private void enforceRequireContentUriPermissionFromCallerOnIntentExtraStream(Intent intent,
-            int contentUserHint, int mode, int callingUid,
-            @RequiredContentUriPermission Integer requireContentUriPermissionFromCaller) {
+    private void enforceRequireContentUriPermissionFromCallerOnIntentExtraStreamUnlocked(
+            Intent intent, int contentUserHint, int mode, int callingUid,
+            @RequiredContentUriPermission Integer requireContentUriPermissionFromCaller,
+            Integer requestHashCode) {
         try {
             final Uri uri = intent.getParcelableExtra(Intent.EXTRA_STREAM, Uri.class);
             if (uri != null) {
                 final GrantUri grantUri = GrantUri.resolve(contentUserHint, uri, mode);
-                enforceRequireContentUriPermissionFromCaller(
-                        requireContentUriPermissionFromCaller, grantUri, callingUid);
+                enforceRequireContentUriPermissionFromCallerUnlocked(
+                        requireContentUriPermissionFromCaller, grantUri, callingUid,
+                        requestHashCode);
             }
         } catch (BadParcelableException e) {
             Slog.w(TAG, "Failed to unparcel an URI in EXTRA_STREAM, skipping"
@@ -759,8 +821,9 @@
             if (uris != null) {
                 for (int i = uris.size() - 1; i >= 0; i--) {
                     final GrantUri grantUri = GrantUri.resolve(contentUserHint, uris.get(i), mode);
-                    enforceRequireContentUriPermissionFromCaller(
-                            requireContentUriPermissionFromCaller, grantUri, callingUid);
+                    enforceRequireContentUriPermissionFromCallerUnlocked(
+                            requireContentUriPermissionFromCaller, grantUri, callingUid,
+                            requestHashCode);
                 }
             }
         } catch (BadParcelableException e) {
@@ -769,6 +832,37 @@
         }
     }
 
+    private void notifyActivityLaunchRequestCompletedUnlocked(Integer requestHashCode,
+            boolean isSuccessfulLaunch, String intentAction, int callingUid,
+            String callingActivityName, int calleeUid, String calleeActivityName,
+            boolean isStartActivityForResult) {
+        ArraySet<Uri> contentUris;
+        synchronized (mLaunchToContentUrisWithoutCallerReadPermission) {
+            contentUris = mLaunchToContentUrisWithoutCallerReadPermission.get(requestHashCode);
+            mLaunchToContentUrisWithoutCallerReadPermission.remove(requestHashCode);
+        }
+        if (!isSuccessfulLaunch || contentUris == null) return;
+
+        final String[] authorities = new String[contentUris.size()];
+        final String[] schemes = new String[contentUris.size()];
+        for (int i = contentUris.size() - 1; i >= 0; i--) {
+            Uri uri = contentUris.valueAt(i);
+            authorities[i] = uri.getAuthority();
+            schemes[i] = uri.getScheme();
+        }
+        FrameworkStatsLog.write(CONTENT_OR_FILE_URI_EVENT_REPORTED,
+                CONTENT_OR_FILE_URI_EVENT_REPORTED__EVENT_TYPE__CONTENT_URI_WITHOUT_CALLER_READ_PERMISSION,
+                intentAction,
+                callingUid,
+                callingActivityName,
+                calleeUid,
+                calleeActivityName,
+                isStartActivityForResult,
+                String.join(",", authorities),
+                String.join(",", schemes),
+                /* uri_mime_type */ null);
+    }
+
     @GuardedBy("mLock")
     private void readGrantedUriPermissionsLocked() {
         if (DEBUG) Slog.v(TAG, "readGrantedUriPermissions()");
@@ -1645,23 +1739,36 @@
         public NeededUriGrants checkGrantUriPermissionFromIntent(Intent intent, int callingUid,
                 String targetPkg, int targetUserId) {
             return internalCheckGrantUriPermissionFromIntent(intent, callingUid, targetPkg,
-                    targetUserId, /* requireContentUriPermissionFromCaller */ null);
+                    targetUserId, /* requireContentUriPermissionFromCaller */ null,
+                    /* requestHashCode */ null);
         }
 
         @Override
         public NeededUriGrants checkGrantUriPermissionFromIntent(Intent intent, int callingUid,
-                String targetPkg, int targetUserId, int requireContentUriPermissionFromCaller) {
+                String targetPkg, int targetUserId, int requireContentUriPermissionFromCaller,
+                int requestHashCode) {
             return internalCheckGrantUriPermissionFromIntent(intent, callingUid, targetPkg,
-                    targetUserId, requireContentUriPermissionFromCaller);
+                    targetUserId, requireContentUriPermissionFromCaller, requestHashCode);
+        }
+
+        @Override
+        public void notifyActivityLaunchRequestCompleted(int requestHashCode,
+                boolean isSuccessfulLaunch, String intentAction, int callingUid,
+                String callingActivityName, int calleeUid, String calleeActivityName,
+                boolean isStartActivityForResult) {
+            UriGrantsManagerService.this.notifyActivityLaunchRequestCompletedUnlocked(
+                    requestHashCode, isSuccessfulLaunch, intentAction, callingUid,
+                    callingActivityName, calleeUid, calleeActivityName,
+                    isStartActivityForResult);
         }
 
         private NeededUriGrants internalCheckGrantUriPermissionFromIntent(Intent intent,
                 int callingUid, String targetPkg, int targetUserId,
-                @Nullable Integer requireContentUriPermissionFromCaller) {
+                @Nullable Integer requireContentUriPermissionFromCaller, Integer requestHashCode) {
             final int mode = (intent != null) ? intent.getFlags() : 0;
             return UriGrantsManagerService.this.checkGrantUriPermissionFromIntentUnlocked(
                     callingUid, targetPkg, intent, mode, null, targetUserId,
-                    requireContentUriPermissionFromCaller);
+                    requireContentUriPermissionFromCaller, requestHashCode);
         }
 
         @Override
diff --git a/services/core/java/com/android/server/vibrator/VibrationThread.java b/services/core/java/com/android/server/vibrator/VibrationThread.java
index ab4a4d8..4c1e16c 100644
--- a/services/core/java/com/android/server/vibrator/VibrationThread.java
+++ b/services/core/java/com/android/server/vibrator/VibrationThread.java
@@ -128,15 +128,20 @@
      *  before the release callback.
      */
     boolean runVibrationOnVibrationThread(VibrationStepConductor conductor) {
-        synchronized (mLock) {
-            if (mRequestedActiveConductor != null) {
-                Slog.wtf(TAG, "Attempt to start vibration when one already running");
-                return false;
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "runVibrationOnVibrationThread");
+        try {
+            synchronized (mLock) {
+                if (mRequestedActiveConductor != null) {
+                    Slog.wtf(TAG, "Attempt to start vibration when one already running");
+                    return false;
+                }
+                mRequestedActiveConductor = conductor;
+                mLock.notifyAll();
             }
-            mRequestedActiveConductor = conductor;
-            mLock.notifyAll();
+            return true;
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
-        return true;
     }
 
     @Override
diff --git a/services/core/java/com/android/server/vibrator/VibratorController.java b/services/core/java/com/android/server/vibrator/VibratorController.java
index 4fc0b74..3c47850 100644
--- a/services/core/java/com/android/server/vibrator/VibratorController.java
+++ b/services/core/java/com/android/server/vibrator/VibratorController.java
@@ -23,6 +23,7 @@
 import android.os.Parcel;
 import android.os.RemoteCallbackList;
 import android.os.RemoteException;
+import android.os.Trace;
 import android.os.VibrationEffect;
 import android.os.VibratorInfo;
 import android.os.vibrator.PrebakedSegment;
@@ -123,21 +124,26 @@
 
     /** Reruns the query to the vibrator to load the {@link VibratorInfo}, if not yet successful. */
     public void reloadVibratorInfoIfNeeded() {
-        // Early check outside lock, for quick return.
-        if (mVibratorInfoLoadSuccessful) {
-            return;
-        }
-        synchronized (mLock) {
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#reloadVibratorInfoIfNeeded");
+        try {
+            // Early check outside lock, for quick return.
             if (mVibratorInfoLoadSuccessful) {
                 return;
             }
-            int vibratorId = mVibratorInfo.getId();
-            VibratorInfo.Builder vibratorInfoBuilder = new VibratorInfo.Builder(vibratorId);
-            mVibratorInfoLoadSuccessful = mNativeWrapper.getInfo(vibratorInfoBuilder);
-            mVibratorInfo = vibratorInfoBuilder.build();
-            if (!mVibratorInfoLoadSuccessful) {
-                Slog.e(TAG, "Failed retry of HAL getInfo for vibrator " + vibratorId);
+            synchronized (mLock) {
+                if (mVibratorInfoLoadSuccessful) {
+                    return;
+                }
+                int vibratorId = mVibratorInfo.getId();
+                VibratorInfo.Builder vibratorInfoBuilder = new VibratorInfo.Builder(vibratorId);
+                mVibratorInfoLoadSuccessful = mNativeWrapper.getInfo(vibratorInfoBuilder);
+                mVibratorInfo = vibratorInfoBuilder.build();
+                if (!mVibratorInfoLoadSuccessful) {
+                    Slog.e(TAG, "Failed retry of HAL getInfo for vibrator " + vibratorId);
+                }
             }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -193,8 +199,13 @@
 
     /** Return {@code true} if the underlying vibrator is currently available, false otherwise. */
     public boolean isAvailable() {
-        synchronized (mLock) {
-            return mNativeWrapper.isAvailable();
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#isAvailable");
+        try {
+            synchronized (mLock) {
+                return mNativeWrapper.isAvailable();
+            }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -204,12 +215,17 @@
      * <p>This will affect the state of {@link #isUnderExternalControl()}.
      */
     public void setExternalControl(boolean externalControl) {
-        if (!mVibratorInfo.hasCapability(IVibrator.CAP_EXTERNAL_CONTROL)) {
-            return;
-        }
-        synchronized (mLock) {
-            mIsUnderExternalControl = externalControl;
-            mNativeWrapper.setExternalControl(externalControl);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "setExternalControl(" + externalControl + ")");
+        try {
+            if (!mVibratorInfo.hasCapability(IVibrator.CAP_EXTERNAL_CONTROL)) {
+                return;
+            }
+            synchronized (mLock) {
+                mIsUnderExternalControl = externalControl;
+                mNativeWrapper.setExternalControl(externalControl);
+            }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -218,28 +234,38 @@
      * if given {@code effect} is {@code null}.
      */
     public void updateAlwaysOn(int id, @Nullable PrebakedSegment prebaked) {
-        if (!mVibratorInfo.hasCapability(IVibrator.CAP_ALWAYS_ON_CONTROL)) {
-            return;
-        }
-        synchronized (mLock) {
-            if (prebaked == null) {
-                mNativeWrapper.alwaysOnDisable(id);
-            } else {
-                mNativeWrapper.alwaysOnEnable(id, prebaked.getEffectId(),
-                        prebaked.getEffectStrength());
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#updateAlwaysOn");
+        try {
+            if (!mVibratorInfo.hasCapability(IVibrator.CAP_ALWAYS_ON_CONTROL)) {
+                return;
             }
+            synchronized (mLock) {
+                if (prebaked == null) {
+                    mNativeWrapper.alwaysOnDisable(id);
+                } else {
+                    mNativeWrapper.alwaysOnEnable(id, prebaked.getEffectId(),
+                            prebaked.getEffectStrength());
+                }
+            }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
     /** Set the vibration amplitude. This will NOT affect the state of {@link #isVibrating()}. */
     public void setAmplitude(float amplitude) {
-        synchronized (mLock) {
-            if (mVibratorInfo.hasCapability(IVibrator.CAP_AMPLITUDE_CONTROL)) {
-                mNativeWrapper.setAmplitude(amplitude);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#setAmplitude");
+        try {
+            synchronized (mLock) {
+                if (mVibratorInfo.hasCapability(IVibrator.CAP_AMPLITUDE_CONTROL)) {
+                    mNativeWrapper.setAmplitude(amplitude);
+                }
+                if (mIsVibrating) {
+                    mCurrentAmplitude = amplitude;
+                }
             }
-            if (mIsVibrating) {
-                mCurrentAmplitude = amplitude;
-            }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -253,13 +279,18 @@
      * do not support the input or a negative number if the operation failed.
      */
     public long on(long milliseconds, long vibrationId) {
-        synchronized (mLock) {
-            long duration = mNativeWrapper.on(milliseconds, vibrationId);
-            if (duration > 0) {
-                mCurrentAmplitude = -1;
-                notifyListenerOnVibrating(true);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#on");
+        try {
+            synchronized (mLock) {
+                long duration = mNativeWrapper.on(milliseconds, vibrationId);
+                if (duration > 0) {
+                    mCurrentAmplitude = -1;
+                    notifyListenerOnVibrating(true);
+                }
+                return duration;
             }
-            return duration;
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -273,6 +304,7 @@
      * do not support the input or a negative number if the operation failed.
      */
     public long on(VibrationEffect.VendorEffect vendorEffect, long vibrationId) {
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#on (vendor)");
         synchronized (mLock) {
             Parcel vendorData = Parcel.obtain();
             try {
@@ -288,6 +320,7 @@
                 return duration;
             } finally {
                 vendorData.recycle();
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
         }
     }
@@ -302,14 +335,19 @@
      * do not support the input or a negative number if the operation failed.
      */
     public long on(PrebakedSegment prebaked, long vibrationId) {
-        synchronized (mLock) {
-            long duration = mNativeWrapper.perform(prebaked.getEffectId(),
-                    prebaked.getEffectStrength(), vibrationId);
-            if (duration > 0) {
-                mCurrentAmplitude = -1;
-                notifyListenerOnVibrating(true);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#on (Prebaked)");
+        try {
+            synchronized (mLock) {
+                long duration = mNativeWrapper.perform(prebaked.getEffectId(),
+                        prebaked.getEffectStrength(), vibrationId);
+                if (duration > 0) {
+                    mCurrentAmplitude = -1;
+                    notifyListenerOnVibrating(true);
+                }
+                return duration;
             }
-            return duration;
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -323,16 +361,21 @@
      * do not support the input or a negative number if the operation failed.
      */
     public long on(PrimitiveSegment[] primitives, long vibrationId) {
-        if (!mVibratorInfo.hasCapability(IVibrator.CAP_COMPOSE_EFFECTS)) {
-            return 0;
-        }
-        synchronized (mLock) {
-            long duration = mNativeWrapper.compose(primitives, vibrationId);
-            if (duration > 0) {
-                mCurrentAmplitude = -1;
-                notifyListenerOnVibrating(true);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#on (Primitive)");
+        try {
+            if (!mVibratorInfo.hasCapability(IVibrator.CAP_COMPOSE_EFFECTS)) {
+                return 0;
             }
-            return duration;
+            synchronized (mLock) {
+                long duration = mNativeWrapper.compose(primitives, vibrationId);
+                if (duration > 0) {
+                    mCurrentAmplitude = -1;
+                    notifyListenerOnVibrating(true);
+                }
+                return duration;
+            }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -345,17 +388,22 @@
      * @return The duration of the effect playing, or 0 if unsupported.
      */
     public long on(RampSegment[] primitives, long vibrationId) {
-        if (!mVibratorInfo.hasCapability(IVibrator.CAP_COMPOSE_PWLE_EFFECTS)) {
-            return 0;
-        }
-        synchronized (mLock) {
-            int braking = mVibratorInfo.getDefaultBraking();
-            long duration = mNativeWrapper.composePwle(primitives, braking, vibrationId);
-            if (duration > 0) {
-                mCurrentAmplitude = -1;
-                notifyListenerOnVibrating(true);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#on (PWLE)");
+        try {
+            if (!mVibratorInfo.hasCapability(IVibrator.CAP_COMPOSE_PWLE_EFFECTS)) {
+                return 0;
             }
-            return duration;
+            synchronized (mLock) {
+                int braking = mVibratorInfo.getDefaultBraking();
+                long duration = mNativeWrapper.composePwle(primitives, braking, vibrationId);
+                if (duration > 0) {
+                    mCurrentAmplitude = -1;
+                    notifyListenerOnVibrating(true);
+                }
+                return duration;
+            }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
@@ -365,10 +413,15 @@
      * <p>This will affect the state of {@link #isVibrating()}.
      */
     public void off() {
-        synchronized (mLock) {
-            mNativeWrapper.off();
-            mCurrentAmplitude = 0;
-            notifyListenerOnVibrating(false);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "VibratorController#off");
+        try {
+            synchronized (mLock) {
+                mNativeWrapper.off();
+                mCurrentAmplitude = 0;
+                notifyListenerOnVibrating(false);
+            }
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
         }
     }
 
diff --git a/services/core/java/com/android/server/vibrator/VibratorManagerService.java b/services/core/java/com/android/server/vibrator/VibratorManagerService.java
index 799934a..899f0b1 100644
--- a/services/core/java/com/android/server/vibrator/VibratorManagerService.java
+++ b/services/core/java/com/android/server/vibrator/VibratorManagerService.java
@@ -462,20 +462,31 @@
     @Override // Binder call
     public void performHapticFeedback(int uid, int deviceId, String opPkg, int constant,
             String reason, int flags, int privFlags) {
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "performHapticFeedback");
         // Note that the `performHapticFeedback` method does not take a token argument from the
         // caller, and instead, uses this service as the token. This is to mitigate performance
         // impact that would otherwise be caused due to marshal latency. Haptic feedback effects are
         // short-lived, so we don't need to cancel when the process dies.
-        performHapticFeedbackInternal(uid, deviceId, opPkg, constant, reason, /* token= */
-                this, flags, privFlags);
+        try {
+            performHapticFeedbackInternal(uid, deviceId, opPkg, constant, reason, /* token= */
+                    this, flags, privFlags);
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+        }
     }
 
     @Override // Binder call
     public void performHapticFeedbackForInputDevice(int uid, int deviceId, String opPkg,
             int constant, int inputDeviceId, int inputSource, String reason, int flags,
             int privFlags) {
-        performHapticFeedbackForInputDeviceInternal(uid, deviceId, opPkg, constant, inputDeviceId,
-                inputSource, reason, /* token= */ this, flags, privFlags);
+        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "performHapticFeedbackForInputDevice");
+        try {
+            performHapticFeedbackForInputDeviceInternal(uid, deviceId, opPkg, constant,
+                    inputDeviceId,
+                    inputSource, reason, /* token= */ this, flags, privFlags);
+        } finally {
+            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+        }
     }
 
     /**
@@ -919,30 +930,25 @@
     @GuardedBy("mLock")
     @Nullable
     private Vibration.EndInfo startVibrationOnThreadLocked(VibrationStepConductor conductor) {
-        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "startVibrationThreadLocked");
-        try {
-            HalVibration vib = conductor.getVibration();
-            int mode = startAppOpModeLocked(vib.callerInfo);
-            switch (mode) {
-                case AppOpsManager.MODE_ALLOWED:
-                    Trace.asyncTraceBegin(Trace.TRACE_TAG_VIBRATOR, "vibration", 0);
-                    // Make sure mCurrentVibration is set while triggering the VibrationThread.
-                    mCurrentVibration = conductor;
-                    if (!mVibrationThread.runVibrationOnVibrationThread(mCurrentVibration)) {
-                        // Shouldn't happen. The method call already logs a wtf.
-                        mCurrentVibration = null;  // Aborted.
-                        return new Vibration.EndInfo(Status.IGNORED_ERROR_SCHEDULING);
-                    }
-                    return null;
-                case AppOpsManager.MODE_ERRORED:
-                    Slog.w(TAG, "Start AppOpsManager operation errored for uid "
-                            + vib.callerInfo.uid);
-                    return new Vibration.EndInfo(Status.IGNORED_ERROR_APP_OPS);
-                default:
-                    return new Vibration.EndInfo(Status.IGNORED_APP_OPS);
-            }
-        } finally {
-            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+        HalVibration vib = conductor.getVibration();
+        int mode = startAppOpModeLocked(vib.callerInfo);
+        switch (mode) {
+            case AppOpsManager.MODE_ALLOWED:
+                Trace.asyncTraceBegin(Trace.TRACE_TAG_VIBRATOR, "vibration", 0);
+                // Make sure mCurrentVibration is set while triggering the VibrationThread.
+                mCurrentVibration = conductor;
+                if (!mVibrationThread.runVibrationOnVibrationThread(mCurrentVibration)) {
+                    // Shouldn't happen. The method call already logs a wtf.
+                    mCurrentVibration = null;  // Aborted.
+                    return new Vibration.EndInfo(Status.IGNORED_ERROR_SCHEDULING);
+                }
+                return null;
+            case AppOpsManager.MODE_ERRORED:
+                Slog.w(TAG, "Start AppOpsManager operation errored for uid "
+                        + vib.callerInfo.uid);
+                return new Vibration.EndInfo(Status.IGNORED_ERROR_APP_OPS);
+            default:
+                return new Vibration.EndInfo(Status.IGNORED_APP_OPS);
         }
     }
 
@@ -1050,21 +1056,16 @@
 
     @GuardedBy("mLock")
     private void reportFinishedVibrationLocked(Vibration.EndInfo vibrationEndInfo) {
-        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "reportFinishVibrationLocked");
         Trace.asyncTraceEnd(Trace.TRACE_TAG_VIBRATOR, "vibration", 0);
-        try {
-            HalVibration vib = mCurrentVibration.getVibration();
-            if (DEBUG) {
-                Slog.d(TAG, "Reporting vibration " + vib.id + " finished with "
-                        + vibrationEndInfo);
-            }
-            // DO NOT write metrics at this point, wait for the VibrationThread to report the
-            // vibration was released, after all cleanup. The metrics will be reported then.
-            endVibrationLocked(vib, vibrationEndInfo, /* shouldWriteStats= */ false);
-            finishAppOpModeLocked(vib.callerInfo);
-        } finally {
-            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+        HalVibration vib = mCurrentVibration.getVibration();
+        if (DEBUG) {
+            Slog.d(TAG, "Reporting vibration " + vib.id + " finished with "
+                    + vibrationEndInfo);
         }
+        // DO NOT write metrics at this point, wait for the VibrationThread to report the
+        // vibration was released, after all cleanup. The metrics will be reported then.
+        endVibrationLocked(vib, vibrationEndInfo, /* shouldWriteStats= */ false);
+        finishAppOpModeLocked(vib.callerInfo);
     }
 
     private void onSyncedVibrationComplete(long vibrationId) {
@@ -1418,40 +1419,34 @@
 
     @GuardedBy("mLock")
     @Nullable
-    private SparseArray<PrebakedSegment> fixupAlwaysOnEffectsLocked(
-            CombinedVibration effect) {
-        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "fixupAlwaysOnEffectsLocked");
-        try {
-            SparseArray<VibrationEffect> effects;
-            if (effect instanceof CombinedVibration.Mono) {
-                VibrationEffect syncedEffect = ((CombinedVibration.Mono) effect).getEffect();
-                effects = transformAllVibratorsLocked(unused -> syncedEffect);
-            } else if (effect instanceof CombinedVibration.Stereo) {
-                effects = ((CombinedVibration.Stereo) effect).getEffects();
-            } else {
-                // Only synced combinations can be used for always-on effects.
-                return null;
-            }
-            SparseArray<PrebakedSegment> result = new SparseArray<>();
-            for (int i = 0; i < effects.size(); i++) {
-                PrebakedSegment prebaked = extractPrebakedSegment(effects.valueAt(i));
-                if (prebaked == null) {
-                    Slog.e(TAG, "Only prebaked effects supported for always-on.");
-                    return null;
-                }
-                int vibratorId = effects.keyAt(i);
-                VibratorController vibrator = mVibrators.get(vibratorId);
-                if (vibrator != null && vibrator.hasCapability(IVibrator.CAP_ALWAYS_ON_CONTROL)) {
-                    result.put(vibratorId, prebaked);
-                }
-            }
-            if (result.size() == 0) {
-                return null;
-            }
-            return result;
-        } finally {
-            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+    private SparseArray<PrebakedSegment> fixupAlwaysOnEffectsLocked(CombinedVibration effect) {
+        SparseArray<VibrationEffect> effects;
+        if (effect instanceof CombinedVibration.Mono) {
+            VibrationEffect syncedEffect = ((CombinedVibration.Mono) effect).getEffect();
+            effects = transformAllVibratorsLocked(unused -> syncedEffect);
+        } else if (effect instanceof CombinedVibration.Stereo) {
+            effects = ((CombinedVibration.Stereo) effect).getEffects();
+        } else {
+            // Only synced combinations can be used for always-on effects.
+            return null;
         }
+        SparseArray<PrebakedSegment> result = new SparseArray<>();
+        for (int i = 0; i < effects.size(); i++) {
+            PrebakedSegment prebaked = extractPrebakedSegment(effects.valueAt(i));
+            if (prebaked == null) {
+                Slog.e(TAG, "Only prebaked effects supported for always-on.");
+                return null;
+            }
+            int vibratorId = effects.keyAt(i);
+            VibratorController vibrator = mVibrators.get(vibratorId);
+            if (vibrator != null && vibrator.hasCapability(IVibrator.CAP_ALWAYS_ON_CONTROL)) {
+                result.put(vibratorId, prebaked);
+            }
+        }
+        if (result.size() == 0) {
+            return null;
+        }
+        return result;
     }
 
     @Nullable
@@ -1580,25 +1575,42 @@
 
         @Override
         public boolean prepareSyncedVibration(long requiredCapabilities, int[] vibratorIds) {
-            if ((mCapabilities & requiredCapabilities) != requiredCapabilities) {
-                // This sync step requires capabilities this device doesn't have, skipping sync...
-                return false;
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "prepareSyncedVibration");
+            try {
+                if ((mCapabilities & requiredCapabilities) != requiredCapabilities) {
+                    // This sync step requires capabilities this device doesn't have, skipping
+                    // sync...
+                    return false;
+                }
+                return mNativeWrapper.prepareSynced(vibratorIds);
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
-            return mNativeWrapper.prepareSynced(vibratorIds);
         }
 
         @Override
         public boolean triggerSyncedVibration(long vibrationId) {
-            return mNativeWrapper.triggerSynced(vibrationId);
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "triggerSyncedVibration");
+            try {
+                return mNativeWrapper.triggerSynced(vibrationId);
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+            }
         }
 
         @Override
         public void cancelSyncedVibration() {
-            mNativeWrapper.cancelSynced();
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "cancelSyncedVibration");
+            try {
+                mNativeWrapper.cancelSynced();
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+            }
         }
 
         @Override
         public void noteVibratorOn(int uid, long duration) {
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "noteVibratorOn");
             try {
                 if (duration <= 0) {
                     // Tried to turn vibrator ON and got:
@@ -1616,16 +1628,21 @@
                 mFrameworkStatsLogger.writeVibratorStateOnAsync(uid, duration);
             } catch (RemoteException e) {
                 Slog.e(TAG, "Error logging VibratorStateChanged to ON", e);
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
         }
 
         @Override
         public void noteVibratorOff(int uid) {
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "noteVibratorOff");
             try {
                 mBatteryStatsService.noteVibratorOff(uid);
                 mFrameworkStatsLogger.writeVibratorStateOffAsync(uid);
             } catch (RemoteException e) {
                 Slog.e(TAG, "Error logging VibratorStateChanged to OFF", e);
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
         }
 
@@ -1634,11 +1651,16 @@
             if (DEBUG) {
                 Slog.d(TAG, "Vibration " + vibrationId + " finished with " + vibrationEndInfo);
             }
-            synchronized (mLock) {
-                if (mCurrentVibration != null
-                        && mCurrentVibration.getVibration().id == vibrationId) {
-                    reportFinishedVibrationLocked(vibrationEndInfo);
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "onVibrationCompleted");
+            try {
+                synchronized (mLock) {
+                    if (mCurrentVibration != null
+                            && mCurrentVibration.getVibration().id == vibrationId) {
+                        reportFinishedVibrationLocked(vibrationEndInfo);
+                    }
                 }
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
         }
 
@@ -1647,34 +1669,40 @@
             if (DEBUG) {
                 Slog.d(TAG, "VibrationThread released after finished vibration");
             }
-            synchronized (mLock) {
-                if (DEBUG) {
-                    Slog.d(TAG, "Processing VibrationThread released callback");
-                }
-                if (Build.IS_DEBUGGABLE && mCurrentVibration != null
-                        && mCurrentVibration.getVibration().id != vibrationId) {
-                    Slog.wtf(TAG, TextUtils.formatSimple(
-                            "VibrationId mismatch on release. expected=%d, released=%d",
-                            mCurrentVibration.getVibration().id, vibrationId));
-                }
-                if (mCurrentVibration != null) {
-                    // This is when we consider the current vibration complete, so report metrics.
-                    mFrameworkStatsLogger.writeVibrationReportedAsync(
-                            mCurrentVibration.getVibration().getStatsInfo(
-                                    /* completionUptimeMillis= */ SystemClock.uptimeMillis()));
-                    mCurrentVibration = null;
-                }
-                if (mNextVibration != null) {
-                    VibrationStepConductor nextConductor = mNextVibration;
-                    mNextVibration = null;
-                    Vibration.EndInfo vibrationEndInfo = startVibrationOnThreadLocked(
-                            nextConductor);
-                    if (vibrationEndInfo != null) {
-                        // Failed to start the vibration, end it and report metrics right away.
-                        endVibrationLocked(nextConductor.getVibration(),
-                                vibrationEndInfo, /* shouldWriteStats= */ true);
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "onVibrationThreadReleased: " + vibrationId);
+            try {
+                synchronized (mLock) {
+                    if (DEBUG) {
+                        Slog.d(TAG, "Processing VibrationThread released callback");
+                    }
+                    if (Build.IS_DEBUGGABLE && mCurrentVibration != null
+                            && mCurrentVibration.getVibration().id != vibrationId) {
+                        Slog.wtf(TAG, TextUtils.formatSimple(
+                                "VibrationId mismatch on release. expected=%d, released=%d",
+                                mCurrentVibration.getVibration().id, vibrationId));
+                    }
+                    if (mCurrentVibration != null) {
+                        // This is when we consider the current vibration complete, so report
+                        // metrics.
+                        mFrameworkStatsLogger.writeVibrationReportedAsync(
+                                mCurrentVibration.getVibration().getStatsInfo(
+                                        /* completionUptimeMillis= */ SystemClock.uptimeMillis()));
+                        mCurrentVibration = null;
+                    }
+                    if (mNextVibration != null) {
+                        VibrationStepConductor nextConductor = mNextVibration;
+                        mNextVibration = null;
+                        Vibration.EndInfo vibrationEndInfo = startVibrationOnThreadLocked(
+                                nextConductor);
+                        if (vibrationEndInfo != null) {
+                            // Failed to start the vibration, end it and report metrics right away.
+                            endVibrationLocked(nextConductor.getVibration(),
+                                    vibrationEndInfo, /* shouldWriteStats= */ true);
+                        }
                     }
                 }
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
         }
     }
@@ -1917,22 +1945,17 @@
     @GuardedBy("mLock")
     private void endExternalVibrateLocked(Vibration.EndInfo vibrationEndInfo,
             boolean continueExternalControl) {
-        Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "endExternalVibrateLocked");
-        try {
-            if (mCurrentExternalVibration == null) {
-                return;
-            }
-            mCurrentExternalVibration.unlinkToDeath();
-            if (!continueExternalControl) {
-                setExternalControl(false, mCurrentExternalVibration.stats);
-            }
-            // The external control was turned off, end it and report metrics right away.
-            endVibrationLocked(mCurrentExternalVibration, vibrationEndInfo,
-                    /* shouldWriteStats= */ true);
-            mCurrentExternalVibration = null;
-        } finally {
-            Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
+        if (mCurrentExternalVibration == null) {
+            return;
         }
+        mCurrentExternalVibration.unlinkToDeath();
+        if (!continueExternalControl) {
+            setExternalControl(false, mCurrentExternalVibration.stats);
+        }
+        // The external control was turned off, end it and report metrics right away.
+        endVibrationLocked(mCurrentExternalVibration, vibrationEndInfo,
+                /* shouldWriteStats= */ true);
+        mCurrentExternalVibration = null;
     }
 
     private HapticFeedbackVibrationProvider getHapticVibrationProvider() {
@@ -1987,143 +2010,160 @@
 
         @Override
         public ExternalVibrationScale onExternalVibrationStart(ExternalVibration vib) {
-            // Create Vibration.Stats as close to the received request as possible, for tracking.
-            ExternalVibrationSession externalVibration = new ExternalVibrationSession(vib);
-            // Mute the request until we run all the checks and accept the vibration.
-            externalVibration.muteScale();
-            boolean alreadyUnderExternalControl = false;
-            boolean waitForCompletion = false;
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "onExternalVibrationStart");
+            try {
+                // Create Vibration.Stats as close to the received request as possible, for
+                // tracking.
+                ExternalVibrationSession externalVibration = new ExternalVibrationSession(vib);
+                // Mute the request until we run all the checks and accept the vibration.
+                externalVibration.muteScale();
+                boolean alreadyUnderExternalControl = false;
+                boolean waitForCompletion = false;
 
-            synchronized (mLock) {
-                if (!hasExternalControlCapability()) {
-                    endVibrationLocked(externalVibration,
-                            new Vibration.EndInfo(Status.IGNORED_UNSUPPORTED),
-                            /* shouldWriteStats= */ true);
-                    return externalVibration.getScale();
-                }
+                synchronized (mLock) {
+                    if (!hasExternalControlCapability()) {
+                        endVibrationLocked(externalVibration,
+                                new Vibration.EndInfo(Status.IGNORED_UNSUPPORTED),
+                                /* shouldWriteStats= */ true);
+                        return externalVibration.getScale();
+                    }
 
-                if (ActivityManager.checkComponentPermission(android.Manifest.permission.VIBRATE,
-                        vib.getUid(), -1 /*owningUid*/, true /*exported*/)
-                        != PackageManager.PERMISSION_GRANTED) {
-                    Slog.w(TAG, "pkg=" + vib.getPackage() + ", uid=" + vib.getUid()
-                            + " tried to play externally controlled vibration"
-                            + " without VIBRATE permission, ignoring.");
-                    endVibrationLocked(externalVibration,
-                            new Vibration.EndInfo(Status.IGNORED_MISSING_PERMISSION),
-                            /* shouldWriteStats= */ true);
-                    return externalVibration.getScale();
-                }
+                    if (ActivityManager.checkComponentPermission(
+                            android.Manifest.permission.VIBRATE,
+                            vib.getUid(), -1 /*owningUid*/, true /*exported*/)
+                            != PackageManager.PERMISSION_GRANTED) {
+                        Slog.w(TAG, "pkg=" + vib.getPackage() + ", uid=" + vib.getUid()
+                                + " tried to play externally controlled vibration"
+                                + " without VIBRATE permission, ignoring.");
+                        endVibrationLocked(externalVibration,
+                                new Vibration.EndInfo(Status.IGNORED_MISSING_PERMISSION),
+                                /* shouldWriteStats= */ true);
+                        return externalVibration.getScale();
+                    }
 
-                Vibration.EndInfo vibrationEndInfo = shouldIgnoreVibrationLocked(
-                        externalVibration.callerInfo);
+                    Vibration.EndInfo vibrationEndInfo = shouldIgnoreVibrationLocked(
+                            externalVibration.callerInfo);
 
-                if (vibrationEndInfo == null
-                        && mCurrentExternalVibration != null
-                        && mCurrentExternalVibration.isHoldingSameVibration(vib)) {
-                    // We are already playing this external vibration, so we can return the same
-                    // scale calculated in the previous call to this method.
-                    return mCurrentExternalVibration.getScale();
-                }
+                    if (vibrationEndInfo == null
+                            && mCurrentExternalVibration != null
+                            && mCurrentExternalVibration.isHoldingSameVibration(vib)) {
+                        // We are already playing this external vibration, so we can return the same
+                        // scale calculated in the previous call to this method.
+                        return mCurrentExternalVibration.getScale();
+                    }
 
-                if (vibrationEndInfo == null) {
-                    // Check if ongoing vibration is more important than this vibration.
-                    vibrationEndInfo = shouldIgnoreVibrationForOngoingLocked(externalVibration);
-                }
+                    if (vibrationEndInfo == null) {
+                        // Check if ongoing vibration is more important than this vibration.
+                        vibrationEndInfo = shouldIgnoreVibrationForOngoingLocked(externalVibration);
+                    }
 
-                if (vibrationEndInfo != null) {
-                    endVibrationLocked(externalVibration, vibrationEndInfo,
-                            /* shouldWriteStats= */ true);
-                    return externalVibration.getScale();
-                }
+                    if (vibrationEndInfo != null) {
+                        endVibrationLocked(externalVibration, vibrationEndInfo,
+                                /* shouldWriteStats= */ true);
+                        return externalVibration.getScale();
+                    }
 
-                if (mCurrentExternalVibration == null) {
-                    // If we're not under external control right now, then cancel any normal
-                    // vibration that may be playing and ready the vibrator for external control.
-                    if (mCurrentVibration != null) {
+                    if (mCurrentExternalVibration == null) {
+                        // If we're not under external control right now, then cancel any normal
+                        // vibration that may be playing and ready the vibrator for external
+                        // control.
+                        if (mCurrentVibration != null) {
+                            externalVibration.stats.reportInterruptedAnotherVibration(
+                                    mCurrentVibration.getVibration().callerInfo);
+                            clearNextVibrationLocked(
+                                    new Vibration.EndInfo(Status.IGNORED_FOR_EXTERNAL,
+                                            externalVibration.callerInfo));
+                            mCurrentVibration.notifyCancelled(
+                                    new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
+                                            externalVibration.callerInfo),
+                                    /* immediate= */ true);
+                            waitForCompletion = true;
+                        }
+                    } else {
+                        // At this point we have an externally controlled vibration playing already.
+                        // Since the interface defines that only one externally controlled
+                        // vibration can
+                        // play at a time, we need to first mute the ongoing vibration and then
+                        // return
+                        // a scale from this function for the new one, so we can be assured that the
+                        // ongoing will be muted in favor of the new vibration.
+                        //
+                        // Note that this doesn't support multiple concurrent external controls,
+                        // as we would need to mute the old one still if it came from a different
+                        // controller.
+                        alreadyUnderExternalControl = true;
+                        mCurrentExternalVibration.notifyEnded();
                         externalVibration.stats.reportInterruptedAnotherVibration(
-                                mCurrentVibration.getVibration().callerInfo);
-                        clearNextVibrationLocked(
-                                new Vibration.EndInfo(Status.IGNORED_FOR_EXTERNAL,
-                                        externalVibration.callerInfo));
-                        mCurrentVibration.notifyCancelled(
+                                mCurrentExternalVibration.callerInfo);
+                        endExternalVibrateLocked(
                                 new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
                                         externalVibration.callerInfo),
-                                /* immediate= */ true);
-                        waitForCompletion = true;
+                                /* continueExternalControl= */ true);
                     }
-                } else {
-                    // At this point we have an externally controlled vibration playing already.
-                    // Since the interface defines that only one externally controlled vibration can
-                    // play at a time, we need to first mute the ongoing vibration and then return
-                    // a scale from this function for the new one, so we can be assured that the
-                    // ongoing will be muted in favor of the new vibration.
-                    //
-                    // Note that this doesn't support multiple concurrent external controls, as we
-                    // would need to mute the old one still if it came from a different controller.
-                    alreadyUnderExternalControl = true;
-                    mCurrentExternalVibration.notifyEnded();
-                    externalVibration.stats.reportInterruptedAnotherVibration(
-                            mCurrentExternalVibration.callerInfo);
-                    endExternalVibrateLocked(
-                            new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
-                                    externalVibration.callerInfo),
-                            /* continueExternalControl= */ true);
-                }
 
-                VibrationAttributes attrs = fixupVibrationAttributes(vib.getVibrationAttributes(),
-                        /* effect= */ null);
-                if (attrs.isFlagSet(VibrationAttributes.FLAG_INVALIDATE_SETTINGS_CACHE)) {
-                    // Force update of user settings before checking if this vibration effect should
-                    // be ignored or scaled.
-                    mVibrationSettings.update();
-                }
-
-                mCurrentExternalVibration = externalVibration;
-                externalVibration.linkToDeath(this::onExternalVibrationBinderDied);
-                externalVibration.scale(mVibrationScaler, attrs.getUsage());
-            }
-
-            if (waitForCompletion) {
-                if (!mVibrationThread.waitForThreadIdle(VIBRATION_CANCEL_WAIT_MILLIS)) {
-                    Slog.e(TAG, "Timed out waiting for vibration to cancel");
-                    synchronized (mLock) {
-                        // Trigger endExternalVibrateLocked to unlink to death recipient.
-                        endExternalVibrateLocked(
-                                new Vibration.EndInfo(Status.IGNORED_ERROR_CANCELLING),
-                                /* continueExternalControl= */ false);
-                        // Mute the request, vibration will be ignored.
-                        externalVibration.muteScale();
+                    VibrationAttributes attrs = fixupVibrationAttributes(
+                            vib.getVibrationAttributes(),
+                            /* effect= */ null);
+                    if (attrs.isFlagSet(VibrationAttributes.FLAG_INVALIDATE_SETTINGS_CACHE)) {
+                        // Force update of user settings before checking if this vibration effect
+                        // should be ignored or scaled.
+                        mVibrationSettings.update();
                     }
-                    return externalVibration.getScale();
+
+                    mCurrentExternalVibration = externalVibration;
+                    externalVibration.linkToDeath(this::onExternalVibrationBinderDied);
+                    externalVibration.scale(mVibrationScaler, attrs.getUsage());
                 }
-            }
-            if (!alreadyUnderExternalControl) {
+
+                if (waitForCompletion) {
+                    if (!mVibrationThread.waitForThreadIdle(VIBRATION_CANCEL_WAIT_MILLIS)) {
+                        Slog.e(TAG, "Timed out waiting for vibration to cancel");
+                        synchronized (mLock) {
+                            // Trigger endExternalVibrateLocked to unlink to death recipient.
+                            endExternalVibrateLocked(
+                                    new Vibration.EndInfo(Status.IGNORED_ERROR_CANCELLING),
+                                    /* continueExternalControl= */ false);
+                            // Mute the request, vibration will be ignored.
+                            externalVibration.muteScale();
+                        }
+                        return externalVibration.getScale();
+                    }
+                }
+                if (!alreadyUnderExternalControl) {
+                    if (DEBUG) {
+                        Slog.d(TAG, "Vibrator going under external control.");
+                    }
+                    setExternalControl(true, externalVibration.stats);
+                }
                 if (DEBUG) {
-                    Slog.d(TAG, "Vibrator going under external control.");
+                    Slog.d(TAG, "Playing external vibration: " + vib);
                 }
-                setExternalControl(true, externalVibration.stats);
+                // Vibrator will start receiving data from external channels after this point.
+                // Report current time as the vibration start time, for debugging.
+                externalVibration.stats.reportStarted();
+                return externalVibration.getScale();
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
-            if (DEBUG) {
-                Slog.d(TAG, "Playing external vibration: " + vib);
-            }
-            // Vibrator will start receiving data from external channels after this point.
-            // Report current time as the vibration start time, for debugging.
-            externalVibration.stats.reportStarted();
-            return externalVibration.getScale();
         }
 
         @Override
         public void onExternalVibrationStop(ExternalVibration vib) {
-            synchronized (mLock) {
-                if (mCurrentExternalVibration != null
-                        && mCurrentExternalVibration.isHoldingSameVibration(vib)) {
-                    if (DEBUG) {
-                        Slog.d(TAG, "Stopping external vibration: " + vib);
+            Trace.traceBegin(Trace.TRACE_TAG_VIBRATOR, "onExternalVibrationStop");
+            try {
+                synchronized (mLock) {
+                    if (mCurrentExternalVibration != null
+                            && mCurrentExternalVibration.isHoldingSameVibration(vib)) {
+                        if (DEBUG) {
+                            Slog.d(TAG, "Stopping external vibration: " + vib);
+                        }
+                        endExternalVibrateLocked(
+                                new Vibration.EndInfo(Status.FINISHED),
+                                /* continueExternalControl= */ false);
                     }
-                    endExternalVibrateLocked(
-                            new Vibration.EndInfo(Status.FINISHED),
-                            /* continueExternalControl= */ false);
                 }
+            } finally {
+                Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
             }
         }
 
diff --git a/services/core/java/com/android/server/wm/ActivityStarter.java b/services/core/java/com/android/server/wm/ActivityStarter.java
index 1822a80..bc11bac 100644
--- a/services/core/java/com/android/server/wm/ActivityStarter.java
+++ b/services/core/java/com/android/server/wm/ActivityStarter.java
@@ -603,7 +603,8 @@
                             .checkGrantUriPermissionFromIntent(intent, resolvedCallingUid,
                                     activityInfo.applicationInfo.packageName,
                                     UserHandle.getUserId(activityInfo.applicationInfo.uid),
-                                    activityInfo.requireContentUriPermissionFromCaller);
+                                    activityInfo.requireContentUriPermissionFromCaller,
+                                    /* requestHashCode */ this.hashCode());
                 } else {
                     intentGrants = supervisor.mService.mUgmInternal
                             .checkGrantUriPermissionFromIntent(intent, resolvedCallingUid,
@@ -717,6 +718,9 @@
      * @return The starter result.
      */
     int execute() {
+        // Required for logging ContentOrFileUriEventReported in the finally block.
+        String callerActivityName = null;
+        ActivityRecord launchingRecord = null;
         try {
             onExecutionStarted();
 
@@ -737,6 +741,7 @@
                         ?  Binder.getCallingUid() : mRequest.realCallingUid;
                 launchingState = mSupervisor.getActivityMetricsLogger().notifyActivityLaunching(
                         mRequest.intent, caller, callingUid);
+                callerActivityName = caller != null ? caller.info.name : null;
             }
 
             if (mRequest.intent != null) {
@@ -812,7 +817,7 @@
                 final ActivityOptions originalOptions = mRequest.activityOptions != null
                         ? mRequest.activityOptions.getOriginalOptions() : null;
                 // Only track the launch time of activity that will be resumed.
-                final ActivityRecord launchingRecord = mDoResume ? mLastStartActivityRecord : null;
+                launchingRecord = mDoResume ? mLastStartActivityRecord : null;
                 // If the new record is the one that started, a new activity has created.
                 final boolean newActivityCreated = mStartActivity == launchingRecord;
                 // Notify ActivityMetricsLogger that the activity has launched.
@@ -828,6 +833,23 @@
                 return getExternalResult(res);
             }
         } finally {
+            // Notify UriGrantsManagerService that activity launch completed. Required for logging
+            // the ContentOrFileUriEventReported message.
+            mSupervisor.mService.mUgmInternal.notifyActivityLaunchRequestCompleted(
+                    mRequest.hashCode(),
+                    // isSuccessfulLaunch
+                    launchingRecord != null,
+                    // Intent action
+                    mRequest.intent != null ? mRequest.intent.getAction() : null,
+                    mRequest.realCallingUid,
+                    callerActivityName,
+                    // Callee UID
+                    mRequest.activityInfo != null
+                            ? mRequest.activityInfo.applicationInfo.uid : INVALID_UID,
+                    // Callee Activity name
+                    mRequest.activityInfo != null ? mRequest.activityInfo.name : null,
+                    // isStartActivityForResult
+                    launchingRecord != null && launchingRecord.resultTo != null);
             onExecutionComplete();
         }
     }
diff --git a/services/core/java/com/android/server/wm/AppCompatConfigurationPersister.java b/services/core/java/com/android/server/wm/AppCompatConfigurationPersister.java
index 852ce04..9c861fe 100644
--- a/services/core/java/com/android/server/wm/AppCompatConfigurationPersister.java
+++ b/services/core/java/com/android/server/wm/AppCompatConfigurationPersister.java
@@ -16,16 +16,13 @@
 
 package com.android.server.wm;
 
-import static android.os.StrictMode.setThreadPolicy;
-
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
+import android.annotation.MainThread;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.os.Environment;
-import android.os.StrictMode;
-import android.os.StrictMode.ThreadPolicy;
 import android.util.AtomicFile;
 import android.util.Slog;
 
@@ -122,7 +119,7 @@
         final File prefFiles = new File(configFolder, letterboxConfigurationFileName);
         mConfigurationFile = new AtomicFile(prefFiles);
         mPersisterQueue = persisterQueue;
-        runWithDiskReadsThreadPolicy(this::readCurrentConfiguration);
+        readCurrentConfiguration();
     }
 
     /**
@@ -212,6 +209,7 @@
                 mDefaultTabletopModeReachabilitySupplier.get();
     }
 
+    @MainThread
     private void readCurrentConfiguration() {
         if (!mConfigurationFile.exists()) {
             useDefaultValue();
@@ -272,20 +270,6 @@
         }
     }
 
-    // The LetterboxConfigurationDeviceConfig needs to access the
-    // file with the current reachability position once when the
-    // device boots. Because DisplayThread uses allowIo=false
-    // accessing a file triggers a DiskReadViolation.
-    // Here we use StrictMode to allow the current thread to read
-    // the AtomicFile once in the current thread restoring the
-    // original ThreadPolicy after that.
-    private void runWithDiskReadsThreadPolicy(Runnable runnable) {
-        final ThreadPolicy currentPolicy = StrictMode.getThreadPolicy();
-        setThreadPolicy(new ThreadPolicy.Builder().permitDiskReads().build());
-        runnable.run();
-        setThreadPolicy(currentPolicy);
-    }
-
     private static class UpdateValuesCommand implements
             PersisterQueue.WriteQueueItem<UpdateValuesCommand> {
 
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 459a509..33f2dd1 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -1186,9 +1186,13 @@
     public static WindowManagerService main(final Context context, final InputManagerService im,
             final boolean showBootMsgs, WindowManagerPolicy policy,
             ActivityTaskManagerService atm) {
+        // Using SysUI context to have access to Material colors extracted from Wallpaper.
+        final AppCompatConfiguration appCompat = new AppCompatConfiguration(
+                ActivityThread.currentActivityThread().getSystemUiContext());
+
         final WindowManagerService wms = main(context, im, showBootMsgs, policy, atm,
                 new DisplayWindowSettingsProvider(), SurfaceControl.Transaction::new,
-                SurfaceControl.Builder::new);
+                SurfaceControl.Builder::new, appCompat);
         WindowManagerGlobal.setWindowManagerServiceForSystemProcess(wms);
         return wms;
     }
@@ -1202,12 +1206,14 @@
             final boolean showBootMsgs, WindowManagerPolicy policy, ActivityTaskManagerService atm,
             DisplayWindowSettingsProvider displayWindowSettingsProvider,
             Supplier<SurfaceControl.Transaction> transactionFactory,
-            Supplier<SurfaceControl.Builder> surfaceControlFactory) {
+            Supplier<SurfaceControl.Builder> surfaceControlFactory,
+            AppCompatConfiguration appCompat) {
+
         final WindowManagerService[] wms = new WindowManagerService[1];
         DisplayThread.getHandler().runWithScissors(() ->
                 wms[0] = new WindowManagerService(context, im, showBootMsgs, policy, atm,
                         displayWindowSettingsProvider, transactionFactory,
-                        surfaceControlFactory), 0);
+                        surfaceControlFactory, appCompat), 0);
         return wms[0];
     }
 
@@ -1231,7 +1237,8 @@
             boolean showBootMsgs, WindowManagerPolicy policy, ActivityTaskManagerService atm,
             DisplayWindowSettingsProvider displayWindowSettingsProvider,
             Supplier<SurfaceControl.Transaction> transactionFactory,
-            Supplier<SurfaceControl.Builder> surfaceControlFactory) {
+            Supplier<SurfaceControl.Builder> surfaceControlFactory,
+            AppCompatConfiguration appCompat) {
         installLock(this, INDEX_WINDOW);
         mGlobalLock = atm.getGlobalLock();
         mAtmService = atm;
@@ -1283,9 +1290,7 @@
                     | WindowInsets.Type.navigationBars();
         }
 
-        mAppCompatConfiguration = new AppCompatConfiguration(
-                // Using SysUI context to have access to Material colors extracted from Wallpaper.
-                ActivityThread.currentActivityThread().getSystemUiContext());
+        mAppCompatConfiguration = appCompat;
 
         mInputManager = inputManager; // Must be before createDisplayContentLocked.
         mDisplayManagerInternal = LocalServices.getService(DisplayManagerInternal.class);
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index b527630..1640ad3 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2754,10 +2754,16 @@
      * @param outRegion The region to update.
      */
     private void updateRegionForModalActivityWindow(Region outRegion) {
-        // If the inner bounds of letterbox is available, then it will be used as the
-        // touchable region so it won't cover the touchable letterbox and the touch
-        // events can slip to activity from letterbox.
-        mActivityRecord.getLetterboxInnerBounds(mTmpRect);
+        if (Flags.scrollingFromLetterbox()) {
+            // Touchable region expands to the letterbox area to react to scrolls from letterbox.
+            mTmpRect.setEmpty();
+        } else {
+            // If the activity is letterboxed and scrolling from letterbox is disabled, limit the
+            // touchable region to the activity. This way, the letterbox area is exposed to react
+            // to touch events, and the touch events can slip from the activity from letterbox.
+            mActivityRecord.getLetterboxInnerBounds(mTmpRect);
+        }
+
         if (mTmpRect.isEmpty()) {
             final Rect transformedBounds = mActivityRecord.getFixedRotationTransformDisplayBounds();
             if (transformedBounds != null) {
diff --git a/services/tests/appfunctions/Android.bp b/services/tests/appfunctions/Android.bp
index b5cf986..9560ec9 100644
--- a/services/tests/appfunctions/Android.bp
+++ b/services/tests/appfunctions/Android.bp
@@ -45,8 +45,8 @@
     ],
 
     libs: [
-        "android.test.base",
-        "android.test.runner",
+        "android.test.base.stubs.system",
+        "android.test.runner.stubs.system",
     ],
 
     certificate: "platform",
diff --git a/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt b/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
index b938c3c..6930b3c 100644
--- a/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
+++ b/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
@@ -16,26 +16,28 @@
 package com.android.server.appfunctions
 
 import android.app.appfunctions.AppFunctionRuntimeMetadata
-import android.app.appfunctions.AppFunctionRuntimeMetadata.createParentAppFunctionRuntimeSchema
-import android.app.appfunctions.AppFunctionStaticMetadataHelper
+import android.app.appsearch.AppSearchBatchResult
 import android.app.appsearch.AppSearchManager
-import android.app.appsearch.AppSearchManager.SearchContext
+import android.app.appsearch.AppSearchResult
+import android.app.appsearch.AppSearchSchema
 import android.app.appsearch.GenericDocument
+import android.app.appsearch.GetByDocumentIdRequest
+import android.app.appsearch.GetSchemaResponse
 import android.app.appsearch.PutDocumentsRequest
+import android.app.appsearch.RemoveByDocumentIdRequest
 import android.app.appsearch.SearchResult
 import android.app.appsearch.SearchSpec
 import android.app.appsearch.SetSchemaRequest
+import android.app.appsearch.SetSchemaResponse
 import android.util.ArrayMap
 import android.util.ArraySet
+import android.util.Log
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.internal.infra.AndroidFuture
 import com.android.server.appfunctions.FutureAppSearchSession.FutureSearchResults
 import com.google.common.truth.Truth.assertThat
 import com.google.common.util.concurrent.MoreExecutors
-import java.util.concurrent.Executor
 import java.util.concurrent.atomic.AtomicBoolean
-import org.junit.After
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
@@ -47,46 +49,19 @@
     private val testExecutor = MoreExecutors.directExecutor()
     private val packageManager = context.packageManager
 
-    @Before
-    @After
-    fun clearData() {
-        val searchContext = SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
-            val setSchemaRequest = SetSchemaRequest.Builder().setForceOverride(true).build()
-            it.setSchema(setSchemaRequest).get()
-        }
-    }
-
     @Test
     fun getPackageToFunctionIdMap() {
-        val searchContext: SearchContext = SearchContext.Builder(TEST_DB).build()
+        val searchSession = FakeSearchSession()
         val functionRuntimeMetadata =
             AppFunctionRuntimeMetadata.Builder(TEST_TARGET_PKG_NAME, "testFunctionId", "").build()
-        val setSchemaRequest =
-            SetSchemaRequest.Builder()
-                .addSchemas(AppFunctionRuntimeMetadata.createParentAppFunctionRuntimeSchema())
-                .addSchemas(
-                    AppFunctionRuntimeMetadata.createAppFunctionRuntimeSchema(TEST_TARGET_PKG_NAME)
-                )
-                .build()
         val putDocumentsRequest: PutDocumentsRequest =
             PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
-        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
-            val setSchemaResponse = it.setSchema(setSchemaRequest).get()
-            assertThat(setSchemaResponse).isNotNull()
-            val appSearchBatchResult = it.put(putDocumentsRequest).get()
-            assertThat(appSearchBatchResult.isSuccess).isTrue()
-        }
+        searchSession.put(putDocumentsRequest).get()
 
-        val metadataSyncAdapter =
-            MetadataSyncAdapter(
-                testExecutor,
-                FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext),
-                packageManager,
-            )
         val packageToFunctionIdMap =
-            metadataSyncAdapter.getPackageToFunctionIdMap(
-                AppFunctionRuntimeMetadata.RUNTIME_SCHEMA_TYPE,
+            MetadataSyncAdapter.getPackageToFunctionIdMap(
+                searchSession,
+                "fakeSchema",
                 AppFunctionRuntimeMetadata.PROPERTY_FUNCTION_ID,
                 AppFunctionRuntimeMetadata.PROPERTY_PACKAGE_NAME,
             )
@@ -97,7 +72,7 @@
 
     @Test
     fun getPackageToFunctionIdMap_multipleDocuments() {
-        val searchContext: SearchContext = SearchContext.Builder(TEST_DB).build()
+        val searchSession = FakeSearchSession()
         val functionRuntimeMetadata =
             AppFunctionRuntimeMetadata.Builder(TEST_TARGET_PKG_NAME, "testFunctionId", "").build()
         val functionRuntimeMetadata1 =
@@ -106,13 +81,6 @@
             AppFunctionRuntimeMetadata.Builder(TEST_TARGET_PKG_NAME, "testFunctionId2", "").build()
         val functionRuntimeMetadata3 =
             AppFunctionRuntimeMetadata.Builder(TEST_TARGET_PKG_NAME, "testFunctionId3", "").build()
-        val setSchemaRequest =
-            SetSchemaRequest.Builder()
-                .addSchemas(AppFunctionRuntimeMetadata.createParentAppFunctionRuntimeSchema())
-                .addSchemas(
-                    AppFunctionRuntimeMetadata.createAppFunctionRuntimeSchema(TEST_TARGET_PKG_NAME)
-                )
-                .build()
         val putDocumentsRequest: PutDocumentsRequest =
             PutDocumentsRequest.Builder()
                 .addGenericDocuments(
@@ -122,21 +90,11 @@
                     functionRuntimeMetadata3,
                 )
                 .build()
-        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
-            val setSchemaResponse = it.setSchema(setSchemaRequest).get()
-            assertThat(setSchemaResponse).isNotNull()
-            val appSearchBatchResult = it.put(putDocumentsRequest).get()
-            assertThat(appSearchBatchResult.isSuccess).isTrue()
-        }
+        searchSession.put(putDocumentsRequest).get()
 
-        val metadataSyncAdapter =
-            MetadataSyncAdapter(
-                testExecutor,
-                FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext),
-                packageManager,
-            )
         val packageToFunctionIdMap =
-            metadataSyncAdapter.getPackageToFunctionIdMap(
+            MetadataSyncAdapter.getPackageToFunctionIdMap(
+                searchSession,
                 AppFunctionRuntimeMetadata.RUNTIME_SCHEMA_TYPE,
                 AppFunctionRuntimeMetadata.PROPERTY_FUNCTION_ID,
                 AppFunctionRuntimeMetadata.PROPERTY_PACKAGE_NAME,
@@ -172,62 +130,21 @@
 
     @Test
     fun syncMetadata_noDiff() {
-        val searchContext: SearchContext = SearchContext.Builder(TEST_DB).build()
-        val appSearchSession =
-            PartialFakeFutureAppSearchSession(appSearchManager, testExecutor, searchContext)
-        val fakeFunctionId = "syncMetadata_noDiff"
-        val fakeStaticMetadata: GenericDocument =
-            GenericDocument.Builder<GenericDocument.Builder<*>>(
-                    AppFunctionStaticMetadataHelper.APP_FUNCTION_STATIC_NAMESPACE,
-                    AppFunctionStaticMetadataHelper.getDocumentIdForAppFunction(
-                        TEST_TARGET_PKG_NAME,
-                        fakeFunctionId,
-                    ),
-                    AppFunctionStaticMetadataHelper.STATIC_SCHEMA_TYPE,
-                )
-                .setPropertyString(
-                    AppFunctionStaticMetadataHelper.PROPERTY_PACKAGE_NAME,
-                    TEST_TARGET_PKG_NAME,
-                )
-                .setPropertyString(
-                    AppFunctionStaticMetadataHelper.PROPERTY_FUNCTION_ID,
-                    fakeFunctionId,
-                )
-                .build()
-        appSearchSession.overrideStaticMetadataSearchResult = mutableListOf(fakeStaticMetadata)
-        val putCorrespondingSchema =
-            appSearchSession
-                .setSchema(
-                    SetSchemaRequest.Builder()
-                        .addSchemas(
-                            createParentAppFunctionRuntimeSchema(),
-                            AppFunctionRuntimeMetadata.createAppFunctionRuntimeSchema(
-                                TEST_TARGET_PKG_NAME
-                            ),
-                        )
-                        .setForceOverride(true)
-                        .build()
-                )
-                .get()
-        assertThat(putCorrespondingSchema).isNotNull()
-        val putCorrespondingRuntimeMetadata =
-            appSearchSession
-                .put(
-                    PutDocumentsRequest.Builder()
-                        .addGenericDocuments(
-                            AppFunctionRuntimeMetadata.Builder(
-                                    TEST_TARGET_PKG_NAME,
-                                    fakeFunctionId,
-                                    "",
-                                )
-                                .build()
-                        )
-                        .build()
-                )
-                .get()
-        assertThat(putCorrespondingRuntimeMetadata.isSuccess).isTrue()
+        val runtimeSearchSession = FakeSearchSession()
+        val staticSearchSession = FakeSearchSession()
+        val functionRuntimeMetadata =
+            AppFunctionRuntimeMetadata.Builder(TEST_TARGET_PKG_NAME, "testFunctionId", "").build()
+        val putDocumentsRequest: PutDocumentsRequest =
+            PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
+        runtimeSearchSession.put(putDocumentsRequest).get()
+        staticSearchSession.put(putDocumentsRequest).get()
         val metadataSyncAdapter =
-            MetadataSyncAdapter(testExecutor, appSearchSession, context.packageManager)
+            MetadataSyncAdapter(
+                testExecutor,
+                runtimeSearchSession,
+                staticSearchSession,
+                packageManager,
+            )
 
         val submitSyncRequest = metadataSyncAdapter.submitSyncRequest()
 
@@ -257,31 +174,20 @@
 
     @Test
     fun syncMetadata_addedFunction() {
-        val searchContext: SearchContext = SearchContext.Builder(TEST_DB).build()
-        val appSearchSession =
-            PartialFakeFutureAppSearchSession(appSearchManager, testExecutor, searchContext)
-        val fakeFunctionId = "addedFunction1"
-        val fakeStaticMetadata: GenericDocument =
-            GenericDocument.Builder<GenericDocument.Builder<*>>(
-                    AppFunctionStaticMetadataHelper.APP_FUNCTION_STATIC_NAMESPACE,
-                    AppFunctionStaticMetadataHelper.getDocumentIdForAppFunction(
-                        TEST_TARGET_PKG_NAME,
-                        fakeFunctionId,
-                    ),
-                    AppFunctionStaticMetadataHelper.STATIC_SCHEMA_TYPE,
-                )
-                .setPropertyString(
-                    AppFunctionStaticMetadataHelper.PROPERTY_PACKAGE_NAME,
-                    TEST_TARGET_PKG_NAME,
-                )
-                .setPropertyString(
-                    AppFunctionStaticMetadataHelper.PROPERTY_FUNCTION_ID,
-                    fakeFunctionId,
-                )
-                .build()
-        appSearchSession.overrideStaticMetadataSearchResult = mutableListOf(fakeStaticMetadata)
+        val runtimeSearchSession = FakeSearchSession()
+        val staticSearchSession = FakeSearchSession()
+        val functionRuntimeMetadata =
+            AppFunctionRuntimeMetadata.Builder(TEST_TARGET_PKG_NAME, "testFunctionId", "").build()
+        val putDocumentsRequest: PutDocumentsRequest =
+            PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
+        staticSearchSession.put(putDocumentsRequest).get()
         val metadataSyncAdapter =
-            MetadataSyncAdapter(testExecutor, appSearchSession, context.packageManager)
+            MetadataSyncAdapter(
+                testExecutor,
+                runtimeSearchSession,
+                staticSearchSession,
+                packageManager,
+            )
 
         val submitSyncRequest = metadataSyncAdapter.submitSyncRequest()
 
@@ -325,43 +231,20 @@
 
     @Test
     fun syncMetadata_removedFunction() {
-        val searchContext: SearchContext = SearchContext.Builder(TEST_DB).build()
-        val appSearchSession =
-            PartialFakeFutureAppSearchSession(appSearchManager, testExecutor, searchContext)
-        val fakeFunctionId = "syncMetadata_removedFunction"
-        val putCorrespondingSchema =
-            appSearchSession
-                .setSchema(
-                    SetSchemaRequest.Builder()
-                        .addSchemas(
-                            createParentAppFunctionRuntimeSchema(),
-                            AppFunctionRuntimeMetadata.createAppFunctionRuntimeSchema(
-                                TEST_TARGET_PKG_NAME
-                            ),
-                        )
-                        .setForceOverride(true)
-                        .build()
-                )
-                .get()
-        assertThat(putCorrespondingSchema).isNotNull()
-        val putStaleRuntimeMetadata =
-            appSearchSession
-                .put(
-                    PutDocumentsRequest.Builder()
-                        .addGenericDocuments(
-                            AppFunctionRuntimeMetadata.Builder(
-                                    TEST_TARGET_PKG_NAME,
-                                    fakeFunctionId,
-                                    "",
-                                )
-                                .build()
-                        )
-                        .build()
-                )
-                .get()
-        assertThat(putStaleRuntimeMetadata.isSuccess).isTrue()
+        val runtimeSearchSession = FakeSearchSession()
+        val staticSearchSession = FakeSearchSession()
+        val functionRuntimeMetadata =
+            AppFunctionRuntimeMetadata.Builder(TEST_TARGET_PKG_NAME, "testFunctionId", "").build()
+        val putDocumentsRequest: PutDocumentsRequest =
+            PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
+        runtimeSearchSession.put(putDocumentsRequest).get()
         val metadataSyncAdapter =
-            MetadataSyncAdapter(testExecutor, appSearchSession, context.packageManager)
+            MetadataSyncAdapter(
+                testExecutor,
+                runtimeSearchSession,
+                staticSearchSession,
+                packageManager,
+            )
 
         val submitSyncRequest = metadataSyncAdapter.submitSyncRequest()
 
@@ -426,48 +309,99 @@
         const val TEST_TARGET_PKG_NAME = "com.android.frameworks.appfunctionstests"
     }
 
-    class PartialFakeFutureAppSearchSession(
-        appSearchManager: AppSearchManager,
-        executor: Executor,
-        appSearchContext: SearchContext,
-    ) : FutureAppSearchSessionImpl(appSearchManager, executor, appSearchContext) {
-        var overrideStaticMetadataSearchResult: MutableList<GenericDocument> = mutableListOf()
-        private val overrideUsed = AtomicBoolean(false)
+    class FakeSearchSession : FutureAppSearchSession {
+        private val schemas: MutableSet<AppSearchSchema> = mutableSetOf()
+        private val genericDocumentMutableMap: MutableMap<String, GenericDocument> = mutableMapOf()
 
-        // Overriding this method to fake searching for static metadata.
-        // Static metadata is the source of truth for the metadata sync behaviour since the sync is
-        // updating the runtime metadata to match the existing static metadata.
+        override fun close() {
+            Log.d("FakeRuntimeMetadataSearchSession", "Closing session")
+        }
+
+        override fun setSchema(
+            setSchemaRequest: SetSchemaRequest
+        ): AndroidFuture<SetSchemaResponse> {
+            schemas.addAll(setSchemaRequest.schemas)
+            return AndroidFuture.completedFuture(SetSchemaResponse.Builder().build())
+        }
+
+        override fun getSchema(): AndroidFuture<GetSchemaResponse> {
+            val resultBuilder = GetSchemaResponse.Builder()
+            for (schema in schemas) {
+                resultBuilder.addSchema(schema)
+            }
+            return AndroidFuture.completedFuture(resultBuilder.build())
+        }
+
+        override fun put(
+            putDocumentsRequest: PutDocumentsRequest
+        ): AndroidFuture<AppSearchBatchResult<String, Void>> {
+            for (document in putDocumentsRequest.genericDocuments) {
+                genericDocumentMutableMap[document.id] = document
+            }
+            val batchResultBuilder = AppSearchBatchResult.Builder<String, Void>()
+            for (document in putDocumentsRequest.genericDocuments) {
+                batchResultBuilder.setResult(document.id, AppSearchResult.newSuccessfulResult(null))
+            }
+            return AndroidFuture.completedFuture(batchResultBuilder.build())
+        }
+
+        override fun remove(
+            removeRequest: RemoveByDocumentIdRequest
+        ): AndroidFuture<AppSearchBatchResult<String, Void>> {
+            for (documentId in removeRequest.ids) {
+                if (!genericDocumentMutableMap.keys.contains(documentId)) {
+                    throw IllegalStateException("Document $documentId does not exist")
+                }
+            }
+            val batchResultBuilder = AppSearchBatchResult.Builder<String, Void>()
+            for (id in removeRequest.ids) {
+                batchResultBuilder.setResult(id, AppSearchResult.newSuccessfulResult(null))
+            }
+            return AndroidFuture.completedFuture(batchResultBuilder.build())
+        }
+
+        override fun getByDocumentId(
+            getRequest: GetByDocumentIdRequest
+        ): AndroidFuture<AppSearchBatchResult<String, GenericDocument>> {
+            val batchResultBuilder = AppSearchBatchResult.Builder<String, GenericDocument>()
+            for (documentId in getRequest.ids) {
+                if (!genericDocumentMutableMap.keys.contains(documentId)) {
+                    throw IllegalStateException("Document $documentId does not exist")
+                }
+                batchResultBuilder.setResult(
+                    documentId,
+                    AppSearchResult.newSuccessfulResult(genericDocumentMutableMap[documentId]),
+                )
+            }
+            return AndroidFuture.completedFuture(batchResultBuilder.build())
+        }
+
         override fun search(
             queryExpression: String,
             searchSpec: SearchSpec,
         ): AndroidFuture<FutureSearchResults> {
-            if (
-                searchSpec.filterSchemas.contains(
-                    AppFunctionStaticMetadataHelper.STATIC_SCHEMA_TYPE
-                )
-            ) {
-                val futureSearchResults =
-                    object : FutureSearchResults {
-                        override fun getNextPage(): AndroidFuture<MutableList<SearchResult>> {
-                            if (overrideUsed.get()) {
-                                overrideStaticMetadataSearchResult.clear()
-                                return AndroidFuture.completedFuture(mutableListOf())
-                            }
-                            overrideUsed.set(true)
-                            return AndroidFuture.completedFuture(
-                                overrideStaticMetadataSearchResult
-                                    .map {
-                                        SearchResult.Builder(TEST_TARGET_PKG_NAME, TEST_DB)
-                                            .setGenericDocument(it)
-                                            .build()
-                                    }
-                                    .toMutableList()
-                            )
+            val futureSearchResults =
+                object : FutureSearchResults {
+                    val hasNextPage = AtomicBoolean(false)
+
+                    override fun getNextPage(): AndroidFuture<MutableList<SearchResult>> {
+                        val searchResultMutableList: MutableList<SearchResult> =
+                            genericDocumentMutableMap.values
+                                .map {
+                                    SearchResult.Builder(TEST_TARGET_PKG_NAME, TEST_DB)
+                                        .setGenericDocument(it)
+                                        .build()
+                                }
+                                .toMutableList()
+                        if (!hasNextPage.get()) {
+                            hasNextPage.set(true)
+                            return AndroidFuture.completedFuture(searchResultMutableList)
+                        } else {
+                            return AndroidFuture.completedFuture(mutableListOf())
                         }
                     }
-                return AndroidFuture.completedFuture(futureSearchResults)
-            }
-            return super.search(queryExpression, searchSpec)
+                }
+            return AndroidFuture.completedFuture(futureSearchResults)
         }
     }
 }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
index d0aec3b..bf5a692 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
@@ -75,6 +75,7 @@
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import com.android.internal.app.IBatteryStats;
 import com.android.modules.utils.testing.ExtendedMockitoRule;
 import com.android.server.LocalServices;
 import com.android.server.am.BatteryStatsService;
@@ -158,6 +159,8 @@
     private DisplayManagerFlags mDisplayManagerFlagsMock;
     @Mock
     private DisplayManagerInternal.DisplayOffloadSession mDisplayOffloadSession;
+    @Mock
+    private IBatteryStats mMockBatteryStats;
     @Captor
     private ArgumentCaptor<SensorEventListener> mSensorEventListenerCaptor;
 
@@ -204,7 +207,8 @@
 
         doAnswer((Answer<Void>) invocationOnMock -> null).when(() ->
                 SystemProperties.set(anyString(), any()));
-        doAnswer((Answer<Void>) invocationOnMock -> null).when(BatteryStatsService::getService);
+        doAnswer((Answer<IBatteryStats>) invocationOnMock -> mMockBatteryStats)
+                .when(BatteryStatsService::getService);
         doAnswer((Answer<Boolean>) invocationOnMock -> false)
                 .when(ActivityManager::isLowRamDeviceStatic);
 
@@ -2227,6 +2231,52 @@
         verify(mHolder.brightnessSetting).saveIfNeeded();
     }
 
+    @Test
+    public void testBatteryStatNotes_enabledOnDefaultDisplayWhenDisabledOnOthers()
+            throws Exception {
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(false);
+
+        verifyNoteScreenState(Display.DEFAULT_DISPLAY, /* expectNote= */ true);
+    }
+
+    @Test
+    public void testBatteryStatNotes_enabledOnDefaultDisplayWhenEnabledOnOthers() throws Exception {
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(true);
+
+        verifyNoteScreenState(Display.DEFAULT_DISPLAY, /* expectNote= */ true);
+    }
+
+    @Test
+    public void testBatteryStatNotes_flagGuardedOnNonDefaultDisplays() throws Exception {
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(false);
+
+        verifyNoteScreenState(/* displayId= */ 2, /* expectNote= */ false);
+
+        when(mDisplayManagerFlagsMock.isBatteryStatsEnabledForAllDisplays()).thenReturn(true);
+
+        verifyNoteScreenState(/* displayId= */ 2, /* expectNote= */ true);
+    }
+
+    private void verifyNoteScreenState(int displayId, boolean expectNote) throws Exception {
+        mHolder = createDisplayPowerController(displayId, UNIQUE_ID);
+        DisplayPowerRequest dpr = new DisplayPowerRequest();
+        dpr.policy = DisplayPowerRequest.POLICY_BRIGHT;
+        when(mHolder.displayPowerState.getScreenState()).thenReturn(Display.STATE_ON);
+
+        mHolder.dpc.requestPowerState(dpr, /* waitForNegativeProximity= */ false);
+        advanceTime(1); // Run updatePowerState
+
+        if (expectNote) {
+            verify(mMockBatteryStats)
+                    .noteScreenState(
+                            displayId, Display.STATE_ON, Display.STATE_REASON_DEFAULT_POLICY);
+            verify(mMockBatteryStats).noteScreenBrightness(eq(displayId), anyInt());
+        } else {
+            verify(mMockBatteryStats, never()).noteScreenState(anyInt(), anyInt(), anyInt());
+            verify(mMockBatteryStats, never()).noteScreenBrightness(anyInt(), anyInt());
+        }
+    }
+
     /**
      * Creates a mock and registers it to {@link LocalServices}.
      */
diff --git a/services/tests/mockingservicestests/assets/CpuInfoReaderTest/valid_cpuset_2/background/cpus b/services/tests/mockingservicestests/assets/CpuInfoReaderTest/valid_cpuset_2/background/cpus
new file mode 100644
index 0000000..8b0fab8
--- /dev/null
+++ b/services/tests/mockingservicestests/assets/CpuInfoReaderTest/valid_cpuset_2/background/cpus
@@ -0,0 +1 @@
+0-1
diff --git a/services/tests/mockingservicestests/assets/CpuInfoReaderTest/valid_cpuset_2/top-app/cpus b/services/tests/mockingservicestests/assets/CpuInfoReaderTest/valid_cpuset_2/top-app/cpus
new file mode 100644
index 0000000..40c7bb2
--- /dev/null
+++ b/services/tests/mockingservicestests/assets/CpuInfoReaderTest/valid_cpuset_2/top-app/cpus
@@ -0,0 +1 @@
+0-3
diff --git a/services/tests/mockingservicestests/src/com/android/server/cpu/CpuInfoReaderTest.java b/services/tests/mockingservicestests/src/com/android/server/cpu/CpuInfoReaderTest.java
index 2fbe8aa..3fe038a 100644
--- a/services/tests/mockingservicestests/src/com/android/server/cpu/CpuInfoReaderTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/cpu/CpuInfoReaderTest.java
@@ -26,6 +26,7 @@
 
 import android.content.Context;
 import android.content.res.AssetManager;
+import android.util.IntArray;
 import android.util.Log;
 import android.util.SparseArray;
 
@@ -48,6 +49,7 @@
     private static final String TAG = CpuInfoReaderTest.class.getSimpleName();
     private static final String ROOT_DIR_NAME = "CpuInfoReaderTest";
     private static final String VALID_CPUSET_DIR = "valid_cpuset";
+    private static final String VALID_CPUSET_2_DIR = "valid_cpuset_2";
     private static final String VALID_CPUSET_WITH_EMPTY_CPUS = "valid_cpuset_with_empty_cpus";
     private static final String VALID_CPUFREQ_WITH_EMPTY_AFFECTED_CPUS =
             "valid_cpufreq_with_empty_affected_cpus";
@@ -88,54 +90,95 @@
     }
 
     @Test
+    public void testReadCpuInfoWithUpdatedCpuset() throws Exception {
+        CpuInfoReader cpuInfoReader = newCpuInfoReader(getCacheFile(VALID_CPUSET_DIR),
+                getCacheFile(VALID_CPUFREQ_WITH_TIME_IN_STATE_DIR), getCacheFile(VALID_PROC_STAT));
+
+        SparseArray<CpuInfoReader.CpuInfo> actualCpuInfos = cpuInfoReader.readCpuInfos();
+        SparseArray<CpuInfoReader.CpuInfo> expectedCpuInfos =
+                getFirstCpuInfosWithTimeInStateSnapshot();
+
+        compareCpuInfos("CPU infos first snapshot", expectedCpuInfos, actualCpuInfos);
+
+        cpuInfoReader.setCpusetDir(getCacheFile(VALID_CPUSET_2_DIR));
+        cpuInfoReader.setCpuFreqDir(getCacheFile(VALID_CPUFREQ_WITH_TIME_IN_STATE_2_DIR));
+        cpuInfoReader.setProcStatFile(getCacheFile(VALID_PROC_STAT_2));
+
+        actualCpuInfos = cpuInfoReader.readCpuInfos();
+
+        IntArray cpusetCategories = new IntArray();
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP);
+        expectedCpuInfos = getSecondCpuInfosWithTimeInStateSnapshot(cpusetCategories);
+
+        compareCpuInfos("CPU infos second snapshot", expectedCpuInfos, actualCpuInfos);
+    }
+
+    @Test
+    public void testReadCpuInfoWithUpdatedCpusetBeforeStopSignal() throws Exception {
+        CpuInfoReader cpuInfoReader = newCpuInfoReader(getCacheFile(VALID_CPUSET_DIR),
+                getCacheFile(VALID_CPUFREQ_WITH_TIME_IN_STATE_DIR), getCacheFile(VALID_PROC_STAT));
+
+        SparseArray<CpuInfoReader.CpuInfo> actualCpuInfos = cpuInfoReader.readCpuInfos();
+        SparseArray<CpuInfoReader.CpuInfo> expectedCpuInfos =
+                getFirstCpuInfosWithTimeInStateSnapshot();
+
+        compareCpuInfos("CPU infos first snapshot", expectedCpuInfos, actualCpuInfos);
+
+        cpuInfoReader.setCpusetDir(getCacheFile(VALID_CPUSET_2_DIR));
+        // When stopping the periodic cpuset reading, the reader will create a new snapshot.
+        cpuInfoReader.stopPeriodicCpusetReading();
+        // Any cpuset update after the stop signal should be ignored by the reader.
+        cpuInfoReader.setCpusetDir(getCacheFile(VALID_CPUSET_DIR));
+        cpuInfoReader.setCpuFreqDir(getCacheFile(VALID_CPUFREQ_WITH_TIME_IN_STATE_2_DIR));
+        cpuInfoReader.setProcStatFile(getCacheFile(VALID_PROC_STAT_2));
+
+        actualCpuInfos = cpuInfoReader.readCpuInfos();
+
+        IntArray cpusetCategories = new IntArray();
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP);
+        expectedCpuInfos = getSecondCpuInfosWithTimeInStateSnapshot(cpusetCategories);
+
+        compareCpuInfos("CPU infos second snapshot", expectedCpuInfos, actualCpuInfos);
+    }
+
+
+    @Test
+    public void testReadCpuInfoWithUpdatedCpusetAfterStopSignal() throws Exception {
+        CpuInfoReader cpuInfoReader = newCpuInfoReader(getCacheFile(VALID_CPUSET_DIR),
+                getCacheFile(VALID_CPUFREQ_WITH_TIME_IN_STATE_DIR), getCacheFile(VALID_PROC_STAT));
+
+        SparseArray<CpuInfoReader.CpuInfo> actualCpuInfos = cpuInfoReader.readCpuInfos();
+        SparseArray<CpuInfoReader.CpuInfo> expectedCpuInfos =
+                getFirstCpuInfosWithTimeInStateSnapshot();
+
+        compareCpuInfos("CPU infos first snapshot", expectedCpuInfos, actualCpuInfos);
+
+        cpuInfoReader.stopPeriodicCpusetReading();
+        // Any cpuset update after the stop signal should be ignored by the reader.
+        cpuInfoReader.setCpusetDir(getCacheFile(VALID_CPUSET_2_DIR));
+        cpuInfoReader.setCpuFreqDir(getCacheFile(VALID_CPUFREQ_WITH_TIME_IN_STATE_2_DIR));
+        cpuInfoReader.setProcStatFile(getCacheFile(VALID_PROC_STAT_2));
+
+        actualCpuInfos = cpuInfoReader.readCpuInfos();
+
+        expectedCpuInfos = getSecondCpuInfosWithTimeInStateSnapshot();
+        compareCpuInfos("CPU infos second snapshot", expectedCpuInfos, actualCpuInfos);
+    }
+
+    @Test
     public void testReadCpuInfoWithTimeInState() throws Exception {
         CpuInfoReader cpuInfoReader = newCpuInfoReader(getCacheFile(VALID_CPUSET_DIR),
                 getCacheFile(VALID_CPUFREQ_WITH_TIME_IN_STATE_DIR), getCacheFile(VALID_PROC_STAT));
 
         SparseArray<CpuInfoReader.CpuInfo> actualCpuInfos = cpuInfoReader.readCpuInfos();
-        SparseArray<CpuInfoReader.CpuInfo> expectedCpuInfos = new SparseArray<>();
-        expectedCpuInfos.append(0, new CpuInfoReader.CpuInfo(/* cpuCore= */ 0,
-                FLAG_CPUSET_CATEGORY_TOP_APP, /* isOnline= */ true, /* curCpuFreqKHz= */ 1_230_000,
-                /* maxCpuFreqKHz= */ 2_500_000, /* avgTimeInStateCpuFreqKHz= */ 488_095,
-                /* normalizedAvailableCpuFreqKHz= */ 2_402_267,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 32_249_610,
-                        /* niceTimeMillis= */ 7_950_930, /* systemTimeMillis= */ 52_227_050,
-                        /* idleTimeMillis= */ 409_036_950, /* iowaitTimeMillis= */ 1_322_810,
-                        /* irqTimeMillis= */ 8_146_740, /* softirqTimeMillis= */ 428_970,
-                        /* stealTimeMillis= */ 81_950, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
-        expectedCpuInfos.append(1, new CpuInfoReader.CpuInfo(/* cpuCore= */ 1,
-                FLAG_CPUSET_CATEGORY_TOP_APP, /* isOnline= */ true, /* curCpuFreqKHz= */ 1_450_000,
-                /* maxCpuFreqKHz= */ 2_800_000, /* avgTimeInStateCpuFreqKHz= */ 502_380,
-                /* normalizedAvailableCpuFreqKHz= */ 2_693_525,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 28_949_280,
-                        /* niceTimeMillis= */ 7_799_450, /* systemTimeMillis= */ 54_004_020,
-                        /* idleTimeMillis= */ 402_707_120, /* iowaitTimeMillis= */ 1_186_960,
-                        /* irqTimeMillis= */ 14_786_940, /* softirqTimeMillis= */ 1_498_130,
-                        /* stealTimeMillis= */ 78_780, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
-        expectedCpuInfos.append(2, new CpuInfoReader.CpuInfo(/* cpuCore= */ 2,
-                FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND,
-                /* isOnline= */ true, /* curCpuFreqKHz= */ 1_000_000,
-                /* maxCpuFreqKHz= */ 2_000_000, /* avgTimeInStateCpuFreqKHz= */ 464_285,
-                /* normalizedAvailableCpuFreqKHz= */ 1_901_608,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 28_959_280,
-                        /* niceTimeMillis= */ 7_789_450, /* systemTimeMillis= */ 54_014_020,
-                        /* idleTimeMillis= */ 402_717_120, /* iowaitTimeMillis= */ 1_166_960,
-                        /* irqTimeMillis= */ 14_796_940, /* softirqTimeMillis= */ 1_478_130,
-                        /* stealTimeMillis= */ 88_780, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
-        expectedCpuInfos.append(3, new CpuInfoReader.CpuInfo(/* cpuCore= */ 3,
-                FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND,
-                /* isOnline= */ true, /* curCpuFreqKHz= */ 1_000_000,
-                /* maxCpuFreqKHz= */ 2_000_000, /* avgTimeInStateCpuFreqKHz= */ 464_285,
-                /* normalizedAvailableCpuFreqKHz= */ 1_907_125,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 32_349_610,
-                        /* niceTimeMillis= */ 7_850_930, /* systemTimeMillis= */ 52_127_050,
-                        /* idleTimeMillis= */ 409_136_950, /* iowaitTimeMillis= */ 1_332_810,
-                        /* irqTimeMillis= */ 8_136_740, /* softirqTimeMillis= */ 438_970,
-                        /* stealTimeMillis= */ 71_950, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
+        SparseArray<CpuInfoReader.CpuInfo> expectedCpuInfos =
+                getFirstCpuInfosWithTimeInStateSnapshot();
 
         compareCpuInfos("CPU infos first snapshot", expectedCpuInfos, actualCpuInfos);
 
@@ -144,49 +187,7 @@
 
         actualCpuInfos = cpuInfoReader.readCpuInfos();
 
-        expectedCpuInfos.clear();
-        expectedCpuInfos.append(0, new CpuInfoReader.CpuInfo(/* cpuCore= */ 0,
-                FLAG_CPUSET_CATEGORY_TOP_APP, /* isOnline= */ true, /* curCpuFreqKHz= */ 1_000_000,
-                /* maxCpuFreqKHz= */ 2_600_000, /* avgTimeInStateCpuFreqKHz= */ 419_354,
-                /* normalizedAvailableCpuFreqKHz= */ 2_525_919,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 10_000_000,
-                        /* niceTimeMillis= */ 1_000_000, /* systemTimeMillis= */ 10_000_000,
-                        /* idleTimeMillis= */ 110_000_000, /* iowaitTimeMillis= */ 1_100_000,
-                        /* irqTimeMillis= */ 1_400_000, /* softirqTimeMillis= */ 80_000,
-                        /* stealTimeMillis= */ 21_000, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
-        expectedCpuInfos.append(1, new CpuInfoReader.CpuInfo(/* cpuCore= */ 1,
-                FLAG_CPUSET_CATEGORY_TOP_APP, /* isOnline= */ true, /* curCpuFreqKHz= */ 2_800_000,
-                /* maxCpuFreqKHz= */ 2_900_000, /* avgTimeInStateCpuFreqKHz= */ 429_032,
-                /* normalizedAvailableCpuFreqKHz= */ 2_503_009,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 900_000,
-                        /* niceTimeMillis= */ 1_000_000, /* systemTimeMillis= */ 10_000_000,
-                        /* idleTimeMillis= */ 1_000_000, /* iowaitTimeMillis= */ 90_000,
-                        /* irqTimeMillis= */ 200_000, /* softirqTimeMillis= */ 100_000,
-                        /* stealTimeMillis= */ 100_000, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
-        expectedCpuInfos.append(2, new CpuInfoReader.CpuInfo(/* cpuCore= */ 2,
-                FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND,
-                /* isOnline= */ true, /* curCpuFreqKHz= */ 2_000_000,
-                /* maxCpuFreqKHz= */ 2_100_000, /* avgTimeInStateCpuFreqKHz= */ 403_225,
-                /* normalizedAvailableCpuFreqKHz= */ 1_788_209,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 10_000_000,
-                        /* niceTimeMillis= */ 2_000_000, /* systemTimeMillis= */ 0,
-                        /* idleTimeMillis= */ 10_000_000, /* iowaitTimeMillis= */ 1_000_000,
-                        /* irqTimeMillis= */ 20_000_000, /* softirqTimeMillis= */ 1_000_000,
-                        /* stealTimeMillis= */ 100_000, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
-        expectedCpuInfos.append(3, new CpuInfoReader.CpuInfo(/* cpuCore= */ 3,
-                FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND,
-                /* isOnline= */ false, /* curCpuFreqKHz= */ MISSING_FREQUENCY,
-                /* maxCpuFreqKHz= */ 2_100_000, /* avgTimeInStateCpuFreqKHz= */ MISSING_FREQUENCY,
-                /* normalizedAvailableCpuFreqKHz= */ MISSING_FREQUENCY,
-                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 2_000_000,
-                        /* niceTimeMillis= */ 1_000_000, /* systemTimeMillis= */ 1_000_000,
-                        /* idleTimeMillis= */ 100_000, /* iowaitTimeMillis= */ 100_000,
-                        /* irqTimeMillis= */ 100_000, /* softirqTimeMillis= */ 1_000_000,
-                        /* stealTimeMillis= */ 1_000, /* guestTimeMillis= */ 0,
-                        /* guestNiceTimeMillis= */ 0)));
+        expectedCpuInfos = getSecondCpuInfosWithTimeInStateSnapshot();
 
         compareCpuInfos("CPU infos second snapshot", expectedCpuInfos, actualCpuInfos);
     }
@@ -592,4 +593,108 @@
         }
         return rootDir.delete();
     }
+
+    private SparseArray<CpuInfoReader.CpuInfo> getFirstCpuInfosWithTimeInStateSnapshot() {
+        SparseArray<CpuInfoReader.CpuInfo> cpuInfos = new SparseArray<>();
+        cpuInfos.append(0, new CpuInfoReader.CpuInfo(/* cpuCore= */ 0, FLAG_CPUSET_CATEGORY_TOP_APP,
+                /* isOnline= */ true, /* curCpuFreqKHz= */ 1_230_000,
+                /* maxCpuFreqKHz= */ 2_500_000, /* avgTimeInStateCpuFreqKHz= */ 488_095,
+                /* normalizedAvailableCpuFreqKHz= */ 2_402_267,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 32_249_610,
+                        /* niceTimeMillis= */ 7_950_930, /* systemTimeMillis= */ 52_227_050,
+                        /* idleTimeMillis= */ 409_036_950, /* iowaitTimeMillis= */ 1_322_810,
+                        /* irqTimeMillis= */ 8_146_740, /* softirqTimeMillis= */ 428_970,
+                        /* stealTimeMillis= */ 81_950, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        cpuInfos.append(1, new CpuInfoReader.CpuInfo(/* cpuCore= */ 1, FLAG_CPUSET_CATEGORY_TOP_APP,
+                /* isOnline= */ true, /* curCpuFreqKHz= */ 1_450_000,
+                /* maxCpuFreqKHz= */ 2_800_000, /* avgTimeInStateCpuFreqKHz= */ 502_380,
+                /* normalizedAvailableCpuFreqKHz= */ 2_693_525,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 28_949_280,
+                        /* niceTimeMillis= */ 7_799_450, /* systemTimeMillis= */ 54_004_020,
+                        /* idleTimeMillis= */ 402_707_120, /* iowaitTimeMillis= */ 1_186_960,
+                        /* irqTimeMillis= */ 14_786_940, /* softirqTimeMillis= */ 1_498_130,
+                        /* stealTimeMillis= */ 78_780, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        cpuInfos.append(2, new CpuInfoReader.CpuInfo(/* cpuCore= */ 2,
+                FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND,
+                /* isOnline= */ true, /* curCpuFreqKHz= */ 1_000_000,
+                /* maxCpuFreqKHz= */ 2_000_000, /* avgTimeInStateCpuFreqKHz= */ 464_285,
+                /* normalizedAvailableCpuFreqKHz= */ 1_901_608,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 28_959_280,
+                        /* niceTimeMillis= */ 7_789_450, /* systemTimeMillis= */ 54_014_020,
+                        /* idleTimeMillis= */ 402_717_120, /* iowaitTimeMillis= */ 1_166_960,
+                        /* irqTimeMillis= */ 14_796_940, /* softirqTimeMillis= */ 1_478_130,
+                        /* stealTimeMillis= */ 88_780, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        cpuInfos.append(3, new CpuInfoReader.CpuInfo(/* cpuCore= */ 3,
+                FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND,
+                /* isOnline= */ true, /* curCpuFreqKHz= */ 1_000_000,
+                /* maxCpuFreqKHz= */ 2_000_000, /* avgTimeInStateCpuFreqKHz= */ 464_285,
+                /* normalizedAvailableCpuFreqKHz= */ 1_907_125,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 32_349_610,
+                        /* niceTimeMillis= */ 7_850_930, /* systemTimeMillis= */ 52_127_050,
+                        /* idleTimeMillis= */ 409_136_950, /* iowaitTimeMillis= */ 1_332_810,
+                        /* irqTimeMillis= */ 8_136_740, /* softirqTimeMillis= */ 438_970,
+                        /* stealTimeMillis= */ 71_950, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        return cpuInfos;
+    }
+
+    private SparseArray<CpuInfoReader.CpuInfo> getSecondCpuInfosWithTimeInStateSnapshot() {
+        IntArray cpusetCategories = new IntArray();
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND);
+        cpusetCategories.add(FLAG_CPUSET_CATEGORY_TOP_APP | FLAG_CPUSET_CATEGORY_BACKGROUND);
+        return getSecondCpuInfosWithTimeInStateSnapshot(cpusetCategories);
+    }
+
+    private SparseArray<CpuInfoReader.CpuInfo> getSecondCpuInfosWithTimeInStateSnapshot(
+        IntArray cpusetCategories) {
+        SparseArray<CpuInfoReader.CpuInfo> cpuInfos = new SparseArray<>();
+        cpuInfos.append(0, new CpuInfoReader.CpuInfo(/* cpuCore= */ 0, cpusetCategories.get(0),
+                /* isOnline= */ true, /* curCpuFreqKHz= */ 1_000_000,
+                /* maxCpuFreqKHz= */ 2_600_000, /* avgTimeInStateCpuFreqKHz= */ 419_354,
+                /* normalizedAvailableCpuFreqKHz= */ 2_525_919,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 10_000_000,
+                        /* niceTimeMillis= */ 1_000_000, /* systemTimeMillis= */ 10_000_000,
+                        /* idleTimeMillis= */ 110_000_000, /* iowaitTimeMillis= */ 1_100_000,
+                        /* irqTimeMillis= */ 1_400_000, /* softirqTimeMillis= */ 80_000,
+                        /* stealTimeMillis= */ 21_000, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        cpuInfos.append(1, new CpuInfoReader.CpuInfo(/* cpuCore= */ 1, cpusetCategories.get(1),
+                /* isOnline= */ true, /* curCpuFreqKHz= */ 2_800_000,
+                /* maxCpuFreqKHz= */ 2_900_000, /* avgTimeInStateCpuFreqKHz= */ 429_032,
+                /* normalizedAvailableCpuFreqKHz= */ 2_503_009,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 900_000,
+                        /* niceTimeMillis= */ 1_000_000, /* systemTimeMillis= */ 10_000_000,
+                        /* idleTimeMillis= */ 1_000_000, /* iowaitTimeMillis= */ 90_000,
+                        /* irqTimeMillis= */ 200_000, /* softirqTimeMillis= */ 100_000,
+                        /* stealTimeMillis= */ 100_000, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        cpuInfos.append(2, new CpuInfoReader.CpuInfo(/* cpuCore= */ 2, cpusetCategories.get(2),
+                /* isOnline= */ true, /* curCpuFreqKHz= */ 2_000_000,
+                /* maxCpuFreqKHz= */ 2_100_000, /* avgTimeInStateCpuFreqKHz= */ 403_225,
+                /* normalizedAvailableCpuFreqKHz= */ 1_788_209,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 10_000_000,
+                        /* niceTimeMillis= */ 2_000_000, /* systemTimeMillis= */ 0,
+                        /* idleTimeMillis= */ 10_000_000, /* iowaitTimeMillis= */ 1_000_000,
+                        /* irqTimeMillis= */ 20_000_000, /* softirqTimeMillis= */ 1_000_000,
+                        /* stealTimeMillis= */ 100_000, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        cpuInfos.append(3, new CpuInfoReader.CpuInfo(/* cpuCore= */ 3, cpusetCategories.get(3),
+                /* isOnline= */ false,
+                /* curCpuFreqKHz= */ MISSING_FREQUENCY, /* maxCpuFreqKHz= */ 2_100_000,
+                /* avgTimeInStateCpuFreqKHz= */ MISSING_FREQUENCY,
+                /* normalizedAvailableCpuFreqKHz= */ MISSING_FREQUENCY,
+                new CpuInfoReader.CpuUsageStats(/* userTimeMillis= */ 2_000_000,
+                        /* niceTimeMillis= */ 1_000_000, /* systemTimeMillis= */ 1_000_000,
+                        /* idleTimeMillis= */ 100_000, /* iowaitTimeMillis= */ 100_000,
+                        /* irqTimeMillis= */ 100_000, /* softirqTimeMillis= */ 1_000_000,
+                        /* stealTimeMillis= */ 1_000, /* guestTimeMillis= */ 0,
+                        /* guestNiceTimeMillis= */ 0)));
+        return cpuInfos;
+    }
+
 }
diff --git a/services/tests/mockingservicestests/src/com/android/server/cpu/CpuMonitorServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/cpu/CpuMonitorServiceTest.java
index 994313f..d9e09d8 100644
--- a/services/tests/mockingservicestests/src/com/android/server/cpu/CpuMonitorServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/cpu/CpuMonitorServiceTest.java
@@ -26,6 +26,7 @@
 import static com.android.server.cpu.CpuInfoReader.FLAG_CPUSET_CATEGORY_BACKGROUND;
 import static com.android.server.cpu.CpuInfoReader.FLAG_CPUSET_CATEGORY_TOP_APP;
 import static com.android.server.cpu.CpuMonitorService.DEFAULT_MONITORING_INTERVAL_MILLISECONDS;
+import static com.android.server.SystemService.PHASE_BOOT_COMPLETED;
 
 import static com.google.common.truth.Truth.assertWithMessage;
 
@@ -75,6 +76,7 @@
     private static final long TEST_NORMAL_MONITORING_INTERVAL_MILLISECONDS = 100;
     private static final long TEST_DEBUG_MONITORING_INTERVAL_MILLISECONDS = 150;
     private static final long TEST_LATEST_AVAILABILITY_DURATION_MILLISECONDS = 300;
+    private static final long TEST_STOP_PERIODIC_CPUSET_READING_DELAY_MILLISECONDS = 0;
     private static final CpuAvailabilityMonitoringConfig TEST_MONITORING_CONFIG_ALL_CPUSET =
             new CpuAvailabilityMonitoringConfig.Builder(CPUSET_ALL)
                     .addThreshold(30).addThreshold(70).build();
@@ -119,7 +121,8 @@
         mService = new CpuMonitorService(mMockContext, mMockCpuInfoReader, mServiceHandlerThread,
                 /* shouldDebugMonitor= */ true, TEST_NORMAL_MONITORING_INTERVAL_MILLISECONDS,
                 TEST_DEBUG_MONITORING_INTERVAL_MILLISECONDS,
-                TEST_LATEST_AVAILABILITY_DURATION_MILLISECONDS);
+                TEST_LATEST_AVAILABILITY_DURATION_MILLISECONDS,
+                TEST_STOP_PERIODIC_CPUSET_READING_DELAY_MILLISECONDS);
 
         doNothing().when(() -> ServiceManager.addService(eq("cpu_monitor"), any(Binder.class),
                 anyBoolean(), anyInt()));
@@ -535,6 +538,18 @@
     }
 
     @Test
+    public void testBootCompleted() throws Exception {
+        mService.onBootPhase(PHASE_BOOT_COMPLETED);
+
+        // Message to stop periodic cpuset reading is posted on the service handler thread. Sync
+        // with this thread before proceeding.
+        syncWithHandler(mServiceHandler, /* delayMillis= */ 0);
+
+        verify(mMockCpuInfoReader, timeout(ASYNC_CALLBACK_WAIT_TIMEOUT_MILLISECONDS))
+                .stopPeriodicCpusetReading();
+    }
+
+    @Test
     public void testHeavyCpuLoadMonitoring() throws Exception {
         // TODO(b/267500110): Once heavy CPU load detection logic is added, add unittest.
     }
@@ -567,7 +582,8 @@
                 mServiceHandlerThread, /* shouldDebugMonitor= */ false,
                 TEST_NORMAL_MONITORING_INTERVAL_MILLISECONDS,
                 TEST_DEBUG_MONITORING_INTERVAL_MILLISECONDS,
-                TEST_LATEST_AVAILABILITY_DURATION_MILLISECONDS);
+                TEST_LATEST_AVAILABILITY_DURATION_MILLISECONDS,
+                TEST_STOP_PERIODIC_CPUSET_READING_DELAY_MILLISECONDS);
 
         startService();
     }
diff --git a/services/tests/mockingservicestests/src/com/android/server/crashrecovery/Android.bp b/services/tests/mockingservicestests/src/com/android/server/crashrecovery/Android.bp
index 127d3e8..7ac7aca 100644
--- a/services/tests/mockingservicestests/src/com/android/server/crashrecovery/Android.bp
+++ b/services/tests/mockingservicestests/src/com/android/server/crashrecovery/Android.bp
@@ -39,9 +39,9 @@
     ],
 
     libs: [
-        "android.test.mock",
-        "android.test.base",
-        "android.test.runner",
+        "android.test.mock.stubs.system",
+        "android.test.base.stubs.system",
+        "android.test.runner.stubs.system",
     ],
 
     jni_libs: [
diff --git a/services/tests/ondeviceintelligencetests/Android.bp b/services/tests/ondeviceintelligencetests/Android.bp
index aa85942..a31a3fb 100644
--- a/services/tests/ondeviceintelligencetests/Android.bp
+++ b/services/tests/ondeviceintelligencetests/Android.bp
@@ -47,9 +47,9 @@
     ],
 
     libs: [
-        "android.test.mock",
-        "android.test.base",
-        "android.test.runner",
+        "android.test.mock.stubs.system",
+        "android.test.base.stubs.system",
+        "android.test.runner.stubs.system",
     ],
 
     certificate: "platform",
diff --git a/services/tests/performancehinttests/Android.bp b/services/tests/performancehinttests/Android.bp
index 1692921c..c8121fc 100644
--- a/services/tests/performancehinttests/Android.bp
+++ b/services/tests/performancehinttests/Android.bp
@@ -19,7 +19,7 @@
         "truth",
     ],
     libs: [
-        "android.test.base",
+        "android.test.base.stubs.system",
     ],
     test_suites: [
         "automotive-tests",
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerCalculatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerCalculatorTest.java
index f74cfae..c0be865 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerCalculatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerCalculatorTest.java
@@ -56,14 +56,14 @@
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{300_000_000},
                 new int[]{Display.STATE_ON}, 0);
 
-        stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
-                30 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, Display.STATE_DOZE, Display.STATE_REASON_DEFAULT_POLICY,
+                30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS);
 
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{200_000_000},
                 new int[]{Display.STATE_DOZE}, 30 * MINUTE_IN_MS);
 
-        stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
-                120 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, Display.STATE_OFF, Display.STATE_REASON_DEFAULT_POLICY,
+                120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS);
 
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{100_000_000},
                 new int[]{Display.STATE_OFF}, 120 * MINUTE_IN_MS);
@@ -93,37 +93,37 @@
 
         final int[] screenStates = new int[] {Display.STATE_OFF, Display.STATE_OFF};
 
-        stats.noteScreenStateLocked(0, screenStates[0], 0, 0, 0);
-        stats.noteScreenStateLocked(1, screenStates[1], 0, 0, 0);
+        stats.noteScreenStateLocked(0, screenStates[0], Display.STATE_REASON_UNKNOWN, 0, 0, 0);
+        stats.noteScreenStateLocked(1, screenStates[1], Display.STATE_REASON_UNKNOWN, 0, 0, 0);
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{300, 400}, screenStates, 0);
 
         // Switch display0 to doze
         screenStates[0] = Display.STATE_DOZE;
-        stats.noteScreenStateLocked(0, screenStates[0], 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
-                30 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, screenStates[0], Display.STATE_REASON_UNKNOWN,
+                30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS);
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{200, 300},
                 screenStates, 30 * MINUTE_IN_MS);
 
         // Switch display1 to doze
         screenStates[1] = Display.STATE_DOZE;
-        stats.noteScreenStateLocked(1, Display.STATE_DOZE, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS,
-                90 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_DOZE, Display.STATE_REASON_UNKNOWN,
+                90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS);
         // 100,000,000 uC should be attributed to display 0 doze here.
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{100_000_000, 700_000_000},
                 screenStates, 90 * MINUTE_IN_MS);
 
         // Switch display0 to off
         screenStates[0] = Display.STATE_OFF;
-        stats.noteScreenStateLocked(0, screenStates[0], 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
-                120 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, screenStates[0], Display.STATE_REASON_UNKNOWN,
+                120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS);
         // 40,000,000 and 70,000,000 uC should be attributed to display 0 and 1 doze here.
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{40_000_000, 70_000_000},
                 screenStates, 120 * MINUTE_IN_MS);
 
         // Switch display1 to off
         screenStates[1] = Display.STATE_OFF;
-        stats.noteScreenStateLocked(1, screenStates[1], 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS,
-                150 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, screenStates[1], Display.STATE_REASON_UNKNOWN,
+                150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS);
         stats.updateDisplayEnergyConsumerStatsLocked(new long[]{100, 90_000_000}, screenStates,
                 150 * MINUTE_IN_MS);
         // 90,000,000 uC should be attributed to display 1 doze here.
@@ -148,10 +148,10 @@
     public void testPowerProfileBasedModel() {
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
-        stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
-                30 * MINUTE_IN_MS);
-        stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
-                120 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, Display.STATE_DOZE, Display.STATE_REASON_UNKNOWN,
+                30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
+                120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS);
 
         AmbientDisplayPowerCalculator calculator =
                 new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
@@ -174,15 +174,15 @@
 
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
-        stats.noteScreenStateLocked(1, Display.STATE_OFF, 0, 0, 0);
-        stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
-                30 * MINUTE_IN_MS);
-        stats.noteScreenStateLocked(1, Display.STATE_DOZE, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS,
-                90 * MINUTE_IN_MS);
-        stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
-                120 * MINUTE_IN_MS);
-        stats.noteScreenStateLocked(1, Display.STATE_OFF, 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS,
-                150 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN, 0, 0, 0);
+        stats.noteScreenStateLocked(0, Display.STATE_DOZE, Display.STATE_REASON_UNKNOWN,
+                30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_DOZE, Display.STATE_REASON_UNKNOWN,
+                90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
+                120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
+                150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS);
 
         AmbientDisplayPowerCalculator calculator =
                 new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
index afbe9159..2ccb642 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsNoteTest.java
@@ -44,6 +44,10 @@
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.WorkSource;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.DisabledOnRavenwood;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.platform.test.ravenwood.RavenwoodRule;
 import android.telephony.AccessNetworkConstants;
 import android.telephony.ActivityStatsTechSpecificInfo;
@@ -65,6 +69,7 @@
 import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.power.EnergyConsumerStats;
+import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.BatteryStatsImpl.DualTimer;
 
 import org.junit.Rule;
@@ -90,6 +95,8 @@
             .setProvideMainThread(true)
             .build();
 
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
     private static final String TAG = BatteryStatsNoteTest.class.getSimpleName();
 
     private static final int UID = 10500;
@@ -104,6 +111,54 @@
     @Mock
     NetworkStatsManager mNetworkStatsManager;
 
+    @DisabledOnRavenwood
+    @EnableFlags(Flags.FLAG_BATTERY_STATS_SCREEN_STATE_EVENT)
+    @Test
+    public void testScreenStateEvent_screenStateEventFlagOn_eventsRecorded() throws Exception {
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(new MockClock());
+        bi.forceRecordAllHistory();
+
+        bi.noteScreenStateLocked(0, Display.STATE_ON, Display.STATE_REASON_DEFAULT_POLICY,
+                0, 0, 0);
+        bi.noteScreenStateLocked(2, Display.STATE_DOZE_SUSPEND, Display.STATE_REASON_DRAW_WAKE_LOCK,
+                1, 1, 1);
+
+        BatteryStatsHistoryIterator iterator =
+                bi.iterateBatteryStatsHistory(0, MonotonicClock.UNDEFINED);
+        BatteryStats.HistoryItem item =
+                iterateAndFind(iterator, HistoryItem.EVENT_DISPLAY_STATE_CHANGED);
+        assertThat(item).isNotNull();
+        assertThat(item.eventTag).isNotNull();
+        assertThat(item.eventTag.string).isEqualTo("display=0 state=ON reason=DEFAULT_POLICY");
+        assertThat(item.eventTag.uid).isEqualTo(Process.INVALID_UID);
+
+        item = iterateAndFind(iterator, HistoryItem.EVENT_DISPLAY_STATE_CHANGED);
+        assertThat(item).isNotNull();
+        assertThat(item.eventTag).isNotNull();
+        assertThat(item.eventTag.string)
+                .isEqualTo("display=2 state=DOZE_SUSPEND reason=DRAW_WAKE_LOCK");
+        assertThat(item.eventTag.uid).isEqualTo(Process.INVALID_UID);
+
+        // Last check to make sure that we did not record any extra event.
+        assertThat(iterateAndFind(iterator, HistoryItem.EVENT_DISPLAY_STATE_CHANGED)).isNull();
+    }
+
+    @DisableFlags(Flags.FLAG_BATTERY_STATS_SCREEN_STATE_EVENT)
+    @Test
+    public void testScreenStateEvent_screenStateEventFlagOff_eventsNotRecorded() throws Exception {
+        MockBatteryStatsImpl bi = new MockBatteryStatsImpl(new MockClock());
+        bi.forceRecordAllHistory();
+
+        bi.noteScreenStateLocked(0, Display.STATE_ON, Display.STATE_REASON_DEFAULT_POLICY,
+                0, 0, 0);
+        bi.noteScreenStateLocked(2, Display.STATE_DOZE_SUSPEND, Display.STATE_REASON_DRAW_WAKE_LOCK,
+                1, 1, 1);
+
+        BatteryStatsHistoryIterator iterator =
+                bi.iterateBatteryStatsHistory(0, MonotonicClock.UNDEFINED);
+        assertThat(iterateAndFind(iterator, HistoryItem.EVENT_DISPLAY_STATE_CHANGED)).isNull();
+    }
+
     /**
      * Test BatteryStatsImpl.Uid.noteBluetoothScanResultLocked.
      */
@@ -285,20 +340,15 @@
         final BatteryStatsHistoryIterator iterator =
                 bi.iterateBatteryStatsHistory(0, MonotonicClock.UNDEFINED);
 
-        BatteryStats.HistoryItem item;
+        BatteryStats.HistoryItem item =
+                iterateAndFind(iterator, HistoryItem.EVENT_LONG_WAKE_LOCK_START);
 
-        while ((item = iterator.next()) != null) {
-            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_START) break;
-        }
-        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_START);
+        assertThat(item).isNotNull();
         assertThat(item.eventTag).isNotNull();
         assertThat(item.eventTag.string).isEqualTo(historyName);
         assertThat(item.eventTag.uid).isEqualTo(UID);
 
-        while ((item = iterator.next()) != null) {
-            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH) break;
-        }
-        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH);
+        item = iterateAndFind(iterator, HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH);
         assertThat(item.eventTag).isNotNull();
         assertThat(item.eventTag.string).isEqualTo(historyName);
         assertThat(item.eventTag.uid).isEqualTo(UID);
@@ -343,20 +393,15 @@
         final BatteryStatsHistoryIterator iterator =
                 bi.iterateBatteryStatsHistory(0, MonotonicClock.UNDEFINED);
 
-        BatteryStats.HistoryItem item;
-
-        while ((item = iterator.next()) != null) {
-            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_START) break;
-        }
-        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_START);
+        BatteryStats.HistoryItem item =
+                iterateAndFind(iterator, HistoryItem.EVENT_LONG_WAKE_LOCK_START);
+        assertThat(item).isNotNull();
         assertThat(item.eventTag).isNotNull();
         assertThat(item.eventTag.string).isEqualTo(historyName);
         assertThat(item.eventTag.uid).isEqualTo(UID);
 
-        while ((item = iterator.next()) != null) {
-            if (item.eventCode == HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH) break;
-        }
-        assertThat(item.eventCode).isEqualTo(HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH);
+        item = iterateAndFind(iterator, HistoryItem.EVENT_LONG_WAKE_LOCK_FINISH);
+        assertThat(item).isNotNull();
         assertThat(item.eventTag).isNotNull();
         assertThat(item.eventTag.string).isEqualTo(historyName);
         assertThat(item.eventTag.uid).isEqualTo(UID);
@@ -2562,4 +2607,18 @@
                     currentTimeMs, currentTimeMs, mNetworkStatsManager);
         }
     }
+
+    /**
+     * Moves a given {@link BatteryStatsHistoryIterator} until a history item with the given
+     * {@code eventCode} is found and returns the history item. Returns {@code null} if no such item
+     * is found.
+     */
+    private static BatteryStats.HistoryItem iterateAndFind(
+                BatteryStatsHistoryIterator iterator, int eventCode) {
+        BatteryStats.HistoryItem item;
+        while ((item = iterator.next()) != null) {
+            if (item.eventCode == eventCode) return item;
+        }
+        return null;
+    }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerCalculatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerCalculatorTest.java
index 88d4ea7..2da98e8 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerCalculatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerCalculatorTest.java
@@ -61,7 +61,8 @@
         mStatsRule.initMeasuredEnergyStatsLocked();
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
 
-        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
+        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, Display.STATE_REASON_UNKNOWN,
+                0, 0, 0);
         batteryStats.updateDisplayEnergyConsumerStatsLocked(new long[]{0},
                 new int[]{Display.STATE_ON}, 0);
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
@@ -79,7 +80,7 @@
         batteryStats.updateDisplayEnergyConsumerStatsLocked(new long[]{300_000_000},
                 new int[]{Display.STATE_ON}, 60 * MINUTE_IN_MS);
 
-        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
+        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
@@ -150,8 +151,10 @@
 
         final int[] screenStates = new int[]{Display.STATE_ON, Display.STATE_OFF};
 
-        batteryStats.noteScreenStateLocked(0, screenStates[0], 0, 0, 0);
-        batteryStats.noteScreenStateLocked(1, screenStates[1], 0, 0, 0);
+        batteryStats.noteScreenStateLocked(0, screenStates[0], Display.STATE_REASON_UNKNOWN,
+                0, 0, 0);
+        batteryStats.noteScreenStateLocked(1, screenStates[1], Display.STATE_REASON_UNKNOWN,
+                0, 0, 0);
         batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true, 0, 0);
         batteryStats.updateDisplayEnergyConsumerStatsLocked(new long[]{300, 400}, screenStates, 0);
@@ -166,10 +169,10 @@
 
         screenStates[0] = Display.STATE_OFF;
         screenStates[1] = Display.STATE_ON;
-        batteryStats.noteScreenStateLocked(0, screenStates[0],
+        batteryStats.noteScreenStateLocked(0, screenStates[0], Display.STATE_REASON_UNKNOWN,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
-        batteryStats.noteScreenStateLocked(1, screenStates[1], 80 * MINUTE_IN_MS,
-                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, screenStates[1], Display.STATE_REASON_UNKNOWN,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         batteryStats.updateDisplayEnergyConsumerStatsLocked(new long[]{600_000_000, 500},
                 screenStates, 80 * MINUTE_IN_MS);
 
@@ -178,8 +181,8 @@
         batteryStats.noteScreenBrightnessLocked(1, 75, 98 * MINUTE_IN_MS, 98 * MINUTE_IN_MS);
 
         screenStates[1] = Display.STATE_OFF;
-        batteryStats.noteScreenStateLocked(1, screenStates[1], 110 * MINUTE_IN_MS,
-                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, screenStates[1], Display.STATE_REASON_UNKNOWN,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
         batteryStats.updateDisplayEnergyConsumerStatsLocked(new long[]{700, 800_000_000},
                 screenStates, 110 * MINUTE_IN_MS);
 
@@ -240,7 +243,8 @@
     public void testPowerProfileBasedModel() {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
 
-        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
+        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, Display.STATE_REASON_UNKNOWN,
+                0, 0, 0);
         batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
                 0, 0);
@@ -253,7 +257,7 @@
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
 
-        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
+        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
@@ -313,8 +317,10 @@
 
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
 
-        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
-        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, 0, 0, 0);
+        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, Display.STATE_REASON_UNKNOWN,
+                0, 0, 0);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
+                0, 0, 0);
         batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
                 0, 0);
@@ -327,16 +333,16 @@
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
 
-        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
+        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
-        batteryStats.noteScreenStateLocked(1, Display.STATE_ON, 80 * MINUTE_IN_MS,
-                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_ON, Display.STATE_REASON_UNKNOWN,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         batteryStats.noteScreenBrightnessLocked(1, 20, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
 
         batteryStats.noteScreenBrightnessLocked(1, 250, 86 * MINUTE_IN_MS, 86 * MINUTE_IN_MS);
         batteryStats.noteScreenBrightnessLocked(1, 75, 98 * MINUTE_IN_MS, 98 * MINUTE_IN_MS);
-        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, 110 * MINUTE_IN_MS,
-                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, Display.STATE_REASON_UNKNOWN,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
 
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
                 110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
diff --git a/services/tests/selinux/Android.bp b/services/tests/selinux/Android.bp
index 12a7038..048978a 100644
--- a/services/tests/selinux/Android.bp
+++ b/services/tests/selinux/Android.bp
@@ -42,9 +42,9 @@
         "mockito_extended",
     ],
     libs: [
-        "android.test.base",
-        "android.test.mock",
-        "android.test.runner",
+        "android.test.base.stubs.system",
+        "android.test.mock.stubs.system",
+        "android.test.runner.stubs.system",
         "servicestests-core-utils",
     ],
     static_libs: [
diff --git a/services/tests/servicestests/Android.bp b/services/tests/servicestests/Android.bp
index b41b30c..bbe0755 100644
--- a/services/tests/servicestests/Android.bp
+++ b/services/tests/servicestests/Android.bp
@@ -91,6 +91,7 @@
         "net_flags_lib",
         "CtsVirtualDeviceCommonLib",
         "com_android_server_accessibility_flags_lib",
+        "locksettings_flags_lib",
     ],
 
     libs: [
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
index d6f7e21..d071c15 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
@@ -60,6 +60,9 @@
 import android.os.ServiceSpecificException;
 import android.os.UserManager;
 import android.platform.test.annotations.Presubmit;
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
@@ -70,6 +73,7 @@
 import com.android.server.pm.UserManagerInternal;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -108,6 +112,9 @@
             0x26, 0x52, 0x72, 0x63, 0x63, 0x61, 0x78, 0x23,
     };
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
     private Context mContext;
     private UserManager mUserManager;
     private UserManagerInternal mUserManagerInternal;
@@ -145,7 +152,6 @@
         private RebootEscrowProviderInterface mRebootEscrowProviderInUse;
         private ConnectivityManager.NetworkCallback mNetworkCallback;
         private Consumer<ConnectivityManager.NetworkCallback> mNetworkConsumer;
-        private boolean mWaitForInternet;
 
         MockInjector(
                 Context context,
@@ -159,7 +165,6 @@
             super(context, storage, userManagerInternal);
             mRebootEscrow = rebootEscrow;
             mServerBased = false;
-            mWaitForInternet = false;
             RebootEscrowProviderHalImpl.Injector halInjector =
                     new RebootEscrowProviderHalImpl.Injector() {
                         @Override
@@ -185,7 +190,6 @@
             super(context, storage, userManagerInternal);
             mRebootEscrow = null;
             mServerBased = true;
-            mWaitForInternet = false;
             RebootEscrowProviderServerBasedImpl.Injector injector =
                     new RebootEscrowProviderServerBasedImpl.Injector(serviceConnection) {
                         @Override
@@ -227,15 +231,6 @@
         }
 
         @Override
-        public boolean waitForInternet() {
-            return mWaitForInternet;
-        }
-
-        public void setWaitForNetwork(boolean waitForNetworkEnabled) {
-            mWaitForInternet = waitForNetworkEnabled;
-        }
-
-        @Override
         public boolean isNetworkConnected() {
             return false;
         }
@@ -934,10 +929,10 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_serverBasedWaitForInternet_success()
             throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
@@ -987,10 +982,10 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_serverBasedWaitForInternetRemoteException_Failure()
             throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
@@ -1042,10 +1037,10 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_waitForInternet_networkUnavailable()
             throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
@@ -1090,9 +1085,9 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_waitForInternet_networkLost() throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
@@ -1145,10 +1140,10 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_waitForInternet_networkAvailableWithDelay()
             throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
@@ -1204,10 +1199,10 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_waitForInternet_timeoutExhausted()
             throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
@@ -1264,10 +1259,10 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_serverBasedWaitForNetwork_retryCountExhausted()
             throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
@@ -1320,10 +1315,10 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WAIT_FOR_INTERNET_ROR)
     public void loadRebootEscrowDataIfAvailable_ServerBasedWaitForInternet_RetrySuccess()
             throws Exception {
         setServerBasedRebootEscrowProvider();
-        mMockInjector.setWaitForNetwork(true);
 
         when(mInjected.getBootCount()).thenReturn(0);
         RebootEscrowListener mockListener = mock(RebootEscrowListener.class);
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java
index f8ff1f4..efcf027 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java
@@ -756,7 +756,7 @@
         assertEquals("a", fromXml.getPkg());
 
         fromXml.condition = new Condition(Uri.EMPTY, "", Condition.STATE_TRUE);
-        assertTrue(fromXml.isAutomaticActive());
+        assertTrue(fromXml.isActive());
     }
 
     @Test
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
index baa633f..39a9d30 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
@@ -5788,7 +5788,7 @@
 
         // ... but it is NOT active
         ZenRule storedRule = mZenModeHelper.mConfig.automaticRules.get(newRuleId);
-        assertThat(storedRule.isAutomaticActive()).isFalse();
+        assertThat(storedRule.isActive()).isFalse();
         assertThat(storedRule.isTrueOrUnknown()).isFalse();
         assertThat(storedRule.condition).isNull();
         assertThat(mZenModeHelper.getZenMode()).isEqualTo(ZEN_MODE_OFF);
@@ -5841,7 +5841,7 @@
 
         // ... but it is NEITHER active NOR snoozed.
         ZenRule storedRule = mZenModeHelper.mConfig.automaticRules.get(newRuleId);
-        assertThat(storedRule.isAutomaticActive()).isFalse();
+        assertThat(storedRule.isActive()).isFalse();
         assertThat(storedRule.isTrueOrUnknown()).isFalse();
         assertThat(storedRule.condition).isNull();
         assertThat(storedRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
@@ -6619,7 +6619,7 @@
                 new Condition(rule.getConditionId(), "manual-on", STATE_TRUE, SOURCE_USER_ACTION),
                 ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_ACTIVATE);
         assertThat(zenRule.condition).isNull();
 
@@ -6627,14 +6627,14 @@
                 new Condition(rule.getConditionId(), "manual-off", STATE_FALSE, SOURCE_USER_ACTION),
                 ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isFalse();
+        assertThat(zenRule.isActive()).isFalse();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
         assertThat(zenRule.condition).isNull();
 
         // Bonus check: app has resumed control over the rule and can now turn it on.
         mZenModeHelper.setAutomaticZenRuleState(ruleId, autoOn,  ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
         assertThat(zenRule.condition).isEqualTo(autoOn);
     }
@@ -6655,7 +6655,7 @@
 
         mZenModeHelper.setAutomaticZenRuleState(ruleId, autoOn,  ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
         assertThat(zenRule.condition).isEqualTo(autoOn);
 
@@ -6663,7 +6663,7 @@
                 new Condition(rule.getConditionId(), "manual-off", STATE_FALSE, SOURCE_USER_ACTION),
                 ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isFalse();
+        assertThat(zenRule.isActive()).isFalse();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_DEACTIVATE);
         assertThat(zenRule.condition).isEqualTo(autoOn);
 
@@ -6671,14 +6671,14 @@
                 new Condition(rule.getConditionId(), "manual-on", STATE_TRUE, SOURCE_USER_ACTION),
                 ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
         assertThat(zenRule.condition).isEqualTo(autoOn);
 
         // Bonus check: app has resumed control over the rule and can now turn it off.
         mZenModeHelper.setAutomaticZenRuleState(ruleId, autoOff,  ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isFalse();
+        assertThat(zenRule.isActive()).isFalse();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
         assertThat(zenRule.condition).isEqualTo(autoOff);
     }
@@ -6696,7 +6696,7 @@
                 new Condition(rule.getConditionId(), "auto-on", STATE_TRUE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         ZenRule zenRuleOn = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRuleOn.isAutomaticActive()).isTrue();
+        assertThat(zenRuleOn.isActive()).isTrue();
         assertThat(zenRuleOn.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
         assertThat(zenRuleOn.condition).isNotNull();
 
@@ -6704,7 +6704,7 @@
                 new Condition(rule.getConditionId(), "manual-off", STATE_FALSE, SOURCE_USER_ACTION),
                 ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
         ZenRule zenRuleOff = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRuleOff.isAutomaticActive()).isFalse();
+        assertThat(zenRuleOff.isActive()).isFalse();
         assertThat(zenRuleOff.getConditionOverride()).isEqualTo(OVERRIDE_DEACTIVATE);
         assertThat(zenRuleOff.condition).isNotNull();
     }
@@ -6723,27 +6723,27 @@
                 new Condition(rule.getConditionId(), "manual-on", STATE_TRUE, SOURCE_USER_ACTION),
                 ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_ACTIVATE);
 
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-off", STATE_FALSE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_ACTIVATE);
 
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-on", STATE_TRUE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-off", STATE_FALSE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isFalse();
+        assertThat(zenRule.isActive()).isFalse();
     }
 
     @Test
@@ -6760,35 +6760,35 @@
                 new Condition(rule.getConditionId(), "auto-on", STATE_TRUE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
 
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "manual-off", STATE_FALSE, SOURCE_USER_ACTION),
                 ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isFalse();
+        assertThat(zenRule.isActive()).isFalse();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_DEACTIVATE);
 
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-on", STATE_TRUE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isFalse();
+        assertThat(zenRule.isActive()).isFalse();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_DEACTIVATE);
 
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-off", STATE_FALSE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isFalse();
+        assertThat(zenRule.isActive()).isFalse();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
 
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-on", STATE_TRUE, SOURCE_CONTEXT),
                 ORIGIN_APP, CUSTOM_PKG_UID);
         zenRule = mZenModeHelper.mConfig.automaticRules.get(ruleId);
-        assertThat(zenRule.isAutomaticActive()).isTrue();
+        assertThat(zenRule.isActive()).isTrue();
         assertThat(zenRule.getConditionOverride()).isEqualTo(OVERRIDE_NONE);
     }
 
@@ -6805,14 +6805,14 @@
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "manual-on-from-sysui", STATE_TRUE,
                         SOURCE_USER_ACTION), ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
-        assertThat(getZenRule(ruleId).isAutomaticActive()).isTrue();
+        assertThat(getZenRule(ruleId).isActive()).isTrue();
         assertThat(getZenRule(ruleId).getConditionOverride()).isEqualTo(OVERRIDE_ACTIVATE);
 
         // ... and they can turn it off manually from inside the app.
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "manual-off-from-app", STATE_FALSE,
                         SOURCE_USER_ACTION), ORIGIN_USER_IN_APP, CUSTOM_PKG_UID);
-        assertThat(getZenRule(ruleId).isAutomaticActive()).isFalse();
+        assertThat(getZenRule(ruleId).isActive()).isFalse();
         assertThat(getZenRule(ruleId).getConditionOverride()).isEqualTo(OVERRIDE_NONE);
     }
 
@@ -6829,21 +6829,21 @@
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-on-from-app", STATE_TRUE,
                         SOURCE_SCHEDULE), ORIGIN_APP, CUSTOM_PKG_UID);
-        assertThat(getZenRule(ruleId).isAutomaticActive()).isTrue();
+        assertThat(getZenRule(ruleId).isActive()).isTrue();
         assertThat(getZenRule(ruleId).getConditionOverride()).isEqualTo(OVERRIDE_NONE);
 
         // User manually turns off rule from SysUI / Settings...
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "manual-off-from-sysui", STATE_FALSE,
                         SOURCE_USER_ACTION), ORIGIN_USER_IN_SYSTEMUI, SYSTEM_UID);
-        assertThat(getZenRule(ruleId).isAutomaticActive()).isFalse();
+        assertThat(getZenRule(ruleId).isActive()).isFalse();
         assertThat(getZenRule(ruleId).getConditionOverride()).isEqualTo(OVERRIDE_DEACTIVATE);
 
         // ... and they can turn it on manually from inside the app.
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "manual-on-from-app", STATE_TRUE,
                         SOURCE_USER_ACTION), ORIGIN_USER_IN_APP, CUSTOM_PKG_UID);
-        assertThat(getZenRule(ruleId).isAutomaticActive()).isTrue();
+        assertThat(getZenRule(ruleId).isActive()).isTrue();
         assertThat(getZenRule(ruleId).getConditionOverride()).isEqualTo(OVERRIDE_NONE);
     }
 
@@ -6861,14 +6861,14 @@
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "manual-on-from-app", STATE_TRUE,
                         SOURCE_USER_ACTION), ORIGIN_USER_IN_APP, CUSTOM_PKG_UID);
-        assertThat(getZenRule(ruleId).isAutomaticActive()).isTrue();
+        assertThat(getZenRule(ruleId).isActive()).isTrue();
         assertThat(getZenRule(ruleId).getConditionOverride()).isEqualTo(OVERRIDE_NONE);
 
         // ... so the app can turn it off when its schedule is over.
         mZenModeHelper.setAutomaticZenRuleState(ruleId,
                 new Condition(rule.getConditionId(), "auto-off-from-app", STATE_FALSE,
                         SOURCE_SCHEDULE), ORIGIN_APP, CUSTOM_PKG_UID);
-        assertThat(getZenRule(ruleId).isAutomaticActive()).isFalse();
+        assertThat(getZenRule(ruleId).isActive()).isFalse();
         assertThat(getZenRule(ruleId).getConditionOverride()).isEqualTo(OVERRIDE_NONE);
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index 71cfbfd..08622e6 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -133,6 +133,7 @@
     private final ArrayList<DeviceConfig.OnPropertiesChangedListener> mDeviceConfigListeners =
             new ArrayList<>();
 
+    private AppCompatConfiguration mAppCompat;
     private Description mDescription;
     private Context mContext;
     private StaticMockitoSession mMockitoSession;
@@ -379,6 +380,11 @@
                 mock(ActivityManagerService.class, withSettings().stubOnly());
         mAtmService = new TestActivityTaskManagerService(mContext, amService);
         LocalServices.addService(ActivityTaskManagerInternal.class, mAtmService.getAtmInternal());
+
+        // AppCompatConfiguration
+        mAppCompat = new AppCompatConfiguration(
+                ActivityThread.currentActivityThread().getSystemUiContext());
+
         // Create a fake WindowProcessController for the system process.
         final WindowProcessController wpc =
                 addProcess("android", "system", 1485 /* pid */, 1000 /* uid */);
@@ -394,7 +400,7 @@
         mWmService = WindowManagerService.main(
                 mContext, mImService, false, wmPolicy, mAtmService,
                 testDisplayWindowSettingsProvider, StubTransaction::new,
-                MockSurfaceControlBuilder::new);
+                MockSurfaceControlBuilder::new, mAppCompat);
         spyOn(mWmService);
         spyOn(mWmService.mRoot);
         // Invoked during {@link ActivityStack} creation.
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
index 5a54af1..2d5e5da 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowStateTests.java
@@ -80,6 +80,7 @@
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.atMost;
 import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.when;
 
@@ -88,9 +89,12 @@
 import android.graphics.Matrix;
 import android.graphics.Point;
 import android.graphics.Rect;
+import android.graphics.Region;
 import android.os.IBinder;
 import android.os.InputConfig;
 import android.os.RemoteException;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.util.ArraySet;
@@ -116,6 +120,7 @@
 import com.android.server.inputmethod.InputMethodManagerInternal;
 import com.android.server.testutils.StubTransaction;
 import com.android.server.wm.SensitiveContentPackages.PackageInfo;
+import com.android.window.flags.Flags;
 
 import org.junit.After;
 import org.junit.Test;
@@ -965,6 +970,88 @@
         assertTrue(testFlag(handle.inputConfig, InputConfig.NO_INPUT_CHANNEL));
     }
 
+    @DisableFlags(Flags.FLAG_SCROLLING_FROM_LETTERBOX)
+    @Test
+    public void testTouchRegionUsesLetterboxBoundsIfTransformedBoundsAndLetterboxScrolling() {
+        final WindowState win = createWindow(null, TYPE_APPLICATION, "win");
+
+        // Transformed bounds used for size of touchable region if letterbox inner bounds are empty.
+        final Rect transformedBounds = new Rect(0, 0, 300, 500);
+        doReturn(transformedBounds).when(win.mToken).getFixedRotationTransformDisplayBounds();
+
+        // Otherwise, touchable region should match letterbox inner bounds.
+        final Rect letterboxInnerBounds = new Rect(30, 0, 270, 500);
+        doAnswer(invocation -> {
+            Rect rect = invocation.getArgument(0);
+            rect.set(letterboxInnerBounds);
+            return null;
+        }).when(win.mActivityRecord).getLetterboxInnerBounds(any());
+
+        Region outRegion = new Region();
+        win.getSurfaceTouchableRegion(outRegion, win.mAttrs);
+
+        // Because scrollingFromLetterbox flag is disabled and letterboxInnerBounds is not empty,
+        // touchable region should match letterboxInnerBounds always.
+        assertEquals(letterboxInnerBounds, outRegion.getBounds());
+    }
+
+    @DisableFlags(Flags.FLAG_SCROLLING_FROM_LETTERBOX)
+    @Test
+    public void testTouchRegionUsesLetterboxBoundsIfNullTransformedBoundsAndLetterboxScrolling() {
+        final WindowState win = createWindow(null, TYPE_APPLICATION, "win");
+
+        // Fragment bounds used for size of touchable region if letterbox inner bounds are empty
+        // and Transform bounds are null.
+        doReturn(null).when(win.mToken).getFixedRotationTransformDisplayBounds();
+        final Rect fragmentBounds = new Rect(0, 0, 300, 500);
+        final TaskFragment taskFragment = win.mActivityRecord.getTaskFragment();
+        doAnswer(invocation -> {
+            Rect rect = invocation.getArgument(0);
+            rect.set(fragmentBounds);
+            return null;
+        }).when(taskFragment).getDimBounds(any());
+
+        // Otherwise, touchable region should match letterbox inner bounds.
+        final Rect letterboxInnerBounds = new Rect(30, 0, 270, 500);
+        doAnswer(invocation -> {
+            Rect rect = invocation.getArgument(0);
+            rect.set(letterboxInnerBounds);
+            return null;
+        }).when(win.mActivityRecord).getLetterboxInnerBounds(any());
+
+        Region outRegion = new Region();
+        win.getSurfaceTouchableRegion(outRegion, win.mAttrs);
+
+        // Because scrollingFromLetterbox flag is disabled and letterboxInnerBounds is not empty,
+        // touchable region should match letterboxInnerBounds always.
+        assertEquals(letterboxInnerBounds, outRegion.getBounds());
+    }
+
+    @EnableFlags(Flags.FLAG_SCROLLING_FROM_LETTERBOX)
+    @Test
+    public void testTouchRegionUsesTransformedBoundsIfLetterboxScrolling() {
+        final WindowState win = createWindow(null, TYPE_APPLICATION, "win");
+
+        // Transformed bounds used for size of touchable region if letterbox inner bounds are empty.
+        final Rect transformedBounds = new Rect(0, 0, 300, 500);
+        doReturn(transformedBounds).when(win.mToken).getFixedRotationTransformDisplayBounds();
+
+        // Otherwise, touchable region should match letterbox inner bounds.
+        final Rect letterboxInnerBounds = new Rect(30, 0, 270, 500);
+        doAnswer(invocation -> {
+            Rect rect = invocation.getArgument(0);
+            rect.set(letterboxInnerBounds);
+            return null;
+        }).when(win.mActivityRecord).getLetterboxInnerBounds(any());
+
+        Region outRegion = new Region();
+        win.getSurfaceTouchableRegion(outRegion, win.mAttrs);
+
+        // Because scrollingFromLetterbox flag is enabled and transformedBounds are non-null,
+        // touchable region should match transformedBounds.
+        assertEquals(transformedBounds, outRegion.getBounds());
+    }
+
     @Test
     public void testHasActiveVisibleWindow() {
         final int uid = ActivityBuilder.DEFAULT_FAKE_UID;
diff --git a/tests/Input/AndroidManifest.xml b/tests/Input/AndroidManifest.xml
index a05d08c..914adc4 100644
--- a/tests/Input/AndroidManifest.xml
+++ b/tests/Input/AndroidManifest.xml
@@ -32,6 +32,14 @@
              android:process=":externalProcess">
         </activity>
 
+        <activity android:name="com.android.test.input.CaptureEventActivity"
+            android:label="Capture events"
+            android:configChanges="touchscreen|uiMode|orientation|screenSize|screenLayout|keyboardHidden|uiMode|navigation|keyboard|density|fontScale|layoutDirection|locale|mcc|mnc|smallestScreenSize"
+            android:enableOnBackInvokedCallback="false"
+            android:turnScreenOn="true"
+            android:exported="true">
+        </activity>
+
     </application>
     <instrumentation android:name="androidx.test.runner.AndroidJUnitRunner"
                      android:targetPackage="com.android.test.input"
diff --git a/tests/Input/res/raw/google_pixel_tablet_touchscreen.evemu b/tests/Input/res/raw/google_pixel_tablet_touchscreen.evemu
new file mode 100644
index 0000000..1a9112b
--- /dev/null
+++ b/tests/Input/res/raw/google_pixel_tablet_touchscreen.evemu
@@ -0,0 +1,150 @@
+# EVEMU 1.2
+# One finger swipe gesture on the Google Pixel Tablet touchscreen
+N: NVTCapacitiveTouchScreen
+I: 001c 0603 7806 0100
+P: 02 00 00 00 00 00 00 00
+B: 00 0b 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 80 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 04 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 01 00 00 00 00 00 00 00 00
+B: 02 00 00 00 00 00 00 00 00
+B: 03 00 00 00 00 00 80 f3 46
+B: 04 00 00 00 00 00 00 00 00
+B: 05 00 00 00 00 00 00 00 00
+B: 11 00 00 00 00 00 00 00 00
+B: 12 00 00 00 00 00 00 00 00
+A: 2f 0 9 0 0 0
+A: 30 0 2559 0 0 11
+A: 31 0 1599 0 0 11
+A: 34 -4096 4096 0 0 0
+A: 35 0 1599 0 0 11
+A: 36 0 2559 0 0 11
+A: 37 0 2 0 0 0
+A: 39 0 65535 0 0 0
+A: 3a 0 256 0 0 0
+A: 3e 0 8 0 0 0
+E: 0.000001 0001 014a 0001
+E: 0.000001 0003 0039 0003
+E: 0.000001 0003 0035 0810
+E: 0.000001 0003 0036 1650
+E: 0.000001 0003 0030 0082
+E: 0.000001 0003 0031 0077
+E: 0.000001 0003 003a 0215
+E: 0.000001 0003 0034 3218
+E: 0.000001 0000 0000 0000
+E: 0.008818 0003 0035 0825
+E: 0.008818 0003 0036 1645
+E: 0.008818 0003 0034 3217
+E: 0.008818 0000 0000 0000
+E: 0.016306 0003 0035 0841
+E: 0.016306 0003 0036 1639
+E: 0.016306 0003 0034 3102
+E: 0.016306 0000 0000 0000
+E: 0.025653 0003 0035 0862
+E: 0.025653 0003 0036 1630
+E: 0.025653 0003 0034 3092
+E: 0.025653 0000 0000 0000
+E: 0.032936 0003 0035 0883
+E: 0.032936 0003 0036 1619
+E: 0.032936 0003 0034 3030
+E: 0.032936 0000 0000 0000
+E: 0.042072 0003 0035 0905
+E: 0.042072 0003 0036 1604
+E: 0.042072 0003 0034 2848
+E: 0.042072 0000 0000 0000
+E: 0.049569 0003 0035 0924
+E: 0.049569 0003 0036 1591
+E: 0.049569 0003 0034 2830
+E: 0.049569 0000 0000 0000
+E: 0.058706 0003 0035 0942
+E: 0.058706 0003 0036 1573
+E: 0.058706 0000 0000 0000
+E: 0.066207 0003 0035 0954
+E: 0.066207 0003 0036 1557
+E: 0.066207 0003 0034 2790
+E: 0.066207 0000 0000 0000
+E: 0.075337 0003 0035 0966
+E: 0.075337 0003 0036 1535
+E: 0.075337 0000 0000 0000
+E: 0.082841 0003 0035 0973
+E: 0.082841 0003 0036 1511
+E: 0.082841 0003 0034 2788
+E: 0.082841 0000 0000 0000
+E: 0.091972 0003 0035 0971
+E: 0.091972 0003 0036 1480
+E: 0.091972 0003 0034 2770
+E: 0.091972 0000 0000 0000
+E: 0.099474 0003 0035 0961
+E: 0.099474 0003 0036 1445
+E: 0.099474 0003 0034 2644
+E: 0.099474 0000 0000 0000
+E: 0.108631 0003 0035 0937
+E: 0.108631 0003 0036 1400
+E: 0.108631 0003 0030 0083
+E: 0.108631 0003 0034 2461
+E: 0.108631 0000 0000 0000
+E: 0.116109 0003 0035 0909
+E: 0.116109 0003 0036 1361
+E: 0.116109 0003 0034 2278
+E: 0.116109 0000 0000 0000
+E: 0.125263 0003 0035 0865
+E: 0.125263 0003 0036 1311
+E: 0.125263 0003 0034 2096
+E: 0.125263 0000 0000 0000
+E: 0.132741 0003 0035 0820
+E: 0.132741 0003 0036 1261
+E: 0.132741 0003 0034 2083
+E: 0.132741 0000 0000 0000
+E: 0.141876 0003 0035 0755
+E: 0.141876 0003 0036 1193
+E: 0.141876 0003 003a 0216
+E: 0.141876 0003 0034 2266
+E: 0.141876 0000 0000 0000
+E: 0.149376 0003 0035 0691
+E: 0.149376 0003 0036 1124
+E: 0.149376 0003 0034 2448
+E: 0.149376 0000 0000 0000
+E: 0.158510 0003 0035 0609
+E: 0.158510 0003 0036 1033
+E: 0.158510 0003 0034 2631
+E: 0.158510 0000 0000 0000
+E: 0.166011 0003 0035 0543
+E: 0.166011 0003 0036 0957
+E: 0.166011 0003 0034 2813
+E: 0.166011 0000 0000 0000
+E: 0.175182 0003 0035 0471
+E: 0.175182 0003 0036 0864
+E: 0.175182 0003 0031 0076
+E: 0.175182 0003 0034 2996
+E: 0.175182 0000 0000 0000
+E: 0.182683 0003 0035 0417
+E: 0.182683 0003 0036 0792
+E: 0.182683 0003 003a 0214
+E: 0.182683 0003 0034 3178
+E: 0.182683 0000 0000 0000
+E: 0.191777 0003 0035 0361
+E: 0.191777 0003 0036 0719
+E: 0.191777 0003 0031 0075
+E: 0.191777 0003 003a 0213
+E: 0.191777 0003 0034 2996
+E: 0.191777 0000 0000 0000
+E: 0.199431 0003 0035 0271
+E: 0.199431 0003 0036 0603
+E: 0.199431 0003 0030 0060
+E: 0.199431 0003 0031 0029
+E: 0.199431 0003 003a 0060
+E: 0.199431 0003 0034 2813
+E: 0.199431 0000 0000 0000
+E: 0.207943 0003 003a 0000
+E: 0.207943 0003 0039 -001
+E: 0.207943 0001 014a 0000
+E: 0.207943 0000 0000 0000
diff --git a/tests/Input/res/raw/google_pixel_tablet_touchscreen_events.json b/tests/Input/res/raw/google_pixel_tablet_touchscreen_events.json
new file mode 100644
index 0000000..df4f9fb
--- /dev/null
+++ b/tests/Input/res/raw/google_pixel_tablet_touchscreen_events.json
@@ -0,0 +1,34 @@
+[
+  {
+    "name": "One finger swipe",
+    "source": "TOUCHSCREEN",
+    "events": [
+      {"action":"DOWN","axes":{"AXIS_X":810,"AXIS_Y":1650,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.234087586402893},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":825,"AXIS_Y":1645,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.2337040901184082},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":841,"AXIS_Y":1639,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.1896021366119385},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":862,"AXIS_Y":1630,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.1857671737670898},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":883,"AXIS_Y":1619,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.1619905233383179},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":905,"AXIS_Y":1604,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0921943187713623},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":924,"AXIS_Y":1591,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0852913856506348},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":942,"AXIS_Y":1573,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0852913856506348},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":954,"AXIS_Y":1557,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0699516534805298},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":966,"AXIS_Y":1535,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0699516534805298},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":973,"AXIS_Y":1511,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.06918466091156},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":971,"AXIS_Y":1480,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0622817277908325},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":961,"AXIS_Y":1445,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":82,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":82,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0139613151550293},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":937,"AXIS_Y":1400,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":0.9437817335128784},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":909,"AXIS_Y":1361,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":0.8736020922660828},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":865,"AXIS_Y":1311,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":0.803805947303772},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":820,"AXIS_Y":1261,"AXIS_PRESSURE":0.83984375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":0.7988204956054688},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":755,"AXIS_Y":1193,"AXIS_PRESSURE":0.84375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":0.8690001368522644},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":691,"AXIS_Y":1124,"AXIS_PRESSURE":0.84375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":0.9387962818145752},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":609,"AXIS_Y":1033,"AXIS_PRESSURE":0.84375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.008975863456726},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":543,"AXIS_Y":957,"AXIS_PRESSURE":0.84375,"AXIS_SIZE":0.03126221150159836,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":77,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":77,"AXIS_ORIENTATION":1.0787720680236816},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":471,"AXIS_Y":864,"AXIS_PRESSURE":0.84375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":76,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":76,"AXIS_ORIENTATION":1.1489516496658325},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":417,"AXIS_Y":792,"AXIS_PRESSURE":0.8359375,"AXIS_SIZE":0.03106682375073433,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":76,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":76,"AXIS_ORIENTATION":1.2187477350234985},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":361,"AXIS_Y":719,"AXIS_PRESSURE":0.83203125,"AXIS_SIZE":0.03087143413722515,"AXIS_TOUCH_MAJOR":83,"AXIS_TOUCH_MINOR":75,"AXIS_TOOL_MAJOR":83,"AXIS_TOOL_MINOR":75,"AXIS_ORIENTATION":1.1489516496658325},"buttonState":[]},
+      {"action":"MOVE","axes":{"AXIS_X":271,"AXIS_Y":603,"AXIS_PRESSURE":0.234375,"AXIS_SIZE":0.017389604821801186,"AXIS_TOUCH_MAJOR":60,"AXIS_TOUCH_MINOR":29,"AXIS_TOOL_MAJOR":60,"AXIS_TOOL_MINOR":29,"AXIS_ORIENTATION":1.0787720680236816},"buttonState":[]},
+      {"action":"UP","axes":{"AXIS_X":271,"AXIS_Y":603,"AXIS_PRESSURE":0.234375,"AXIS_SIZE":0.017389604821801186,"AXIS_TOUCH_MAJOR":60,"AXIS_TOUCH_MINOR":29,"AXIS_TOOL_MAJOR":60,"AXIS_TOOL_MINOR":29,"AXIS_ORIENTATION":1.0787720680236816},"buttonState":[]}
+    ]
+  }
+]
diff --git a/tests/Input/src/com/android/test/input/AnrTest.kt b/tests/Input/src/com/android/test/input/AnrTest.kt
index d32cedb..cd6ab30 100644
--- a/tests/Input/src/com/android/test/input/AnrTest.kt
+++ b/tests/Input/src/com/android/test/input/AnrTest.kt
@@ -166,12 +166,12 @@
         val displayManager =
             instrumentation.context.getSystemService(Context.DISPLAY_SERVICE) as DisplayManager
         val display = displayManager.getDisplay(obj.getDisplayId())
-        val touchScreen = UinputTouchScreen(instrumentation, display)
-
         val rect: Rect = obj.visibleBounds
-        val pointer = touchScreen.touchDown(rect.centerX(), rect.centerY())
-        pointer.lift()
-        touchScreen.close()
+        UinputTouchScreen(instrumentation, display).use { touchScreen ->
+            touchScreen
+                .touchDown(rect.centerX(), rect.centerY())
+                .lift()
+        }
     }
 
     private fun triggerAnr() {
diff --git a/tests/Input/src/com/android/test/input/CaptureEventActivity.kt b/tests/Input/src/com/android/test/input/CaptureEventActivity.kt
new file mode 100644
index 0000000..d54e3470
--- /dev/null
+++ b/tests/Input/src/com/android/test/input/CaptureEventActivity.kt
@@ -0,0 +1,85 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.test.input
+
+import android.app.Activity
+import android.os.Bundle
+import android.view.InputEvent
+import android.view.KeyEvent
+import android.view.MotionEvent
+import java.util.concurrent.LinkedBlockingQueue
+import java.util.concurrent.TimeUnit
+import org.junit.Assert.assertNull
+
+class CaptureEventActivity : Activity() {
+    private val events = LinkedBlockingQueue<InputEvent>()
+    var shouldHandleKeyEvents = true
+
+    override fun onCreate(savedInstanceState: Bundle?) {
+        super.onCreate(savedInstanceState)
+
+        // Set the fixed orientation if requested
+        if (intent.hasExtra(EXTRA_FIXED_ORIENTATION)) {
+            val orientation = intent.getIntExtra(EXTRA_FIXED_ORIENTATION, 0)
+            setRequestedOrientation(orientation)
+        }
+
+        // Set the flag if requested
+        if (intent.hasExtra(EXTRA_WINDOW_FLAGS)) {
+            val flags = intent.getIntExtra(EXTRA_WINDOW_FLAGS, 0)
+            window.addFlags(flags)
+        }
+    }
+
+    override fun dispatchGenericMotionEvent(ev: MotionEvent?): Boolean {
+        events.add(MotionEvent.obtain(ev))
+        return true
+    }
+
+    override fun dispatchTouchEvent(ev: MotionEvent?): Boolean {
+        events.add(MotionEvent.obtain(ev))
+        return true
+    }
+
+    override fun dispatchKeyEvent(event: KeyEvent?): Boolean {
+        events.add(KeyEvent(event))
+        return shouldHandleKeyEvents
+    }
+
+    override fun dispatchTrackballEvent(ev: MotionEvent?): Boolean {
+        events.add(MotionEvent.obtain(ev))
+        return true
+    }
+
+    fun getInputEvent(): InputEvent? {
+        return events.poll(5, TimeUnit.SECONDS)
+    }
+
+    fun hasReceivedEvents(): Boolean {
+        return !events.isEmpty()
+    }
+
+    fun assertNoEvents() {
+        val event = events.poll(100, TimeUnit.MILLISECONDS)
+        assertNull("Expected no events, but received $event", event)
+    }
+
+    companion object {
+        const val EXTRA_FIXED_ORIENTATION = "fixed_orientation"
+        const val EXTRA_WINDOW_FLAGS = "window_flags"
+    }
+}
diff --git a/tests/Input/src/com/android/test/input/UinputRecordingIntegrationTests.kt b/tests/Input/src/com/android/test/input/UinputRecordingIntegrationTests.kt
new file mode 100644
index 0000000..5d45c44
--- /dev/null
+++ b/tests/Input/src/com/android/test/input/UinputRecordingIntegrationTests.kt
@@ -0,0 +1,170 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.test.input
+
+import android.Manifest.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY
+import android.app.Instrumentation
+import android.cts.input.EventVerifier
+import android.graphics.PointF
+import android.hardware.input.InputManager
+import android.os.ParcelFileDescriptor
+import android.util.Size
+import android.view.MotionEvent
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.cts.input.BatchedEventSplitter
+import com.android.cts.input.InputJsonParser
+import com.android.cts.input.VirtualDisplayActivityScenario
+import com.android.cts.input.inputeventmatchers.isResampled
+import com.android.cts.input.inputeventmatchers.withButtonState
+import com.android.cts.input.inputeventmatchers.withHistorySize
+import com.android.cts.input.inputeventmatchers.withMotionAction
+import com.android.cts.input.inputeventmatchers.withPressure
+import com.android.cts.input.inputeventmatchers.withRawCoords
+import com.android.cts.input.inputeventmatchers.withSource
+import java.io.InputStream
+import org.hamcrest.Matchers.allOf
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TestName
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+
+/**
+ * Integration tests for the input pipeline that replays recording taken from physical input devices
+ * at the evdev interface level, and makes assertions on the events that are received by a test app.
+ *
+ * These tests associate the playback input device with a virtual display to make these tests
+ * agnostic to the device form factor.
+ *
+ * New recordings can be taken using the `evemu-record` shell command.
+ */
+@RunWith(Parameterized::class)
+class UinputRecordingIntegrationTests {
+
+    companion object {
+        /**
+         * Add new test cases by adding a new [TestData] to the following list.
+         */
+        @JvmStatic
+        @Parameterized.Parameters(name = "{0}")
+        fun data(): Iterable<Any> =
+            listOf(
+                TestData(
+                    "GooglePixelTabletTouchscreen", R.raw.google_pixel_tablet_touchscreen,
+                    R.raw.google_pixel_tablet_touchscreen_events, Size(1600, 2560),
+                    vendorId = 0x0603, productId = 0x7806
+                ),
+            )
+
+        const val INPUT_DEVICE_SOURCE_ALL = -1
+    }
+
+    class TestData(
+        val name: String,
+        val uinputRecordingResource: Int,
+        val expectedEventsResource: Int,
+        val displaySize: Size,
+        val vendorId: Int,
+        val productId: Int,
+    ) {
+        override fun toString(): String = name
+    }
+
+    private lateinit var instrumentation: Instrumentation
+    private lateinit var parser: InputJsonParser
+
+    @get:Rule
+    val testName = TestName()
+
+    @Parameterized.Parameter(0)
+    lateinit var testData: TestData
+
+    @Before
+    fun setUp() {
+        instrumentation = InstrumentationRegistry.getInstrumentation()
+        parser = InputJsonParser(instrumentation.context)
+    }
+
+    @Test
+    fun testEvemuRecording() {
+        VirtualDisplayActivityScenario.AutoClose<CaptureEventActivity>(
+            testName,
+            size = testData.displaySize
+        ).use { scenario ->
+            scenario.activity.window.decorView.requestUnbufferedDispatch(INPUT_DEVICE_SOURCE_ALL)
+
+            try {
+                instrumentation.uiAutomation.adoptShellPermissionIdentity(
+                    ASSOCIATE_INPUT_DEVICE_TO_DISPLAY,
+                )
+
+                val inputPort = "uinput:1:${testData.vendorId}:${testData.productId}"
+                val inputManager =
+                    instrumentation.context.getSystemService(InputManager::class.java)!!
+                try {
+                    inputManager.addUniqueIdAssociationByPort(
+                        inputPort,
+                        scenario.virtualDisplay.display.uniqueId!!,
+                    )
+
+                    val verifier =
+                        EventVerifier(BatchedEventSplitter { scenario.activity.getInputEvent() })
+                    injectUinputEvents()
+                    verifyEvents(verifier)
+                    scenario.activity.assertNoEvents()
+                } finally {
+                    inputManager.removeUniqueIdAssociationByPort(inputPort)
+                }
+            } finally {
+                instrumentation.uiAutomation.dropShellPermissionIdentity()
+            }
+        }
+    }
+
+    private fun injectUinputEvents() {
+        val fds = instrumentation.uiAutomation!!.executeShellCommandRw("uinput -")
+
+        ParcelFileDescriptor.AutoCloseOutputStream(fds[1]).use { stdIn ->
+            val inputStream: InputStream = instrumentation.context.resources.openRawResource(
+                testData.uinputRecordingResource,
+            )
+            stdIn.write(inputStream.readBytes())
+        }
+    }
+
+    private fun verifyEvents(verifier: EventVerifier) {
+        val uinputTestData = parser.getUinputTestData(testData.expectedEventsResource)
+        for (test in uinputTestData) {
+            for ((index, expectedEvent) in test.events.withIndex()) {
+                if (expectedEvent is MotionEvent) {
+                    verifier.assertReceivedMotion(
+                        allOf(
+                            withMotionAction(expectedEvent.action),
+                            withSource(expectedEvent.source),
+                            withButtonState(expectedEvent.buttonState),
+                            withRawCoords(PointF(expectedEvent.rawX, expectedEvent.rawY)),
+                            withPressure(expectedEvent.pressure),
+                            isResampled(false),
+                            withHistorySize(0),
+                        ),
+                        "${test.name}: Expected event at index $index",
+                    )
+                }
+            }
+        }
+    }
+}
diff --git a/tests/UsbManagerTests/Android.bp b/tests/UsbManagerTests/Android.bp
index 2909e66..331a21a 100644
--- a/tests/UsbManagerTests/Android.bp
+++ b/tests/UsbManagerTests/Android.bp
@@ -44,7 +44,7 @@
         "libstaticjvmtiagent",
     ],
     libs: [
-        "android.test.mock",
+        "android.test.mock.stubs.system",
     ],
     certificate: "platform",
     platform_apis: true,
diff --git a/tools/systemfeatures/Android.bp b/tools/systemfeatures/Android.bp
index aca25eb..a9e6328 100644
--- a/tools/systemfeatures/Android.bp
+++ b/tools/systemfeatures/Android.bp
@@ -5,6 +5,7 @@
     // to get the below license kinds:
     //   SPDX-license-identifier-Apache-2.0
     default_applicable_licenses: ["frameworks_base_license"],
+    default_team: "trendy_team_system_performance",
 }
 
 java_library_host {
@@ -25,8 +26,6 @@
     static_libs: ["systemfeatures-gen-lib"],
 }
 
-// TODO(b/203143243): Add golden diff test for generated sources.
-// Functional runtime behavior is covered in systemfeatures-gen-tests.
 genrule {
     name: "systemfeatures-gen-tests-srcs",
     cmd: "$(location systemfeatures-gen-tool) com.android.systemfeatures.RwNoFeatures --readonly=false > $(location RwNoFeatures.java) && " +
@@ -42,11 +41,12 @@
     tools: ["systemfeatures-gen-tool"],
 }
 
+// Functional runtime behavior testing.
 java_test_host {
     name: "systemfeatures-gen-tests",
     test_suites: ["general-tests"],
     srcs: [
-        "tests/**/*.java",
+        "tests/src/**/*.java",
         ":systemfeatures-gen-tests-srcs",
     ],
     test_options: {
@@ -61,3 +61,33 @@
         "truth",
     ],
 }
+
+// Rename the goldens as they may be copied into the source tree, and we don't
+// need or want the usual `.java` linting (e.g., copyright checks).
+genrule {
+    name: "systemfeatures-gen-tests-golden-srcs",
+    cmd: "for f in $(in); do cp $$f $(genDir)/tests/gen/$$(basename $$f).gen; done",
+    srcs: [":systemfeatures-gen-tests-srcs"],
+    out: [
+        "tests/gen/RwNoFeatures.java.gen",
+        "tests/gen/RoNoFeatures.java.gen",
+        "tests/gen/RwFeatures.java.gen",
+        "tests/gen/RoFeatures.java.gen",
+    ],
+}
+
+// Golden output testing. Golden sources can be updated via:
+//   $ANDROID_BUILD_TOP/frameworks/base/tools/systemfeatures/tests/golden_test.sh --update
+sh_test_host {
+    name: "systemfeatures-gen-golden-tests",
+    src: "tests/golden_test.sh",
+    filename: "systemfeatures-gen-golden-tests.sh",
+    test_config: "tests/systemfeatures-gen-golden-tests.xml",
+    data: [
+        "tests/golden/**/*.java*",
+        ":systemfeatures-gen-tests-golden-srcs",
+    ],
+    test_options: {
+        unit_test: true,
+    },
+}
diff --git a/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesGenerator.kt b/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesGenerator.kt
index e537ffc..5df453d 100644
--- a/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesGenerator.kt
+++ b/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesGenerator.kt
@@ -142,6 +142,10 @@
 
         // TODO(b/203143243): Add validation of build vs runtime values to ensure consistency.
         JavaFile.builder(outputClassName.packageName(), classBuilder.build())
+            .indent("    ")
+            .skipJavaLangImports(true)
+            .addFileComment("This file is auto-generated. DO NOT MODIFY.\n")
+            .addFileComment("Args: ${args.joinToString(" \\\n           ")}")
             .build()
             .writeTo(System.out)
     }
@@ -178,6 +182,7 @@
             val methodBuilder =
                 MethodSpec.methodBuilder(methodName)
                     .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                    .addJavadoc("Check for ${feature.name}.\n\n@hide")
                     .returns(Boolean::class.java)
                     .addParameter(CONTEXT_CLASS, "context")
 
@@ -228,6 +233,7 @@
             MethodSpec.methodBuilder("maybeHasFeature")
                 .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                 .addAnnotation(ClassName.get("android.annotation", "Nullable"))
+                .addJavadoc("@hide")
                 .returns(Boolean::class.javaObjectType) // Use object type for nullability
                 .addParameter(String::class.java, "featureName")
                 .addParameter(Int::class.java, "version")
diff --git a/tools/systemfeatures/tests/golden/RoFeatures.java.gen b/tools/systemfeatures/tests/golden/RoFeatures.java.gen
new file mode 100644
index 0000000..724639b
--- /dev/null
+++ b/tools/systemfeatures/tests/golden/RoFeatures.java.gen
@@ -0,0 +1,88 @@
+// This file is auto-generated. DO NOT MODIFY.
+// Args: com.android.systemfeatures.RoFeatures \
+//            --readonly=true \
+//            --feature=WATCH:1 \
+//            --feature=WIFI:0 \
+//            --feature=VULKAN:-1 \
+//            --feature=AUTO: \
+//            --feature-apis=WATCH,PC
+package com.android.systemfeatures;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import com.android.aconfig.annotations.AssumeFalseForR8;
+import com.android.aconfig.annotations.AssumeTrueForR8;
+
+/**
+ * @hide
+ */
+public final class RoFeatures {
+    /**
+     * Check for FEATURE_WATCH.
+     *
+     * @hide
+     */
+    @AssumeTrueForR8
+    public static boolean hasFeatureWatch(Context context) {
+        return true;
+    }
+
+    /**
+     * Check for FEATURE_PC.
+     *
+     * @hide
+     */
+    public static boolean hasFeaturePc(Context context) {
+        return hasFeatureFallback(context, PackageManager.FEATURE_PC);
+    }
+
+    /**
+     * Check for FEATURE_WIFI.
+     *
+     * @hide
+     */
+    @AssumeTrueForR8
+    public static boolean hasFeatureWifi(Context context) {
+        return true;
+    }
+
+    /**
+     * Check for FEATURE_VULKAN.
+     *
+     * @hide
+     */
+    @AssumeFalseForR8
+    public static boolean hasFeatureVulkan(Context context) {
+        return false;
+    }
+
+    /**
+     * Check for FEATURE_AUTO.
+     *
+     * @hide
+     */
+    @AssumeFalseForR8
+    public static boolean hasFeatureAuto(Context context) {
+        return false;
+    }
+
+    private static boolean hasFeatureFallback(Context context, String featureName) {
+        return context.getPackageManager().hasSystemFeature(featureName, 0);
+    }
+
+    /**
+     * @hide
+     */
+    @Nullable
+    public static Boolean maybeHasFeature(String featureName, int version) {
+        switch (featureName) {
+            case PackageManager.FEATURE_WATCH: return 1 >= version;
+            case PackageManager.FEATURE_WIFI: return 0 >= version;
+            case PackageManager.FEATURE_VULKAN: return -1 >= version;
+            case PackageManager.FEATURE_AUTO: return false;
+            default: break;
+        }
+        return null;
+    }
+}
diff --git a/tools/systemfeatures/tests/golden/RoNoFeatures.java.gen b/tools/systemfeatures/tests/golden/RoNoFeatures.java.gen
new file mode 100644
index 0000000..59c5b4e
--- /dev/null
+++ b/tools/systemfeatures/tests/golden/RoNoFeatures.java.gen
@@ -0,0 +1,35 @@
+// This file is auto-generated. DO NOT MODIFY.
+// Args: com.android.systemfeatures.RoNoFeatures \
+//            --readonly=true \
+//            --feature-apis=WATCH
+package com.android.systemfeatures;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.pm.PackageManager;
+
+/**
+ * @hide
+ */
+public final class RoNoFeatures {
+    /**
+     * Check for FEATURE_WATCH.
+     *
+     * @hide
+     */
+    public static boolean hasFeatureWatch(Context context) {
+        return hasFeatureFallback(context, PackageManager.FEATURE_WATCH);
+    }
+
+    private static boolean hasFeatureFallback(Context context, String featureName) {
+        return context.getPackageManager().hasSystemFeature(featureName, 0);
+    }
+
+    /**
+     * @hide
+     */
+    @Nullable
+    public static Boolean maybeHasFeature(String featureName, int version) {
+        return null;
+    }
+}
diff --git a/tools/systemfeatures/tests/golden/RwFeatures.java.gen b/tools/systemfeatures/tests/golden/RwFeatures.java.gen
new file mode 100644
index 0000000..6f89759
--- /dev/null
+++ b/tools/systemfeatures/tests/golden/RwFeatures.java.gen
@@ -0,0 +1,65 @@
+// This file is auto-generated. DO NOT MODIFY.
+// Args: com.android.systemfeatures.RwFeatures \
+//            --readonly=false \
+//            --feature=WATCH:1 \
+//            --feature=WIFI:0 \
+//            --feature=VULKAN:-1 \
+//            --feature=AUTO:
+package com.android.systemfeatures;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.pm.PackageManager;
+
+/**
+ * @hide
+ */
+public final class RwFeatures {
+    /**
+     * Check for FEATURE_WATCH.
+     *
+     * @hide
+     */
+    public static boolean hasFeatureWatch(Context context) {
+        return hasFeatureFallback(context, PackageManager.FEATURE_WATCH);
+    }
+
+    /**
+     * Check for FEATURE_WIFI.
+     *
+     * @hide
+     */
+    public static boolean hasFeatureWifi(Context context) {
+        return hasFeatureFallback(context, PackageManager.FEATURE_WIFI);
+    }
+
+    /**
+     * Check for FEATURE_VULKAN.
+     *
+     * @hide
+     */
+    public static boolean hasFeatureVulkan(Context context) {
+        return hasFeatureFallback(context, PackageManager.FEATURE_VULKAN);
+    }
+
+    /**
+     * Check for FEATURE_AUTO.
+     *
+     * @hide
+     */
+    public static boolean hasFeatureAuto(Context context) {
+        return hasFeatureFallback(context, PackageManager.FEATURE_AUTO);
+    }
+
+    private static boolean hasFeatureFallback(Context context, String featureName) {
+        return context.getPackageManager().hasSystemFeature(featureName, 0);
+    }
+
+    /**
+     * @hide
+     */
+    @Nullable
+    public static Boolean maybeHasFeature(String featureName, int version) {
+        return null;
+    }
+}
diff --git a/tools/systemfeatures/tests/golden/RwNoFeatures.java.gen b/tools/systemfeatures/tests/golden/RwNoFeatures.java.gen
new file mode 100644
index 0000000..2111d56
--- /dev/null
+++ b/tools/systemfeatures/tests/golden/RwNoFeatures.java.gen
@@ -0,0 +1,24 @@
+// This file is auto-generated. DO NOT MODIFY.
+// Args: com.android.systemfeatures.RwNoFeatures \
+//            --readonly=false
+package com.android.systemfeatures;
+
+import android.annotation.Nullable;
+import android.content.Context;
+
+/**
+ * @hide
+ */
+public final class RwNoFeatures {
+    private static boolean hasFeatureFallback(Context context, String featureName) {
+        return context.getPackageManager().hasSystemFeature(featureName, 0);
+    }
+
+    /**
+     * @hide
+     */
+    @Nullable
+    public static Boolean maybeHasFeature(String featureName, int version) {
+        return null;
+    }
+}
diff --git a/tools/systemfeatures/tests/golden_test.sh b/tools/systemfeatures/tests/golden_test.sh
new file mode 100755
index 0000000..c249254
--- /dev/null
+++ b/tools/systemfeatures/tests/golden_test.sh
@@ -0,0 +1,52 @@
+#!/usr/bin/env bash
+
+# Copyright (C) 2024 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -e
+
+GEN_DIR="tests/gen"
+GOLDEN_DIR="tests/golden"
+
+if [[ $(basename $0) == "golden_test.sh" ]]; then
+  # We're running via command-line, so we need to:
+  #   1) manually update generated srcs
+  #   2) use absolute paths
+  if [ -z $ANDROID_BUILD_TOP ]; then
+    echo "You need to source and lunch before you can use this script directly."
+    exit 1
+  fi
+  GEN_DIR="$ANDROID_BUILD_TOP/out/soong/.intermediates/frameworks/base/tools/systemfeatures/systemfeatures-gen-tests-golden-srcs/gen/$GEN_DIR"
+  GOLDEN_DIR="$ANDROID_BUILD_TOP/frameworks/base/tools/systemfeatures/$GOLDEN_DIR"
+  rm -rf "$GEN_DIR"
+  "$ANDROID_BUILD_TOP"/build/soong/soong_ui.bash --make-mode systemfeatures-gen-tests-golden-srcs
+fi
+
+if [[ "$1" == "--update" ]]; then
+  rm -rf "$GOLDEN_DIR"
+  cp -R "$GEN_DIR" "$GOLDEN_DIR"
+  echo "Updated golden test files."
+else
+  echo "Running diff from test output against golden test files..."
+  if diff -ruN "$GOLDEN_DIR" "$GEN_DIR" ; then
+    echo "No changes."
+  else
+    echo
+    echo "----------------------------------------------------------------------------------------"
+    echo "If changes look OK, run:"
+    echo "  \$ANDROID_BUILD_TOP/frameworks/base/tools/systemfeatures/tests/golden_test.sh --update"
+    echo "----------------------------------------------------------------------------------------"
+    exit 1
+  fi
+fi
diff --git a/tools/systemfeatures/tests/Context.java b/tools/systemfeatures/tests/src/Context.java
similarity index 100%
rename from tools/systemfeatures/tests/Context.java
rename to tools/systemfeatures/tests/src/Context.java
diff --git a/tools/systemfeatures/tests/PackageManager.java b/tools/systemfeatures/tests/src/PackageManager.java
similarity index 100%
rename from tools/systemfeatures/tests/PackageManager.java
rename to tools/systemfeatures/tests/src/PackageManager.java
diff --git a/tools/systemfeatures/tests/SystemFeaturesGeneratorTest.java b/tools/systemfeatures/tests/src/SystemFeaturesGeneratorTest.java
similarity index 100%
rename from tools/systemfeatures/tests/SystemFeaturesGeneratorTest.java
rename to tools/systemfeatures/tests/src/SystemFeaturesGeneratorTest.java
diff --git a/tools/systemfeatures/tests/systemfeatures-gen-golden-tests.xml b/tools/systemfeatures/tests/systemfeatures-gen-golden-tests.xml
new file mode 100644
index 0000000..e3a5841
--- /dev/null
+++ b/tools/systemfeatures/tests/systemfeatures-gen-golden-tests.xml
@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration description="Runs systemfeatures-gen golden diff test">
+    <test class="com.android.tradefed.testtype.binary.ExecutableHostTest" >
+      <option name="binary" value="systemfeatures-gen-golden-tests.sh"/>
+    </test>
+</configuration>