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/java/android/os/IpcDataCache.java b/core/java/android/os/IpcDataCache.java
index bf44d65..0776cf4 100644
--- a/core/java/android/os/IpcDataCache.java
+++ b/core/java/android/os/IpcDataCache.java
@@ -16,6 +16,7 @@
 
 package android.os;
 
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.StringDef;
@@ -551,7 +552,7 @@
     }
 
     /**
-     * An interface suitable for a lambda expression instead of a QueryHandler.
+     * An interface suitable for a lambda expression instead of a QueryHandler applying remote call.
      * @hide
      */
     public interface RemoteCall<Query, Result> {
@@ -559,6 +560,14 @@
     }
 
     /**
+     * An interface suitable for a lambda expression instead of a QueryHandler bypassing the cache.
+     * @hide
+     */
+    public interface BypassCall<Query> {
+        Boolean apply(Query query);
+    }
+
+    /**
      * This is a query handler that is created with a lambda expression that is invoked
      * every time the handler is called.  The handler is specifically meant for services
      * hosted by system_server; the handler automatically rethrows RemoteException as a
@@ -580,11 +589,54 @@
         }
     }
 
+
     /**
      * Create a cache using a config and a lambda expression.
+     * @param config The configuration for the cache.
+     * @param remoteCall The lambda expression that will be invoked to fetch the data.
      * @hide
      */
-    public IpcDataCache(@NonNull Config config, @NonNull RemoteCall<Query, Result> computer) {
-        this(config, new SystemServerCallHandler<>(computer));
+    public IpcDataCache(@NonNull Config config, @NonNull RemoteCall<Query, Result> remoteCall) {
+      this(config, android.multiuser.Flags.cachingDevelopmentImprovements() ?
+        new QueryHandler<Query, Result>() {
+            @Override
+            public Result apply(Query query) {
+                try {
+                    return remoteCall.apply(query);
+                } catch (RemoteException e) {
+                    throw e.rethrowFromSystemServer();
+                }
+            }
+        } : new SystemServerCallHandler<>(remoteCall));
+    }
+
+
+    /**
+     * Create a cache using a config and a lambda expression.
+     * @param config The configuration for the cache.
+     * @param remoteCall The lambda expression that will be invoked to fetch the data.
+     * @param bypass The lambda expression that will be invoked to determine if the cache should be
+     *     bypassed.
+     * @hide
+     */
+    @FlaggedApi(android.multiuser.Flags.FLAG_CACHING_DEVELOPMENT_IMPROVEMENTS)
+    public IpcDataCache(@NonNull Config config,
+            @NonNull RemoteCall<Query, Result> remoteCall,
+            @NonNull BypassCall<Query> bypass) {
+        this(config, new QueryHandler<Query, Result>() {
+            @Override
+            public Result apply(Query query) {
+                try {
+                    return remoteCall.apply(query);
+                } catch (RemoteException e) {
+                    throw e.rethrowFromSystemServer();
+                }
+            }
+
+            @Override
+            public boolean shouldBypassCache(Query query) {
+                return bypass.apply(query);
+            }
+        });
     }
 }
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/window/OnBackInvokedDispatcher.java b/core/java/android/window/OnBackInvokedDispatcher.java
index bccee92..0632a37 100644
--- a/core/java/android/window/OnBackInvokedDispatcher.java
+++ b/core/java/android/window/OnBackInvokedDispatcher.java
@@ -76,7 +76,7 @@
      * @param callback The callback to be registered. If the callback instance has been already
      *                 registered, the existing instance (no matter its priority) will be
      *                 unregistered and registered again.
-     * @throws {@link IllegalArgumentException} if the priority is negative.
+     * @throws IllegalArgumentException if the priority is negative.
      */
     @SuppressLint({"ExecutorRegistration"})
     void registerOnBackInvokedCallback(
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index 03e230c..69b91fd 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -144,6 +144,13 @@
 }
 
 flag {
+  name: "universal_resizable_by_default"
+  namespace: "windowing_frontend"
+  description: "The orientation, aspect ratio, resizability of activity will follow system behavior by default"
+  bug: "357141415"
+}
+
+flag {
   name: "respect_non_top_visible_fixed_orientation"
   namespace: "windowing_frontend"
   description: "If top activity is not opaque, respect the fixed orientation of activity behind it"
diff --git a/core/tests/coretests/src/android/os/IpcDataCacheTest.java b/core/tests/coretests/src/android/os/IpcDataCacheTest.java
index b03fd64..64f77b3 100644
--- a/core/tests/coretests/src/android/os/IpcDataCacheTest.java
+++ b/core/tests/coretests/src/android/os/IpcDataCacheTest.java
@@ -18,7 +18,9 @@
 
 import static org.junit.Assert.assertEquals;
 
+import android.multiuser.Flags;
 import android.platform.test.annotations.IgnoreUnderRavenwood;
+import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.ravenwood.RavenwoodRule;
 
 import androidx.test.filters.SmallTest;
@@ -151,8 +153,6 @@
         tester.verify(9);
     }
 
-    // This test is disabled pending an sepolicy change that allows any app to set the
-    // test property.
     @Test
     public void testRemoteCall() {
 
@@ -193,6 +193,44 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_CACHING_DEVELOPMENT_IMPROVEMENTS)
+    public void testRemoteCallBypass() {
+
+        // A stand-in for the binder.  The test verifies that calls are passed through to
+        // this class properly.
+        ServerProxy tester = new ServerProxy();
+
+        // Create a cache that uses simple arithmetic to computer its values.
+        IpcDataCache.Config config = new IpcDataCache.Config(4, MODULE, API, "testCache3");
+        IpcDataCache<Integer, Boolean> testCache =
+                new IpcDataCache<>(config, (x) -> tester.query(x), (x) -> x % 9 == 0);
+
+        IpcDataCache.setTestMode(true);
+        testCache.testPropertyName();
+
+        tester.verify(0);
+        assertEquals(tester.value(3), testCache.query(3));
+        tester.verify(1);
+        assertEquals(tester.value(3), testCache.query(3));
+        tester.verify(2);
+        testCache.invalidateCache();
+        assertEquals(tester.value(3), testCache.query(3));
+        tester.verify(3);
+        assertEquals(tester.value(5), testCache.query(5));
+        tester.verify(4);
+        assertEquals(tester.value(5), testCache.query(5));
+        tester.verify(4);
+        assertEquals(tester.value(3), testCache.query(3));
+        tester.verify(4);
+        assertEquals(tester.value(9), testCache.query(9));
+        tester.verify(5);
+        assertEquals(tester.value(3), testCache.query(3));
+        tester.verify(5);
+        assertEquals(tester.value(5), testCache.query(5));
+        tester.verify(5);
+    }
+
+    @Test
     public void testDisableCache() {
 
         // A stand-in for the binder.  The test verifies that calls are passed through to
diff --git a/libs/appfunctions/OWNERS b/libs/appfunctions/OWNERS
new file mode 100644
index 0000000..c093675
--- /dev/null
+++ b/libs/appfunctions/OWNERS
@@ -0,0 +1,3 @@
+set noparent
+
+include /core/java/android/app/appfunctions/OWNERS
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/Graph/Android.bp b/packages/SettingsLib/Graph/Android.bp
index e2ed1e4..163b689 100644
--- a/packages/SettingsLib/Graph/Android.bp
+++ b/packages/SettingsLib/Graph/Android.bp
@@ -4,7 +4,7 @@
 
 filegroup {
     name: "SettingsLibGraph-srcs",
-    srcs: ["src/**/*"],
+    srcs: ["src/**/*.kt"],
 }
 
 android_library {
@@ -14,8 +14,24 @@
     ],
     srcs: [":SettingsLibGraph-srcs"],
     static_libs: [
+        "SettingsLibGraph-proto-lite",
+        "SettingsLibIpc",
+        "SettingsLibMetadata",
+        "SettingsLibPreference",
         "androidx.annotation_annotation",
+        "androidx.fragment_fragment",
         "androidx.preference_preference",
     ],
     kotlincflags: ["-Xjvm-default=all"],
 }
+
+java_library {
+    name: "SettingsLibGraph-proto-lite",
+    srcs: ["graph.proto"],
+    proto: {
+        type: "lite",
+        canonical_path_from_root: false,
+    },
+    sdk_version: "core_current",
+    static_libs: ["libprotobuf-java-lite"],
+}
diff --git a/packages/SettingsLib/Graph/graph.proto b/packages/SettingsLib/Graph/graph.proto
new file mode 100644
index 0000000..e93d756
--- /dev/null
+++ b/packages/SettingsLib/Graph/graph.proto
@@ -0,0 +1,156 @@
+syntax = "proto3";
+
+package com.android.settingslib.graph;
+
+option java_package = "com.android.settingslib.graph.proto";
+option java_multiple_files = true;
+
+// Proto represents preference graph.
+message PreferenceGraphProto {
+  // Preference screens appear in the graph.
+  // Key: preference key of the PreferenceScreen. Value: PreferenceScreen.
+  map<string, PreferenceScreenProto> screens = 1;
+  // Roots of the graph.
+  // Each element is a preference key of the PreferenceScreen.
+  repeated string roots = 2;
+  // Activities appear in the graph.
+  // Key: activity class. Value: preference key of associated PreferenceScreen.
+  map<string, string> activity_screens = 3;
+}
+
+// Proto of PreferenceScreen.
+message PreferenceScreenProto {
+  // Intent to show the PreferenceScreen.
+  optional IntentProto intent = 1;
+  // Root of the PreferenceScreen hierarchy.
+  optional PreferenceGroupProto root = 2;
+  // If the preference screen provides complete hierarchy by source code.
+  optional bool complete_hierarchy = 3;
+}
+
+// Proto of PreferenceGroup.
+message PreferenceGroupProto {
+  // Self information of PreferenceGroup.
+  optional PreferenceProto preference = 1;
+  // A list of children.
+  repeated PreferenceOrGroupProto preferences = 2;
+}
+
+// Proto represents either PreferenceProto or PreferenceGroupProto.
+message PreferenceOrGroupProto {
+  oneof kind {
+    // It is a Preference.
+    PreferenceProto preference = 1;
+    // It is a PreferenceGroup.
+    PreferenceGroupProto group = 2;
+  }
+}
+
+// Proto of Preference.
+message PreferenceProto {
+  // Key of the preference.
+  optional string key = 1;
+  // Title of the preference.
+  optional TextProto title = 2;
+  // Summary of the preference.
+  optional TextProto summary = 3;
+  // Icon of the preference.
+  optional int32 icon = 4;
+  // Additional keywords for indexing.
+  optional int32 keywords = 5;
+  // Extras of the preference.
+  optional BundleProto extras = 6;
+  // Whether the preference is indexable.
+  optional bool indexable = 7;
+  // Whether the preference is enabled.
+  optional bool enabled = 8;
+  // Whether the preference is available/visible.
+  optional bool available = 9;
+  // Whether the preference is persistent.
+  optional bool persistent = 10;
+  // Whether the preference is restricted by managed configurations.
+  optional bool restricted = 11;
+  // Target of the preference action.
+  optional ActionTarget action_target = 12;
+  // Preference value (if present, it means `persistent` is true).
+  optional PreferenceValueProto value = 13;
+
+  // Target of an Intent
+  message ActionTarget {
+    oneof kind {
+      // Resolved key of the preference screen located in current app.
+      // This is resolved from android:fragment or activity of current app.
+      string key = 1;
+      // Unresolvable Intent that is either an unrecognized activity of current
+      // app or activity belongs to other app.
+      IntentProto intent = 2;
+    }
+  }
+}
+
+// Proto of string or string resource id.
+message TextProto {
+  oneof text {
+    int32 resource_id = 1;
+    string string = 2;
+  }
+}
+
+// Proto of preference value.
+message PreferenceValueProto {
+  oneof value {
+    bool boolean_value = 1;
+  }
+}
+
+// Proto of android.content.Intent
+message IntentProto {
+  // The action of the Intent.
+  optional string action = 1;
+
+  // The data attribute of the Intent, expressed as a URI.
+  optional string data = 2;
+
+  // The package attribute of the Intent, which may be set to force the
+  // detection of a particular application package that can handle the event.
+  optional string pkg = 3;
+
+  // The component attribute of the Intent, which may be set to force the
+  // detection of a particular component (app). If present, this must be a
+  // package name followed by a '/' and then followed by the class name.
+  optional string component = 4;
+
+  // Flags controlling how intent is handled. The value must be bitwise OR of
+  // intent flag constants defined by Android.
+  // http://developer.android.com/reference/android/content/Intent.html#setFlags(int)
+  optional int32 flags = 5;
+
+  // Extended data from the intent.
+  optional BundleProto extras = 6;
+
+  // The MIME type of the Intent (e.g. "text/plain").
+  //
+  // For more information, see
+  // https://developer.android.com/reference/android/content/Intent#setType(java.lang.String).
+  optional string mime_type = 7;
+}
+
+// Proto of android.os.Bundle
+message BundleProto {
+  // Bundle data.
+  map<string, BundleValue> values = 1;
+
+  message BundleValue {
+    // Bundle data value for the associated key name.
+    // Can be extended to support other types of bundled data.
+    oneof value {
+      string string_value = 1;
+      bytes bytes_value = 2;
+      int32 int_value = 3;
+      int64 long_value = 4;
+      bool boolean_value = 5;
+      double double_value = 6;
+      BundleProto bundle_value = 7;
+    }
+  }
+}
diff --git a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/GetPreferenceGraphApiHandler.kt b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/GetPreferenceGraphApiHandler.kt
new file mode 100644
index 0000000..04c2968
--- /dev/null
+++ b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/GetPreferenceGraphApiHandler.kt
@@ -0,0 +1,100 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.graph
+
+import android.app.Application
+import android.os.Bundle
+import com.android.settingslib.graph.proto.PreferenceGraphProto
+import com.android.settingslib.ipc.ApiHandler
+import com.android.settingslib.ipc.MessageCodec
+import java.util.Locale
+
+/** API to get preference graph. */
+abstract class GetPreferenceGraphApiHandler(private val activityClasses: Set<String>) :
+    ApiHandler<GetPreferenceGraphRequest, PreferenceGraphProto> {
+
+    override val requestCodec: MessageCodec<GetPreferenceGraphRequest>
+        get() = GetPreferenceGraphRequestCodec
+
+    override val responseCodec: MessageCodec<PreferenceGraphProto>
+        get() = PreferenceGraphProtoCodec
+
+    override suspend fun invoke(
+        application: Application,
+        myUid: Int,
+        callingUid: Int,
+        request: GetPreferenceGraphRequest,
+    ): PreferenceGraphProto {
+        val builderRequest =
+            if (request.activityClasses.isEmpty()) {
+                GetPreferenceGraphRequest(activityClasses, request.visitedScreens, request.locale)
+            } else {
+                request
+            }
+        return PreferenceGraphBuilder.of(application, builderRequest).build()
+    }
+}
+
+/**
+ * Request of [GetPreferenceGraphApiHandler].
+ *
+ * @param activityClasses activities of the preference graph
+ * @param visitedScreens keys of the visited preference screen
+ * @param locale locale of the preference graph
+ */
+data class GetPreferenceGraphRequest
+@JvmOverloads
+constructor(
+    val activityClasses: Set<String> = setOf(),
+    val visitedScreens: Set<String> = setOf(),
+    val locale: Locale? = null,
+    val includeValue: Boolean = true,
+)
+
+object GetPreferenceGraphRequestCodec : MessageCodec<GetPreferenceGraphRequest> {
+    override fun encode(data: GetPreferenceGraphRequest): Bundle =
+        Bundle(3).apply {
+            putStringArray(KEY_ACTIVITIES, data.activityClasses.toTypedArray())
+            putStringArray(KEY_PREF_KEYS, data.visitedScreens.toTypedArray())
+            putString(KEY_LOCALE, data.locale?.toLanguageTag())
+        }
+
+    override fun decode(data: Bundle): GetPreferenceGraphRequest {
+        val activities = data.getStringArray(KEY_ACTIVITIES) ?: arrayOf()
+        val visitedScreens = data.getStringArray(KEY_PREF_KEYS) ?: arrayOf()
+        fun String?.toLocale() = if (this != null) Locale.forLanguageTag(this) else null
+        return GetPreferenceGraphRequest(
+            activities.toSet(),
+            visitedScreens.toSet(),
+            data.getString(KEY_LOCALE).toLocale(),
+        )
+    }
+
+    private const val KEY_ACTIVITIES = "activities"
+    private const val KEY_PREF_KEYS = "keys"
+    private const val KEY_LOCALE = "locale"
+}
+
+object PreferenceGraphProtoCodec : MessageCodec<PreferenceGraphProto> {
+    override fun encode(data: PreferenceGraphProto): Bundle =
+        Bundle(1).apply { putByteArray(KEY_GRAPH, data.toByteArray()) }
+
+    override fun decode(data: Bundle): PreferenceGraphProto =
+        PreferenceGraphProto.parseFrom(data.getByteArray(KEY_GRAPH)!!)
+
+    private const val KEY_GRAPH = "graph"
+}
diff --git a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceGraphBuilder.kt b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceGraphBuilder.kt
new file mode 100644
index 0000000..8c5d877
--- /dev/null
+++ b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceGraphBuilder.kt
@@ -0,0 +1,445 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+@file:Suppress("DEPRECATION")
+
+package com.android.settingslib.graph
+
+import android.annotation.SuppressLint
+import android.app.Activity
+import android.content.Context
+import android.content.Intent
+import android.content.pm.PackageManager
+import android.content.res.Configuration
+import android.os.Build
+import android.os.Bundle
+import android.preference.PreferenceActivity
+import android.util.Log
+import androidx.fragment.app.Fragment
+import androidx.preference.Preference
+import androidx.preference.PreferenceGroup
+import androidx.preference.PreferenceScreen
+import androidx.preference.TwoStatePreference
+import com.android.settingslib.graph.proto.PreferenceGraphProto
+import com.android.settingslib.graph.proto.PreferenceGroupProto
+import com.android.settingslib.graph.proto.PreferenceProto
+import com.android.settingslib.graph.proto.PreferenceProto.ActionTarget
+import com.android.settingslib.graph.proto.PreferenceScreenProto
+import com.android.settingslib.graph.proto.TextProto
+import com.android.settingslib.metadata.BooleanValue
+import com.android.settingslib.metadata.PersistentPreference
+import com.android.settingslib.metadata.PreferenceAvailabilityProvider
+import com.android.settingslib.metadata.PreferenceHierarchy
+import com.android.settingslib.metadata.PreferenceHierarchyNode
+import com.android.settingslib.metadata.PreferenceMetadata
+import com.android.settingslib.metadata.PreferenceRestrictionProvider
+import com.android.settingslib.metadata.PreferenceScreenBindingKeyProvider
+import com.android.settingslib.metadata.PreferenceScreenMetadata
+import com.android.settingslib.metadata.PreferenceScreenRegistry
+import com.android.settingslib.metadata.PreferenceSummaryProvider
+import com.android.settingslib.metadata.PreferenceTitleProvider
+import com.android.settingslib.preference.PreferenceScreenFactory
+import com.android.settingslib.preference.PreferenceScreenProvider
+import java.util.Locale
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.withContext
+
+private const val TAG = "PreferenceGraphBuilder"
+
+/**
+ * Builder of preference graph.
+ *
+ * Only activity in current application is supported. To create preference graph across
+ * applications, use [crawlPreferenceGraph].
+ */
+class PreferenceGraphBuilder
+private constructor(private val context: Context, private val request: GetPreferenceGraphRequest) {
+    private val preferenceScreenFactory by lazy {
+        PreferenceScreenFactory(context.ofLocale(request.locale))
+    }
+    private val builder by lazy { PreferenceGraphProto.newBuilder() }
+    private val visitedScreens = mutableSetOf<String>().apply { addAll(request.visitedScreens) }
+    private val includeValue = request.includeValue
+
+    private suspend fun init() {
+        for (activityClass in request.activityClasses) {
+            add(activityClass)
+        }
+    }
+
+    fun build() = builder.build()
+
+    /** Adds an activity to the graph. */
+    suspend fun <T> add(activityClass: Class<T>) where T : Activity, T : PreferenceScreenProvider =
+        addPreferenceScreenProvider(activityClass)
+
+    /**
+     * Adds an activity to the graph.
+     *
+     * Reflection is used to create the instance. To avoid security vulnerability, the code ensures
+     * given [activityClassName] must be declared as an <activity> entry in AndroidManifest.xml.
+     */
+    suspend fun add(activityClassName: String) {
+        try {
+            val intent = Intent()
+            intent.setClassName(context, activityClassName)
+            if (context.packageManager.resolveActivity(intent, PackageManager.MATCH_DEFAULT_ONLY) ==
+                null) {
+                Log.e(TAG, "$activityClassName is not activity")
+                return
+            }
+            val activityClass = context.classLoader.loadClass(activityClassName)
+            if (addPreferenceScreenKeyProvider(activityClass)) return
+            if (PreferenceScreenProvider::class.java.isAssignableFrom(activityClass)) {
+                addPreferenceScreenProvider(activityClass)
+            } else {
+                Log.w(TAG, "$activityClass does not implement PreferenceScreenProvider")
+            }
+        } catch (e: Exception) {
+            Log.e(TAG, "Fail to add $activityClassName", e)
+        }
+    }
+
+    private suspend fun addPreferenceScreenKeyProvider(activityClass: Class<*>): Boolean {
+        if (!PreferenceScreenBindingKeyProvider::class.java.isAssignableFrom(activityClass)) {
+            return false
+        }
+        val key = getPreferenceScreenKey { activityClass.newInstance() } ?: return false
+        if (addPreferenceScreenFromRegistry(key, activityClass)) {
+            builder.addRoots(key)
+            return true
+        }
+        return false
+    }
+
+    private suspend fun getPreferenceScreenKey(newInstance: () -> Any): String? =
+        withContext(Dispatchers.Main) {
+            try {
+                val instance = newInstance()
+                if (instance is PreferenceScreenBindingKeyProvider) {
+                    return@withContext instance.getPreferenceScreenBindingKey(context)
+                } else {
+                    Log.w(TAG, "$instance is not PreferenceScreenKeyProvider")
+                }
+            } catch (e: Exception) {
+                Log.e(TAG, "getPreferenceScreenKey failed", e)
+            }
+            null
+        }
+
+    private suspend fun addPreferenceScreenFromRegistry(
+        key: String,
+        activityClass: Class<*>,
+    ): Boolean {
+        val metadata = PreferenceScreenRegistry[key] ?: return false
+        if (!metadata.hasCompleteHierarchy()) return false
+        return addPreferenceScreenMetadata(metadata, activityClass)
+    }
+
+    private suspend fun addPreferenceScreenMetadata(
+        metadata: PreferenceScreenMetadata,
+        activityClass: Class<*>,
+    ): Boolean =
+        addPreferenceScreen(metadata.key, activityClass) {
+            preferenceScreenProto {
+                completeHierarchy = true
+                root = metadata.getPreferenceHierarchy(context).toProto(activityClass, true)
+            }
+        }
+
+    private suspend fun addPreferenceScreenProvider(activityClass: Class<*>) {
+        Log.d(TAG, "add $activityClass")
+        createPreferenceScreen { activityClass.newInstance() }
+            ?.let {
+                addPreferenceScreen(Intent(context, activityClass), activityClass, it)
+                builder.addRoots(it.key)
+            }
+    }
+
+    /**
+     * Creates [PreferenceScreen].
+     *
+     * Androidx Activity/Fragment instance must be created in main thread, otherwise an exception is
+     * raised.
+     */
+    private suspend fun createPreferenceScreen(newInstance: () -> Any): PreferenceScreen? =
+        withContext(Dispatchers.Main) {
+            try {
+                val instance = newInstance()
+                Log.d(TAG, "createPreferenceScreen $instance")
+                if (instance is PreferenceScreenProvider) {
+                    return@withContext instance.createPreferenceScreen(preferenceScreenFactory)
+                } else {
+                    Log.w(TAG, "$instance is not PreferenceScreenProvider")
+                }
+            } catch (e: Exception) {
+                Log.e(TAG, "createPreferenceScreen failed", e)
+            }
+            return@withContext null
+        }
+
+    private suspend fun addPreferenceScreen(
+        intent: Intent,
+        activityClass: Class<*>,
+        preferenceScreen: PreferenceScreen?,
+    ) {
+        val key = preferenceScreen?.key
+        if (key.isNullOrEmpty()) {
+            Log.e(TAG, "$activityClass \"$preferenceScreen\" has no key")
+            return
+        }
+        @Suppress("CheckReturnValue")
+        addPreferenceScreen(key, activityClass) { preferenceScreen.toProto(intent, activityClass) }
+    }
+
+    private suspend fun addPreferenceScreen(
+        key: String,
+        activityClass: Class<*>,
+        preferenceScreenProvider: suspend () -> PreferenceScreenProto,
+    ): Boolean {
+        if (!visitedScreens.add(key)) {
+            Log.w(TAG, "$activityClass $key visited")
+            return false
+        }
+        val activityClassName = activityClass.name
+        val associatedKey = builder.getActivityScreensOrDefault(activityClassName, null)
+        if (associatedKey == null) {
+            builder.putActivityScreens(activityClassName, key)
+        } else if (associatedKey != key) {
+            Log.w(TAG, "Dup $activityClassName association, old: $associatedKey, new: $key")
+        }
+        builder.putScreens(key, preferenceScreenProvider())
+        return true
+    }
+
+    private suspend fun PreferenceScreen.toProto(
+        intent: Intent,
+        activityClass: Class<*>,
+    ): PreferenceScreenProto = preferenceScreenProto {
+        this.intent = intent.toProto()
+        root = (this@toProto as PreferenceGroup).toProto(activityClass)
+    }
+
+    private suspend fun PreferenceGroup.toProto(activityClass: Class<*>): PreferenceGroupProto =
+        preferenceGroupProto {
+            preference = (this@toProto as Preference).toProto(activityClass)
+            for (index in 0 until preferenceCount) {
+                val child = getPreference(index)
+                addPreferences(
+                    preferenceOrGroupProto {
+                        if (child is PreferenceGroup) {
+                            group = child.toProto(activityClass)
+                        } else {
+                            preference = child.toProto(activityClass)
+                        }
+                    })
+            }
+        }
+
+    private suspend fun Preference.toProto(activityClass: Class<*>): PreferenceProto =
+        preferenceProto {
+            this@toProto.key?.let { key = it }
+            this@toProto.title?.let { title = textProto { string = it.toString() } }
+            this@toProto.summary?.let { summary = textProto { string = it.toString() } }
+            val preferenceExtras = peekExtras()
+            preferenceExtras?.let { extras = it.toProto() }
+            enabled = isEnabled
+            available = isVisible
+            persistent = isPersistent
+            if (includeValue && isPersistent && this@toProto is TwoStatePreference) {
+                value = preferenceValueProto { booleanValue = this@toProto.isChecked }
+            }
+            this@toProto.fragment.toActionTarget(activityClass, preferenceExtras)?.let {
+                actionTarget = it
+                return@preferenceProto
+            }
+            this@toProto.intent?.let { actionTarget = it.toActionTarget() }
+        }
+
+    private suspend fun PreferenceHierarchy.toProto(
+        activityClass: Class<*>,
+        isRoot: Boolean,
+    ): PreferenceGroupProto = preferenceGroupProto {
+        preference = toProto(this@toProto, activityClass, isRoot)
+        forEachAsync {
+            addPreferences(
+                preferenceOrGroupProto {
+                    if (it is PreferenceHierarchy) {
+                        group = it.toProto(activityClass, false)
+                    } else {
+                        preference = toProto(it, activityClass, false)
+                    }
+                })
+        }
+    }
+
+    private suspend fun toProto(
+        node: PreferenceHierarchyNode,
+        activityClass: Class<*>,
+        isRoot: Boolean,
+    ) = preferenceProto {
+        val metadata = node.metadata
+        key = metadata.key
+        metadata.getTitleTextProto(isRoot)?.let { title = it }
+        if (metadata.summary != 0) {
+            summary = textProto { resourceId = metadata.summary }
+        } else {
+            (metadata as? PreferenceSummaryProvider)?.getSummary(context)?.let {
+                summary = textProto { string = it.toString() }
+            }
+        }
+        if (metadata.icon != 0) icon = metadata.icon
+        if (metadata.keywords != 0) keywords = metadata.keywords
+        val preferenceExtras = metadata.extras(context)
+        preferenceExtras?.let { extras = it.toProto() }
+        indexable = metadata.isIndexable(context)
+        enabled = metadata.isEnabled(context)
+        if (metadata is PreferenceAvailabilityProvider) {
+            available = metadata.isAvailable(context)
+        }
+        if (metadata is PreferenceRestrictionProvider) {
+            restricted = metadata.isRestricted(context)
+        }
+        persistent = metadata.isPersistent(context)
+        if (includeValue &&
+            persistent &&
+            metadata is BooleanValue &&
+            metadata is PersistentPreference<*>) {
+            metadata.storage(context).getValue(metadata.key, Boolean::class.javaObjectType)?.let {
+                value = preferenceValueProto { booleanValue = it }
+            }
+        }
+        if (metadata is PreferenceScreenMetadata) {
+            if (metadata.hasCompleteHierarchy()) {
+                @Suppress("CheckReturnValue") addPreferenceScreenMetadata(metadata, activityClass)
+            } else {
+                metadata.fragmentClass()?.toActionTarget(activityClass, preferenceExtras)?.let {
+                    actionTarget = it
+                }
+            }
+        }
+        metadata.intent(context)?.let { actionTarget = it.toActionTarget() }
+    }
+
+    private fun PreferenceMetadata.getTitleTextProto(isRoot: Boolean): TextProto? {
+        if (isRoot && this is PreferenceScreenMetadata) {
+            val titleRes = screenTitle
+            if (titleRes != 0) {
+                return textProto { resourceId = titleRes }
+            } else {
+                getScreenTitle(context)?.let {
+                    return textProto { string = it.toString() }
+                }
+            }
+        } else {
+            val titleRes = title
+            if (titleRes != 0) {
+                return textProto { resourceId = titleRes }
+            }
+        }
+        return (this as? PreferenceTitleProvider)?.getTitle(context)?.let {
+            textProto { string = it.toString() }
+        }
+    }
+
+    private suspend fun String?.toActionTarget(
+        activityClass: Class<*>,
+        extras: Bundle?,
+    ): ActionTarget? {
+        if (this.isNullOrEmpty()) return null
+        try {
+            val fragmentClass = context.classLoader.loadClass(this)
+            if (Fragment::class.java.isAssignableFrom(fragmentClass)) {
+                @Suppress("UNCHECKED_CAST")
+                return (fragmentClass as Class<out Fragment>).toActionTarget(activityClass, extras)
+            }
+        } catch (e: Exception) {
+            Log.e(TAG, "Cannot loadClass $this", e)
+        }
+        return null
+    }
+
+    private suspend fun Class<out Fragment>.toActionTarget(
+        activityClass: Class<*>,
+        extras: Bundle?,
+    ): ActionTarget {
+        val startIntent = Intent(context, activityClass)
+        startIntent.putExtra(PreferenceActivity.EXTRA_SHOW_FRAGMENT, name)
+        extras?.let { startIntent.putExtra(PreferenceActivity.EXTRA_SHOW_FRAGMENT_ARGUMENTS, it) }
+        if (!PreferenceScreenProvider::class.java.isAssignableFrom(this) &&
+            !PreferenceScreenBindingKeyProvider::class.java.isAssignableFrom(this)) {
+            return actionTargetProto { intent = startIntent.toProto() }
+        }
+        val fragment =
+            withContext(Dispatchers.Main) {
+                return@withContext try {
+                    newInstance().apply { arguments = extras }
+                } catch (e: Exception) {
+                    Log.e(TAG, "Fail to instantiate fragment ${this@toActionTarget}", e)
+                    null
+                }
+            }
+        if (fragment is PreferenceScreenBindingKeyProvider) {
+            val screenKey = fragment.getPreferenceScreenBindingKey(context)
+            if (screenKey != null && addPreferenceScreenFromRegistry(screenKey, activityClass)) {
+                return actionTargetProto { key = screenKey }
+            }
+        }
+        if (fragment is PreferenceScreenProvider) {
+            val screen = fragment.createPreferenceScreen(preferenceScreenFactory)
+            if (screen != null) {
+                addPreferenceScreen(startIntent, activityClass, screen)
+                return actionTargetProto { key = screen.key }
+            }
+        }
+        return actionTargetProto { intent = startIntent.toProto() }
+    }
+
+    private suspend fun Intent.toActionTarget(): ActionTarget {
+        if (component?.packageName == "") {
+            setClassName(context, component!!.className)
+        }
+        resolveActivity(context.packageManager)?.let {
+            if (it.packageName == context.packageName) {
+                add(it.className)
+            }
+        }
+        return actionTargetProto { intent = toProto() }
+    }
+
+    companion object {
+        suspend fun of(context: Context, request: GetPreferenceGraphRequest) =
+            PreferenceGraphBuilder(context, request).also { it.init() }
+    }
+}
+
+@SuppressLint("AppBundleLocaleChanges")
+internal fun Context.ofLocale(locale: Locale?): Context {
+    if (locale == null) return this
+    val baseConfig: Configuration = resources.configuration
+    val baseLocale =
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
+            baseConfig.locales[0]
+        } else {
+            baseConfig.locale
+        }
+    if (locale == baseLocale) {
+        return this
+    }
+    val newConfig = Configuration(baseConfig)
+    newConfig.setLocale(locale)
+    return createConfigurationContext(newConfig)
+}
diff --git a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceScreenManager.kt b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceScreenManager.kt
deleted file mode 100644
index 9231f40..0000000
--- a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceScreenManager.kt
+++ /dev/null
@@ -1,70 +0,0 @@
-package com.android.settingslib.graph
-
-import androidx.annotation.StringRes
-import androidx.annotation.XmlRes
-import androidx.preference.Preference
-import androidx.preference.PreferenceManager
-import androidx.preference.PreferenceScreen
-
-/** Manager to create and initialize preference screen. */
-class PreferenceScreenManager(private val preferenceManager: PreferenceManager) {
-    private val context = preferenceManager.context
-    // the map will preserve order
-    private val updaters = mutableMapOf<String, PreferenceUpdater>()
-    private val screenUpdaters = mutableListOf<PreferenceScreenUpdater>()
-
-    /** Creates an empty [PreferenceScreen]. */
-    fun createPreferenceScreen(): PreferenceScreen =
-        preferenceManager.createPreferenceScreen(context)
-
-    /** Creates [PreferenceScreen] from resource. */
-    fun createPreferenceScreen(@XmlRes xmlRes: Int): PreferenceScreen =
-        preferenceManager.inflateFromResource(context, xmlRes, null)
-
-    /** Adds updater for given preference. */
-    fun addPreferenceUpdater(@StringRes key: Int, updater: PreferenceUpdater) =
-        addPreferenceUpdater(context.getString(key), updater)
-
-    /** Adds updater for given preference. */
-    fun addPreferenceUpdater(
-        key: String,
-        updater: PreferenceUpdater,
-    ): PreferenceScreenManager {
-        updaters.put(key, updater)?.let { if (it != updater) throw IllegalArgumentException() }
-        return this
-    }
-
-    /** Adds updater for preference screen. */
-    fun addPreferenceScreenUpdater(updater: PreferenceScreenUpdater): PreferenceScreenManager {
-        screenUpdaters.add(updater)
-        return this
-    }
-
-    /** Adds a list of updaters for preference screen. */
-    fun addPreferenceScreenUpdater(
-        vararg updaters: PreferenceScreenUpdater,
-    ): PreferenceScreenManager {
-        screenUpdaters.addAll(updaters)
-        return this
-    }
-
-    /** Updates preference screen with registered updaters. */
-    fun updatePreferenceScreen(preferenceScreen: PreferenceScreen) {
-        for ((key, updater) in updaters) {
-            preferenceScreen.findPreference<Preference>(key)?.let { updater.updatePreference(it) }
-        }
-        for (updater in screenUpdaters) {
-            updater.updatePreferenceScreen(preferenceScreen)
-        }
-    }
-}
-
-/** Updater of [Preference]. */
-interface PreferenceUpdater {
-    fun updatePreference(preference: Preference)
-}
-
-/** Updater of [PreferenceScreen]. */
-interface PreferenceScreenUpdater {
-    fun updatePreferenceScreen(preferenceScreen: PreferenceScreen)
-}
diff --git a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceScreenProvider.kt b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceScreenProvider.kt
deleted file mode 100644
index 9e4c1f6..0000000
--- a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/PreferenceScreenProvider.kt
+++ /dev/null
@@ -1,26 +0,0 @@
-package com.android.settingslib.graph
-
-import android.content.Context
-import androidx.preference.PreferenceScreen
-
-/**
- * Interface to provide [PreferenceScreen].
- *
- * It is expected to be implemented by Activity/Fragment and the implementation needs to use
- * [Context] APIs (e.g. `getContext()`, `getActivity()`) with caution: preference screen creation
- * could happen in background service, where the Activity/Fragment lifecycle callbacks (`onCreate`,
- * `onDestroy`, etc.) are not invoked.
- */
-interface PreferenceScreenProvider {
-
-    /**
-     * Creates [PreferenceScreen].
-     *
-     * Preference screen creation could happen in background service. The implementation MUST use
-     * given [context] instead of APIs like `getContext()`, `getActivity()`, etc.
-     */
-    fun createPreferenceScreen(
-        context: Context,
-        preferenceScreenManager: PreferenceScreenManager,
-    ): PreferenceScreen?
-}
diff --git a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/ProtoConverters.kt b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/ProtoConverters.kt
new file mode 100644
index 0000000..d9b9590
--- /dev/null
+++ b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/ProtoConverters.kt
@@ -0,0 +1,74 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.graph
+
+import android.content.Context
+import android.content.Intent
+import android.os.Bundle
+import com.android.settingslib.graph.proto.BundleProto
+import com.android.settingslib.graph.proto.BundleProto.BundleValue
+import com.android.settingslib.graph.proto.IntentProto
+import com.android.settingslib.graph.proto.TextProto
+import com.google.protobuf.ByteString
+
+fun TextProto.getText(context: Context): String? =
+    when {
+        hasResourceId() -> context.getString(resourceId)
+        hasString() -> string
+        else -> null
+    }
+
+fun Intent.toProto(): IntentProto = intentProto {
+    this@toProto.action?.let { action = it }
+    this@toProto.dataString?.let { data = it }
+    this@toProto.`package`?.let { pkg = it }
+    this@toProto.component?.let { component = it.flattenToShortString() }
+    this@toProto.flags.let { if (it != 0) flags = it }
+    this@toProto.extras?.let { extras = it.toProto() }
+    this@toProto.type?.let { mimeType = it }
+}
+
+fun Bundle.toProto(): BundleProto = bundleProto {
+    fun toProto(value: Any): BundleValue = bundleValueProto {
+        when (value) {
+            is String -> stringValue = value
+            is ByteArray -> bytesValue = ByteString.copyFrom(value)
+            is Int -> intValue = value
+            is Long -> longValue = value
+            is Boolean -> booleanValue = value
+            is Double -> doubleValue = value
+            is Bundle -> bundleValue = value.toProto()
+            else -> throw IllegalArgumentException("Unknown type: ${value.javaClass} $value")
+        }
+    }
+
+    for (key in keySet()) {
+        @Suppress("DEPRECATION") get(key)?.let { putValues(key, toProto(it)) }
+    }
+}
+
+fun BundleValue.stringify(): String =
+    when {
+        hasBooleanValue() -> "$valueCase"
+        hasBytesValue() -> "$bytesValue"
+        hasIntValue() -> "$intValue"
+        hasLongValue() -> "$longValue"
+        hasStringValue() -> stringValue
+        hasDoubleValue() -> "$doubleValue"
+        hasBundleValue() -> "$bundleValue"
+        else -> "Unknown"
+    }
diff --git a/packages/SettingsLib/Graph/src/com/android/settingslib/graph/ProtoDsl.kt b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/ProtoDsl.kt
new file mode 100644
index 0000000..d7dae77
--- /dev/null
+++ b/packages/SettingsLib/Graph/src/com/android/settingslib/graph/ProtoDsl.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.settingslib.graph
+
+import com.android.settingslib.graph.proto.BundleProto
+import com.android.settingslib.graph.proto.BundleProto.BundleValue
+import com.android.settingslib.graph.proto.IntentProto
+import com.android.settingslib.graph.proto.PreferenceGroupProto
+import com.android.settingslib.graph.proto.PreferenceOrGroupProto
+import com.android.settingslib.graph.proto.PreferenceProto
+import com.android.settingslib.graph.proto.PreferenceProto.ActionTarget
+import com.android.settingslib.graph.proto.PreferenceScreenProto
+import com.android.settingslib.graph.proto.PreferenceValueProto
+import com.android.settingslib.graph.proto.TextProto
+
+/** Returns root or null. */
+val PreferenceScreenProto.rootOrNull
+    get() = if (hasRoot()) root else null
+
+/** Kotlin DSL-style builder for [PreferenceScreenProto]. */
+@JvmSynthetic
+inline fun preferenceScreenProto(init: PreferenceScreenProto.Builder.() -> Unit) =
+    PreferenceScreenProto.newBuilder().also(init).build()
+
+/** Returns preference or null. */
+val PreferenceOrGroupProto.preferenceOrNull
+    get() = if (hasPreference()) preference else null
+
+/** Returns group or null. */
+val PreferenceOrGroupProto.groupOrNull
+    get() = if (hasGroup()) group else null
+
+/** Kotlin DSL-style builder for [PreferenceOrGroupProto]. */
+@JvmSynthetic
+inline fun preferenceOrGroupProto(init: PreferenceOrGroupProto.Builder.() -> Unit) =
+    PreferenceOrGroupProto.newBuilder().also(init).build()
+
+/** Returns preference or null. */
+val PreferenceGroupProto.preferenceOrNull
+    get() = if (hasPreference()) preference else null
+
+/** Kotlin DSL-style builder for [PreferenceGroupProto]. */
+@JvmSynthetic
+inline fun preferenceGroupProto(init: PreferenceGroupProto.Builder.() -> Unit) =
+    PreferenceGroupProto.newBuilder().also(init).build()
+
+/** Returns title or null. */
+val PreferenceProto.titleOrNull
+    get() = if (hasTitle()) title else null
+
+/** Returns summary or null. */
+val PreferenceProto.summaryOrNull
+    get() = if (hasSummary()) summary else null
+
+/** Returns actionTarget or null. */
+val PreferenceProto.actionTargetOrNull
+    get() = if (hasActionTarget()) actionTarget else null
+
+/** Kotlin DSL-style builder for [PreferenceProto]. */
+@JvmSynthetic
+inline fun preferenceProto(init: PreferenceProto.Builder.() -> Unit) =
+    PreferenceProto.newBuilder().also(init).build()
+
+/** Returns intent or null. */
+val ActionTarget.intentOrNull
+    get() = if (hasIntent()) intent else null
+
+/** Kotlin DSL-style builder for [ActionTarget]. */
+@JvmSynthetic
+inline fun actionTargetProto(init: ActionTarget.Builder.() -> Unit) =
+    ActionTarget.newBuilder().also(init).build()
+
+/** Kotlin DSL-style builder for [PreferenceValueProto]. */
+@JvmSynthetic
+inline fun preferenceValueProto(init: PreferenceValueProto.Builder.() -> Unit) =
+    PreferenceValueProto.newBuilder().also(init).build()
+
+/** Kotlin DSL-style builder for [TextProto]. */
+@JvmSynthetic
+inline fun textProto(init: TextProto.Builder.() -> Unit) = TextProto.newBuilder().also(init).build()
+
+/** Kotlin DSL-style builder for [IntentProto]. */
+@JvmSynthetic
+inline fun intentProto(init: IntentProto.Builder.() -> Unit) =
+    IntentProto.newBuilder().also(init).build()
+
+/** Kotlin DSL-style builder for [BundleProto]. */
+@JvmSynthetic
+inline fun bundleProto(init: BundleProto.Builder.() -> Unit) =
+    BundleProto.newBuilder().also(init).build()
+
+/** Kotlin DSL-style builder for [BundleValue]. */
+@JvmSynthetic
+inline fun bundleValueProto(init: BundleValue.Builder.() -> Unit) =
+    BundleValue.newBuilder().also(init).build()
diff --git a/packages/SettingsLib/Ipc/Android.bp b/packages/SettingsLib/Ipc/Android.bp
index 61adb46..2c7209a 100644
--- a/packages/SettingsLib/Ipc/Android.bp
+++ b/packages/SettingsLib/Ipc/Android.bp
@@ -20,3 +20,15 @@
     ],
     kotlincflags: ["-Xjvm-default=all"],
 }
+
+android_library {
+    name: "SettingsLibIpc-testutils",
+    srcs: ["testutils/**/*.kt"],
+    static_libs: [
+        "Robolectric_all-target_upstream",
+        "SettingsLibIpc",
+        "androidx.test.core",
+        "flag-junit",
+        "kotlinx-coroutines-android",
+    ],
+}
diff --git a/packages/SettingsLib/Ipc/testutils/com/android/settingslib/ipc/MessengerServiceRule.kt b/packages/SettingsLib/Ipc/testutils/com/android/settingslib/ipc/MessengerServiceRule.kt
new file mode 100644
index 0000000..8b2deaf
--- /dev/null
+++ b/packages/SettingsLib/Ipc/testutils/com/android/settingslib/ipc/MessengerServiceRule.kt
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.ipc
+
+import android.app.Application
+import android.app.Service
+import android.content.ComponentName
+import android.content.Intent
+import android.os.Build
+import android.os.Looper
+import androidx.test.core.app.ApplicationProvider
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withContext
+import org.junit.rules.TestWatcher
+import org.junit.runner.Description
+import org.robolectric.Robolectric
+import org.robolectric.Shadows
+import org.robolectric.android.controller.ServiceController
+
+/** Rule for messenger service testing. */
+open class MessengerServiceRule<C : MessengerServiceClient>(
+    private val serviceClass: Class<out MessengerService>,
+    val client: C,
+) : TestWatcher() {
+    val application: Application = ApplicationProvider.getApplicationContext()
+    val isRobolectric = Build.FINGERPRINT.contains("robolectric")
+
+    private var serviceController: ServiceController<out Service>? = null
+
+    override fun starting(description: Description) {
+        if (isRobolectric) {
+            runBlocking { setupRobolectricService() }
+        }
+    }
+
+    override fun finished(description: Description) {
+        client.close()
+        if (isRobolectric) {
+            runBlocking {
+                withContext(Dispatchers.Main) { serviceController?.run { unbind().destroy() } }
+            }
+        }
+    }
+
+    private suspend fun setupRobolectricService() {
+        if (Thread.currentThread() == Looper.getMainLooper().thread) {
+            throw IllegalStateException(
+                "To avoid deadlock, run test with @LooperMode(LooperMode.Mode.INSTRUMENTATION_TEST)"
+            )
+        }
+        withContext(Dispatchers.Main) {
+            serviceController = Robolectric.buildService(serviceClass)
+            val service = serviceController!!.create().get()
+            Shadows.shadowOf(application).apply {
+                setComponentNameAndServiceForBindService(
+                    ComponentName(application, serviceClass),
+                    service.onBind(Intent(application, serviceClass)),
+                )
+                setBindServiceCallsOnServiceConnectedDirectly(true)
+                setUnbindServiceCallsOnServiceDisconnected(false)
+            }
+        }
+    }
+}
diff --git a/packages/SettingsLib/res/values/strings.xml b/packages/SettingsLib/res/values/strings.xml
index 4d771c0..feee89a 100644
--- a/packages/SettingsLib/res/values/strings.xml
+++ b/packages/SettingsLib/res/values/strings.xml
@@ -1411,6 +1411,8 @@
     <string name="media_transfer_this_device_name_tablet">This tablet</string>
     <!-- Name of the default media output of the TV. [CHAR LIMIT=30] -->
     <string name="media_transfer_this_device_name_tv">@string/tv_media_transfer_default</string>
+    <!-- Name of the internal mic. [CHAR LIMIT=30] -->
+    <string name="media_transfer_internal_mic">Microphone (internal)</string>
     <!-- Name of the dock device. [CHAR LIMIT=30] -->
     <string name="media_transfer_dock_speaker_device_name">Dock speaker</string>
     <!-- Default name of the external device. [CHAR LIMIT=30] -->
@@ -1637,6 +1639,12 @@
     <!-- Name of the 3.5mm and usb audio device. [CHAR LIMIT=50] -->
     <string name="media_transfer_wired_usb_device_name">Wired headphone</string>
 
+    <!-- Name of the 3.5mm audio device mic. [CHAR LIMIT=50] -->
+    <string name="media_transfer_wired_device_mic_name">Mic jack</string>
+
+    <!-- Name of the usb audio device mic. [CHAR LIMIT=50] -->
+    <string name="media_transfer_usb_device_mic_name">USB mic</string>
+
     <!-- Label for Wifi hotspot switch on. Toggles hotspot on [CHAR LIMIT=30] -->
     <string name="wifi_hotspot_switch_on_text">On</string>
     <!-- Label for Wifi hotspot switch off. Toggles hotspot off [CHAR LIMIT=30] -->
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/InputMediaDevice.java b/packages/SettingsLib/src/com/android/settingslib/media/InputMediaDevice.java
new file mode 100644
index 0000000..766cd43
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/media/InputMediaDevice.java
@@ -0,0 +1,161 @@
+/*
+ * 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.settingslib.media;
+
+import static android.media.AudioDeviceInfo.TYPE_BUILTIN_MIC;
+import static android.media.AudioDeviceInfo.TYPE_USB_ACCESSORY;
+import static android.media.AudioDeviceInfo.TYPE_USB_DEVICE;
+import static android.media.AudioDeviceInfo.TYPE_USB_HEADSET;
+import static android.media.AudioDeviceInfo.TYPE_WIRED_HEADSET;
+
+import static com.android.settingslib.media.MediaDevice.SelectionBehavior.SELECTION_BEHAVIOR_TRANSFER;
+
+import android.content.Context;
+import android.graphics.drawable.Drawable;
+import android.media.AudioDeviceInfo.AudioDeviceType;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+
+import com.android.settingslib.R;
+
+/** {@link MediaDevice} implementation that represents an input device. */
+public class InputMediaDevice extends MediaDevice {
+
+    private static final String TAG = "InputMediaDevice";
+
+    private final String mId;
+
+    private final @AudioDeviceType int mAudioDeviceInfoType;
+
+    private final int mMaxVolume;
+
+    private final int mCurrentVolume;
+
+    private final boolean mIsVolumeFixed;
+
+    private InputMediaDevice(
+            @NonNull Context context,
+            @NonNull String id,
+            @AudioDeviceType int audioDeviceInfoType,
+            int maxVolume,
+            int currentVolume,
+            boolean isVolumeFixed) {
+        super(context, /* info= */ null, /* item= */ null);
+        mId = id;
+        mAudioDeviceInfoType = audioDeviceInfoType;
+        mMaxVolume = maxVolume;
+        mCurrentVolume = currentVolume;
+        mIsVolumeFixed = isVolumeFixed;
+        initDeviceRecord();
+    }
+
+    @Nullable
+    public static InputMediaDevice create(
+            @NonNull Context context,
+            @NonNull String id,
+            @AudioDeviceType int audioDeviceInfoType,
+            int maxVolume,
+            int currentVolume,
+            boolean isVolumeFixed) {
+        if (!isSupportedInputDevice(audioDeviceInfoType)) {
+            return null;
+        }
+
+        return new InputMediaDevice(
+                context, id, audioDeviceInfoType, maxVolume, currentVolume, isVolumeFixed);
+    }
+
+    public static boolean isSupportedInputDevice(@AudioDeviceType int audioDeviceInfoType) {
+        return switch (audioDeviceInfoType) {
+            case TYPE_BUILTIN_MIC,
+                            TYPE_WIRED_HEADSET,
+                            TYPE_USB_DEVICE,
+                            TYPE_USB_HEADSET,
+                            TYPE_USB_ACCESSORY ->
+                    true;
+            default -> false;
+        };
+    }
+
+    @Override
+    public @NonNull String getName() {
+        CharSequence name =
+                switch (mAudioDeviceInfoType) {
+                    case TYPE_WIRED_HEADSET ->
+                            mContext.getString(R.string.media_transfer_wired_device_mic_name);
+                    case TYPE_USB_DEVICE, TYPE_USB_HEADSET, TYPE_USB_ACCESSORY ->
+                            mContext.getString(R.string.media_transfer_usb_device_mic_name);
+                    default -> mContext.getString(R.string.media_transfer_internal_mic);
+                };
+        return name.toString();
+    }
+
+    @Override
+    public @SelectionBehavior int getSelectionBehavior() {
+        // We don't allow apps to override the selection behavior of system routes.
+        return SELECTION_BEHAVIOR_TRANSFER;
+    }
+
+    @Override
+    public @NonNull String getSummary() {
+        return "";
+    }
+
+    @Override
+    public @Nullable Drawable getIcon() {
+        return getIconWithoutBackground();
+    }
+
+    @Override
+    public @Nullable Drawable getIconWithoutBackground() {
+        return mContext.getDrawable(getDrawableResId());
+    }
+
+    @VisibleForTesting
+    int getDrawableResId() {
+        // TODO(b/357122624): check with UX to obtain the icon for desktop devices.
+        return R.drawable.ic_media_tablet;
+    }
+
+    @Override
+    public @NonNull String getId() {
+        return mId;
+    }
+
+    @Override
+    public boolean isConnected() {
+        // Indicating if the device is connected and thus showing the status of STATE_CONNECTED.
+        // Upon creation, this device is already connected.
+        return true;
+    }
+
+    @Override
+    public int getMaxVolume() {
+        return mMaxVolume;
+    }
+
+    @Override
+    public int getCurrentVolume() {
+        return mCurrentVolume;
+    }
+
+    @Override
+    public boolean isVolumeFixed() {
+        return mIsVolumeFixed;
+    }
+}
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/InputRouteManager.java b/packages/SettingsLib/src/com/android/settingslib/media/InputRouteManager.java
new file mode 100644
index 0000000..548eb3f
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/media/InputRouteManager.java
@@ -0,0 +1,126 @@
+/*
+ * 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.settingslib.media;
+
+import android.content.Context;
+import android.media.AudioDeviceCallback;
+import android.media.AudioDeviceInfo;
+import android.media.AudioManager;
+import android.os.Handler;
+
+import androidx.annotation.NonNull;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+/** Provides functionalities to get/observe input routes, control input routing and volume gain. */
+public final class InputRouteManager {
+
+    private static final String TAG = "InputRouteManager";
+
+    private final Context mContext;
+
+    private final AudioManager mAudioManager;
+
+    @VisibleForTesting final List<MediaDevice> mInputMediaDevices = new CopyOnWriteArrayList<>();
+
+    private final Collection<InputDeviceCallback> mCallbacks = new CopyOnWriteArrayList<>();
+
+    @VisibleForTesting
+    final AudioDeviceCallback mAudioDeviceCallback =
+            new AudioDeviceCallback() {
+                @Override
+                public void onAudioDevicesAdded(@NonNull AudioDeviceInfo[] addedDevices) {
+                    dispatchInputDeviceListUpdate();
+                }
+
+                @Override
+                public void onAudioDevicesRemoved(@NonNull AudioDeviceInfo[] removedDevices) {
+                    dispatchInputDeviceListUpdate();
+                }
+            };
+
+    /* package */ InputRouteManager(@NonNull Context context, @NonNull AudioManager audioManager) {
+        mContext = context;
+        mAudioManager = audioManager;
+        Handler handler = new Handler(context.getMainLooper());
+
+        mAudioManager.registerAudioDeviceCallback(mAudioDeviceCallback, handler);
+    }
+
+    public void registerCallback(@NonNull InputDeviceCallback callback) {
+        if (!mCallbacks.contains(callback)) {
+            mCallbacks.add(callback);
+            dispatchInputDeviceListUpdate();
+        }
+    }
+
+    public void unregisterCallback(@NonNull InputDeviceCallback callback) {
+        mCallbacks.remove(callback);
+    }
+
+    private void dispatchInputDeviceListUpdate() {
+        // TODO (b/360175574): Get selected input device.
+
+        // Get all input devices.
+        AudioDeviceInfo[] audioDeviceInfos =
+                mAudioManager.getDevices(AudioManager.GET_DEVICES_INPUTS);
+        mInputMediaDevices.clear();
+        for (AudioDeviceInfo info : audioDeviceInfos) {
+            MediaDevice mediaDevice =
+                    InputMediaDevice.create(
+                            mContext,
+                            String.valueOf(info.getId()),
+                            info.getType(),
+                            getMaxInputGain(),
+                            getCurrentInputGain(),
+                            isInputGainFixed());
+            if (mediaDevice != null) {
+                mInputMediaDevices.add(mediaDevice);
+            }
+        }
+
+        final List<MediaDevice> inputMediaDevices = new ArrayList<>(mInputMediaDevices);
+        for (InputDeviceCallback callback : mCallbacks) {
+            callback.onInputDeviceListUpdated(inputMediaDevices);
+        }
+    }
+
+    public int getMaxInputGain() {
+        // TODO (b/357123335): use real input gain implementation.
+        // Using 15 for now since it matches the max index for output.
+        return 15;
+    }
+
+    public int getCurrentInputGain() {
+        // TODO (b/357123335): use real input gain implementation.
+        return 8;
+    }
+
+    public boolean isInputGainFixed() {
+        // TODO (b/357123335): use real input gain implementation.
+        return true;
+    }
+
+    /** Callback for listening to input device changes. */
+    public interface InputDeviceCallback {
+        void onInputDeviceListUpdated(@NonNull List<MediaDevice> devices);
+    }
+}
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..11406fa 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
@@ -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/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputMediaDeviceTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputMediaDeviceTest.java
new file mode 100644
index 0000000..bc1ea6c
--- /dev/null
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputMediaDeviceTest.java
@@ -0,0 +1,114 @@
+/*
+ * 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.settingslib.media;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.media.AudioDeviceInfo;
+import android.platform.test.flag.junit.SetFlagsRule;
+
+import com.android.settingslib.R;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+
+@RunWith(RobolectricTestRunner.class)
+public class InputMediaDeviceTest {
+
+    private final int BUILTIN_MIC_ID = 1;
+    private final int WIRED_HEADSET_ID = 2;
+    private final int USB_HEADSET_ID = 3;
+    private final int MAX_VOLUME = 1;
+    private final int CURRENT_VOLUME = 0;
+    private final boolean IS_VOLUME_FIXED = true;
+
+    @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    private Context mContext;
+
+    @Before
+    public void setUp() {
+        mContext = RuntimeEnvironment.application;
+    }
+
+    @Test
+    public void getDrawableResId_returnCorrectResId() {
+        InputMediaDevice builtinMediaDevice =
+                InputMediaDevice.create(
+                        mContext,
+                        String.valueOf(BUILTIN_MIC_ID),
+                        AudioDeviceInfo.TYPE_BUILTIN_MIC,
+                        MAX_VOLUME,
+                        CURRENT_VOLUME,
+                        IS_VOLUME_FIXED);
+        assertThat(builtinMediaDevice).isNotNull();
+        assertThat(builtinMediaDevice.getDrawableResId()).isEqualTo(R.drawable.ic_media_tablet);
+    }
+
+    @Test
+    public void getName_returnCorrectName_builtinMic() {
+        InputMediaDevice builtinMediaDevice =
+                InputMediaDevice.create(
+                        mContext,
+                        String.valueOf(BUILTIN_MIC_ID),
+                        AudioDeviceInfo.TYPE_BUILTIN_MIC,
+                        MAX_VOLUME,
+                        CURRENT_VOLUME,
+                        IS_VOLUME_FIXED);
+        assertThat(builtinMediaDevice).isNotNull();
+        assertThat(builtinMediaDevice.getName())
+                .isEqualTo(mContext.getString(R.string.media_transfer_internal_mic));
+    }
+
+    @Test
+    public void getName_returnCorrectName_wiredHeadset() {
+        InputMediaDevice wiredMediaDevice =
+                InputMediaDevice.create(
+                        mContext,
+                        String.valueOf(WIRED_HEADSET_ID),
+                        AudioDeviceInfo.TYPE_WIRED_HEADSET,
+                        MAX_VOLUME,
+                        CURRENT_VOLUME,
+                        IS_VOLUME_FIXED);
+        assertThat(wiredMediaDevice).isNotNull();
+        assertThat(wiredMediaDevice.getName())
+                .isEqualTo(mContext.getString(R.string.media_transfer_wired_device_mic_name));
+    }
+
+    @Test
+    public void getName_returnCorrectName_usbHeadset() {
+        InputMediaDevice usbMediaDevice =
+                InputMediaDevice.create(
+                        mContext,
+                        String.valueOf(USB_HEADSET_ID),
+                        AudioDeviceInfo.TYPE_USB_HEADSET,
+                        MAX_VOLUME,
+                        CURRENT_VOLUME,
+                        IS_VOLUME_FIXED);
+        assertThat(usbMediaDevice).isNotNull();
+        assertThat(usbMediaDevice.getName())
+                .isEqualTo(mContext.getString(R.string.media_transfer_usb_device_mic_name));
+    }
+}
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputRouteManagerTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputRouteManagerTest.java
new file mode 100644
index 0000000..2501ae6
--- /dev/null
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputRouteManagerTest.java
@@ -0,0 +1,140 @@
+/*
+ * 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.settingslib.media;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.media.AudioDeviceInfo;
+import android.media.AudioManager;
+
+import com.android.settingslib.testutils.shadow.ShadowRouter2Manager;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
+
+@RunWith(RobolectricTestRunner.class)
+@Config(shadows = {ShadowRouter2Manager.class})
+public class InputRouteManagerTest {
+    private static final int BUILTIN_MIC_ID = 1;
+    private static final int INPUT_WIRED_HEADSET_ID = 2;
+    private static final int INPUT_USB_DEVICE_ID = 3;
+    private static final int INPUT_USB_HEADSET_ID = 4;
+    private static final int INPUT_USB_ACCESSORY_ID = 5;
+
+    private final Context mContext = spy(RuntimeEnvironment.application);
+    private InputRouteManager mInputRouteManager;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+
+        final AudioManager audioManager = mock(AudioManager.class);
+        mInputRouteManager = new InputRouteManager(mContext, audioManager);
+    }
+
+    @Test
+    public void onAudioDevicesAdded_shouldUpdateInputMediaDevice() {
+        final AudioDeviceInfo info1 = mock(AudioDeviceInfo.class);
+        when(info1.getType()).thenReturn(AudioDeviceInfo.TYPE_BUILTIN_MIC);
+        when(info1.getId()).thenReturn(BUILTIN_MIC_ID);
+
+        final AudioDeviceInfo info2 = mock(AudioDeviceInfo.class);
+        when(info2.getType()).thenReturn(AudioDeviceInfo.TYPE_WIRED_HEADSET);
+        when(info2.getId()).thenReturn(INPUT_WIRED_HEADSET_ID);
+
+        final AudioDeviceInfo info3 = mock(AudioDeviceInfo.class);
+        when(info3.getType()).thenReturn(AudioDeviceInfo.TYPE_USB_DEVICE);
+        when(info3.getId()).thenReturn(INPUT_USB_DEVICE_ID);
+
+        final AudioDeviceInfo info4 = mock(AudioDeviceInfo.class);
+        when(info4.getType()).thenReturn(AudioDeviceInfo.TYPE_USB_HEADSET);
+        when(info4.getId()).thenReturn(INPUT_USB_HEADSET_ID);
+
+        final AudioDeviceInfo info5 = mock(AudioDeviceInfo.class);
+        when(info5.getType()).thenReturn(AudioDeviceInfo.TYPE_USB_ACCESSORY);
+        when(info5.getId()).thenReturn(INPUT_USB_ACCESSORY_ID);
+
+        final AudioDeviceInfo unsupportedInfo = mock(AudioDeviceInfo.class);
+        when(unsupportedInfo.getType()).thenReturn(AudioDeviceInfo.TYPE_HDMI);
+
+        final AudioManager audioManager = mock(AudioManager.class);
+        AudioDeviceInfo[] devices = {info1, info2, info3, info4, info5, unsupportedInfo};
+        when(audioManager.getDevices(AudioManager.GET_DEVICES_INPUTS)).thenReturn(devices);
+
+        InputRouteManager inputRouteManager = new InputRouteManager(mContext, audioManager);
+
+        assertThat(inputRouteManager.mInputMediaDevices).isEmpty();
+
+        inputRouteManager.mAudioDeviceCallback.onAudioDevicesAdded(devices);
+
+        // The unsupported info should be filtered out.
+        assertThat(inputRouteManager.mInputMediaDevices).hasSize(devices.length - 1);
+        assertThat(inputRouteManager.mInputMediaDevices.get(0).getId())
+                .isEqualTo(String.valueOf(BUILTIN_MIC_ID));
+        assertThat(inputRouteManager.mInputMediaDevices.get(1).getId())
+                .isEqualTo(String.valueOf(INPUT_WIRED_HEADSET_ID));
+        assertThat(inputRouteManager.mInputMediaDevices.get(2).getId())
+                .isEqualTo(String.valueOf(INPUT_USB_DEVICE_ID));
+        assertThat(inputRouteManager.mInputMediaDevices.get(3).getId())
+                .isEqualTo(String.valueOf(INPUT_USB_HEADSET_ID));
+        assertThat(inputRouteManager.mInputMediaDevices.get(4).getId())
+                .isEqualTo(String.valueOf(INPUT_USB_ACCESSORY_ID));
+    }
+
+    @Test
+    public void onAudioDevicesRemoved_shouldUpdateInputMediaDevice() {
+        final AudioManager audioManager = mock(AudioManager.class);
+        when(audioManager.getDevices(AudioManager.GET_DEVICES_INPUTS))
+                .thenReturn(new AudioDeviceInfo[] {});
+
+        InputRouteManager inputRouteManager = new InputRouteManager(mContext, audioManager);
+
+        final MediaDevice device = mock(MediaDevice.class);
+        inputRouteManager.mInputMediaDevices.add(device);
+
+        final AudioDeviceInfo info = mock(AudioDeviceInfo.class);
+        when(info.getType()).thenReturn(AudioDeviceInfo.TYPE_WIRED_HEADSET);
+        inputRouteManager.mAudioDeviceCallback.onAudioDevicesRemoved(new AudioDeviceInfo[] {info});
+
+        assertThat(inputRouteManager.mInputMediaDevices).isEmpty();
+    }
+
+    @Test
+    public void getMaxInputGain_returnMaxInputGain() {
+        assertThat(mInputRouteManager.getMaxInputGain()).isEqualTo(15);
+    }
+
+    @Test
+    public void getCurrentInputGain_returnCurrentInputGain() {
+        assertThat(mInputRouteManager.getCurrentInputGain()).isEqualTo(8);
+    }
+
+    @Test
+    public void isInputGainFixed() {
+        assertThat(mInputRouteManager.isInputGainFixed()).isTrue();
+    }
+}
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/qs/ui/composable/QuickSettings.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
index 671b012..a6d5c1c 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
@@ -48,17 +48,13 @@
 import com.android.systemui.scene.shared.model.Scenes
 
 object QuickSettings {
-    private val SCENES =
-        setOf(
-            Scenes.QuickSettings,
-            Scenes.Shade,
-        )
+    private val SCENES = setOf(Scenes.QuickSettings, Scenes.Shade)
 
     object Elements {
         val Content =
             MovableElementKey(
                 "QuickSettingsContent",
-                contentPicker = MovableElementContentPicker(SCENES)
+                contentPicker = MovableElementContentPicker(SCENES),
             )
         val QuickQuickSettings = ElementKey("QuickQuickSettings")
         val SplitShadeQuickSettings = ElementKey("SplitShadeQuickSettings")
@@ -87,7 +83,7 @@
 
 private fun SceneScope.stateForQuickSettingsContent(
     isSplitShade: Boolean,
-    squishiness: () -> Float = { QuickSettings.SharedValues.SquishinessValues.Default }
+    squishiness: () -> Float = { QuickSettings.SharedValues.SquishinessValues.Default },
 ): QSSceneAdapter.State {
     return when (val transitionState = layoutState.transitionState) {
         is TransitionState.Idle -> {
@@ -122,7 +118,7 @@
                 }
             }
         is TransitionState.Transition.OverlayTransition ->
-            TODO("b/359173565: Handle overlay transitions")
+            error("Bad transition for QuickSettings scene: overlays not supported")
     }
 }
 
@@ -172,7 +168,7 @@
                 val height = heightProvider().coerceAtLeast(0)
 
                 layout(placeable.width, height) { placeable.placeRelative(0, 0) }
-            }
+            },
     ) {
         content { QuickSettingsContent(qsSceneAdapter = qsSceneAdapter, contentState) }
     }
@@ -225,7 +221,7 @@
                             it.addView(view)
                         }
                     },
-                    onRelease = { it.removeAllViews() }
+                    onRelease = { it.removeAllViews() },
                 )
             }
         }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
index 8922224..b85523b 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
@@ -61,23 +61,17 @@
     Box(modifier) {
         Scrim(onClicked = onScrimClicked)
 
-        Box(
-            modifier = Modifier.fillMaxSize().panelPadding(),
-            contentAlignment = Alignment.TopEnd,
-        ) {
+        Box(modifier = Modifier.fillMaxSize().panelPadding(), contentAlignment = Alignment.TopEnd) {
             Panel(
                 modifier = Modifier.element(OverlayShade.Elements.Panel).panelSize(),
-                content = content
+                content = content,
             )
         }
     }
 }
 
 @Composable
-private fun SceneScope.Scrim(
-    onClicked: () -> Unit,
-    modifier: Modifier = Modifier,
-) {
+private fun SceneScope.Scrim(onClicked: () -> Unit, modifier: Modifier = Modifier) {
     Spacer(
         modifier =
             modifier
@@ -89,10 +83,7 @@
 }
 
 @Composable
-private fun SceneScope.Panel(
-    modifier: Modifier = Modifier,
-    content: @Composable () -> Unit,
-) {
+private fun SceneScope.Panel(modifier: Modifier = Modifier, content: @Composable () -> Unit) {
     Box(modifier = modifier.clip(OverlayShade.Shapes.RoundedCornerPanel)) {
         Spacer(
             modifier =
@@ -101,7 +92,7 @@
                     .background(
                         color = OverlayShade.Colors.PanelBackground,
                         shape = OverlayShade.Shapes.RoundedCornerPanel,
-                    ),
+                    )
         )
 
         // This content is intentionally rendered as a separate element from the background in order
@@ -137,7 +128,7 @@
             systemBars.asPaddingValues(),
             displayCutout.asPaddingValues(),
             waterfall.asPaddingValues(),
-            contentPadding
+            contentPadding,
         )
 
     return if (widthSizeClass == WindowWidthSizeClass.Compact) {
@@ -156,14 +147,19 @@
         start = paddingValues.maxOfOrNull { it.calculateStartPadding(layoutDirection) } ?: 0.dp,
         top = paddingValues.maxOfOrNull { it.calculateTopPadding() } ?: 0.dp,
         end = paddingValues.maxOfOrNull { it.calculateEndPadding(layoutDirection) } ?: 0.dp,
-        bottom = paddingValues.maxOfOrNull { it.calculateBottomPadding() } ?: 0.dp
+        bottom = paddingValues.maxOfOrNull { it.calculateBottomPadding() } ?: 0.dp,
     )
 }
 
 object OverlayShade {
     object Elements {
         val Scrim = ElementKey("OverlayShadeScrim", contentPicker = LowestZIndexContentPicker)
-        val Panel = ElementKey("OverlayShadePanel", contentPicker = LowestZIndexContentPicker)
+        val Panel =
+            ElementKey(
+                "OverlayShadePanel",
+                contentPicker = LowestZIndexContentPicker,
+                placeAllCopies = true,
+            )
         val PanelBackground =
             ElementKey("OverlayShadePanelBackground", contentPicker = LowestZIndexContentPicker)
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
index f20548b..cec8883 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -100,6 +100,10 @@
      * By default overlays are centered in their layout but they can be aligned differently using
      * [alignment].
      *
+     * If [isModal] is true (the default), then a protective layer will be added behind the overlay
+     * to prevent swipes from reaching other scenes or overlays behind this one. Clicking this
+     * protective layer will close the overlay.
+     *
      * Important: overlays must be defined after all scenes. Overlay order along the z-axis follows
      * call order. Calling overlay(A) followed by overlay(B) will mean that overlay B renders
      * after/above overlay A.
@@ -109,6 +113,7 @@
         userActions: Map<UserAction, UserActionResult> =
             mapOf(Back to UserActionResult.HideOverlay(key)),
         alignment: Alignment = Alignment.Center,
+        isModal: Boolean = true,
         content: @Composable ContentScope.() -> Unit,
     )
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
index fe05234..65c4043 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
@@ -17,12 +17,16 @@
 package com.android.compose.animation.scene
 
 import androidx.annotation.VisibleForTesting
+import androidx.compose.foundation.clickable
 import androidx.compose.foundation.gestures.Orientation
+import androidx.compose.foundation.interaction.MutableInteractionSource
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.BoxScope
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.key
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.snapshots.SnapshotStateMap
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.ExperimentalComposeUiApi
@@ -253,6 +257,7 @@
                     key: OverlayKey,
                     userActions: Map<UserAction, UserActionResult>,
                     alignment: Alignment,
+                    isModal: Boolean,
                     content: @Composable (ContentScope.() -> Unit),
                 ) {
                     overlaysDefined = true
@@ -266,6 +271,7 @@
                         overlay.zIndex = zIndex
                         overlay.userActions = resolvedUserActions
                         overlay.alignment = alignment
+                        overlay.isModal = isModal
                     } else {
                         // New overlay.
                         overlays[key] =
@@ -276,6 +282,7 @@
                                 resolvedUserActions,
                                 zIndex,
                                 alignment,
+                                isModal,
                             )
                     }
 
@@ -399,12 +406,30 @@
             return
         }
 
-        // We put the overlays inside a Box that is matching the layout size so that overlays are
-        // measured after all scenes and that their max size is the size of the layout without the
-        // overlays.
-        Box(Modifier.matchParentSize().zIndex(overlaysOrderedByZIndex.first().zIndex)) {
-            overlaysOrderedByZIndex.fastForEach { overlay ->
-                key(overlay.key) { overlay.Content(Modifier.align(overlay.alignment)) }
+        overlaysOrderedByZIndex.fastForEach { overlay ->
+            val key = overlay.key
+            key(key) {
+                // We put the overlays inside a Box that is matching the layout size so that they
+                // are measured after all scenes and that their max size is the size of the layout
+                // without the overlays.
+                Box(Modifier.matchParentSize().zIndex(overlay.zIndex)) {
+                    if (overlay.isModal) {
+                        // Add a fullscreen clickable to prevent swipes from reaching the scenes and
+                        // other overlays behind this overlay. Clicking will close the overlay.
+                        Box(
+                            Modifier.fillMaxSize().clickable(
+                                interactionSource = remember { MutableInteractionSource() },
+                                indication = null,
+                            ) {
+                                if (state.canHideOverlay(key)) {
+                                    state.hideOverlay(key, animationScope = animationScope)
+                                }
+                            }
+                        )
+                    }
+
+                    overlay.Content(Modifier.align(overlay.alignment))
+                }
             }
         }
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Overlay.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Overlay.kt
index ccec9e8..d4de559 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Overlay.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/Overlay.kt
@@ -37,8 +37,10 @@
     actions: Map<UserAction.Resolved, UserActionResult>,
     zIndex: Float,
     alignment: Alignment,
+    isModal: Boolean,
 ) : Content(key, layoutImpl, content, actions, zIndex) {
     var alignment by mutableStateOf(alignment)
+    var isModal by mutableStateOf(isModal)
 
     override fun toString(): String {
         return "Overlay(key=$key)"
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
index ffed15b..cae6617 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
@@ -18,10 +18,12 @@
 
 import androidx.compose.animation.core.LinearEasing
 import androidx.compose.animation.core.tween
+import androidx.compose.foundation.ScrollState
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.verticalScroll
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
@@ -31,19 +33,26 @@
 import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.platform.testTag
 import androidx.compose.ui.test.assertIsDisplayed
 import androidx.compose.ui.test.assertIsNotDisplayed
 import androidx.compose.ui.test.assertPositionInRootIsEqualTo
+import androidx.compose.ui.test.click
 import androidx.compose.ui.test.hasTestTag
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.compose.ui.test.onNodeWithTag
+import androidx.compose.ui.test.onRoot
+import androidx.compose.ui.test.performTouchInput
+import androidx.compose.ui.test.swipe
+import androidx.compose.ui.test.swipeUp
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.dp
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.animation.scene.TestOverlays.OverlayA
 import com.android.compose.animation.scene.TestOverlays.OverlayB
 import com.android.compose.animation.scene.TestScenes.SceneA
+import com.android.compose.animation.scene.subjects.assertThat
 import com.android.compose.test.assertSizeIsEqualTo
 import com.android.compose.test.setContentAndCreateMainScope
 import com.android.compose.test.subjects.assertThat
@@ -769,4 +778,59 @@
             .assertSizeIsEqualTo(100.dp)
             .assertIsDisplayed()
     }
+
+    @Test
+    fun overlaysAreModalByDefault() {
+        val state = rule.runOnUiThread { MutableSceneTransitionLayoutStateImpl(SceneA) }
+
+        val scrollState = ScrollState(initial = 0)
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    // Make the scene vertically scrollable.
+                    scene(SceneA) {
+                        Box(Modifier.size(200.dp).verticalScroll(scrollState)) {
+                            Box(Modifier.size(200.dp, 400.dp))
+                        }
+                    }
+
+                    // The overlay is at the center end of the scene.
+                    overlay(OverlayA, alignment = Alignment.CenterEnd) {
+                        Box(Modifier.size(100.dp))
+                    }
+                }
+            }
+
+        fun swipeUp() {
+            rule.onRoot().performTouchInput {
+                swipe(start = Offset(x = 0f, y = bottom), end = Offset(x = 0f, y = top))
+            }
+        }
+
+        // Swiping up on the scene scrolls the list.
+        assertThat(scrollState.value).isEqualTo(0)
+        swipeUp()
+        assertThat(scrollState.value).isNotEqualTo(0)
+
+        // Reset the scroll.
+        scope.launch { scrollState.scrollTo(0) }
+        rule.waitForIdle()
+        assertThat(scrollState.value).isEqualTo(0)
+
+        // Show the overlay.
+        rule.runOnUiThread { state.showOverlay(OverlayA, animationScope = scope) }
+        rule.waitForIdle()
+        assertThat(state.transitionState).isIdle()
+        assertThat(state.transitionState).hasCurrentOverlays(OverlayA)
+
+        // Swiping up does not scroll the scene behind the overlay.
+        swipeUp()
+        assertThat(scrollState.value).isEqualTo(0)
+
+        // Clicking outside the overlay will close it.
+        rule.onRoot().performTouchInput { click(Offset.Zero) }
+        rule.waitForIdle()
+        assertThat(state.transitionState).isIdle()
+        assertThat(state.transitionState).hasCurrentOverlays(/* empty */ )
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractorTest.kt
index edaa3d3..bf97afe 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractorTest.kt
@@ -18,9 +18,12 @@
 
 package com.android.systemui.scene.domain.interactor
 
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.ObservableTransitionState.Transition.ShowOrHideOverlay
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
@@ -29,8 +32,10 @@
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.sceneDataSource
+import com.android.systemui.shade.shared.flag.DualShade
 import com.android.systemui.statusbar.domain.interactor.keyguardOcclusionInteractor
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.mock
@@ -63,10 +68,11 @@
     private val sceneDataSource =
         kosmos.sceneDataSource.apply { changeScene(toScene = Scenes.Lockscreen) }
 
-    private val underTest = kosmos.sceneContainerOcclusionInteractor
+    private val underTest by lazy { kosmos.sceneContainerOcclusionInteractor }
 
     @Test
-    fun invisibleDueToOcclusion() =
+    @DisableFlags(DualShade.FLAG_NAME)
+    fun invisibleDueToOcclusion_dualShadeDisabled() =
         testScope.runTest {
             val invisibleDueToOcclusion by collectLastValue(underTest.invisibleDueToOcclusion)
             val keyguardState by collectLastValue(keyguardTransitionInteractor.currentKeyguardState)
@@ -126,6 +132,68 @@
                 .isFalse()
         }
 
+    @Test
+    @EnableFlags(DualShade.FLAG_NAME)
+    fun invisibleDueToOcclusion_dualShadeEnabled() =
+        testScope.runTest {
+            val invisibleDueToOcclusion by collectLastValue(underTest.invisibleDueToOcclusion)
+            val keyguardState by collectLastValue(keyguardTransitionInteractor.currentKeyguardState)
+
+            // Assert that we have the desired preconditions:
+            assertThat(keyguardState).isEqualTo(KeyguardState.LOCKSCREEN)
+            assertThat(sceneInteractor.currentScene.value).isEqualTo(Scenes.Lockscreen)
+            assertThat(sceneInteractor.transitionState.value)
+                .isEqualTo(ObservableTransitionState.Idle(Scenes.Lockscreen))
+            assertWithMessage("Should start unoccluded").that(invisibleDueToOcclusion).isFalse()
+
+            // Actual testing starts here:
+            showOccludingActivity()
+            assertWithMessage("Should become occluded when occluding activity is shown")
+                .that(invisibleDueToOcclusion)
+                .isTrue()
+
+            transitionIntoAod {
+                assertWithMessage("Should become unoccluded when transitioning into AOD")
+                    .that(invisibleDueToOcclusion)
+                    .isFalse()
+            }
+            assertWithMessage("Should stay unoccluded when in AOD")
+                .that(invisibleDueToOcclusion)
+                .isFalse()
+
+            transitionOutOfAod {
+                assertWithMessage("Should remain unoccluded while transitioning away from AOD")
+                    .that(invisibleDueToOcclusion)
+                    .isFalse()
+            }
+            assertWithMessage("Should become occluded now that no longer in AOD")
+                .that(invisibleDueToOcclusion)
+                .isTrue()
+
+            expandDualShade {
+                assertWithMessage("Should become unoccluded once shade begins to expand")
+                    .that(invisibleDueToOcclusion)
+                    .isFalse()
+            }
+            assertWithMessage("Should be unoccluded when shade is fully expanded")
+                .that(invisibleDueToOcclusion)
+                .isFalse()
+
+            collapseDualShade {
+                assertWithMessage("Should remain unoccluded while shade is collapsing")
+                    .that(invisibleDueToOcclusion)
+                    .isFalse()
+            }
+            assertWithMessage("Should become occluded now that shade is fully collapsed")
+                .that(invisibleDueToOcclusion)
+                .isTrue()
+
+            hideOccludingActivity()
+            assertWithMessage("Should become unoccluded once the occluding activity is hidden")
+                .that(invisibleDueToOcclusion)
+                .isFalse()
+        }
+
     /** Simulates the appearance of a show-when-locked `Activity` in the foreground. */
     private fun TestScope.showOccludingActivity() {
         keyguardOcclusionInteractor.setWmNotifiedShowWhenLockedActivityOnTop(
@@ -138,15 +206,13 @@
     /** Simulates the disappearance of a show-when-locked `Activity` from the foreground. */
     private fun TestScope.hideOccludingActivity() {
         keyguardOcclusionInteractor.setWmNotifiedShowWhenLockedActivityOnTop(
-            showWhenLockedActivityOnTop = false,
+            showWhenLockedActivityOnTop = false
         )
         runCurrent()
     }
 
     /** Simulates a user-driven gradual expansion of the shade. */
-    private fun TestScope.expandShade(
-        assertMidTransition: () -> Unit = {},
-    ) {
+    private fun TestScope.expandShade(assertMidTransition: () -> Unit = {}) {
         val progress = MutableStateFlow(0f)
         mutableTransitionState.value =
             ObservableTransitionState.Transition(
@@ -170,10 +236,41 @@
         runCurrent()
     }
 
+    /** Simulates a user-driven gradual expansion of the dual shade (notifications). */
+    private fun TestScope.expandDualShade(assertMidTransition: () -> Unit = {}) {
+        val progress = MutableStateFlow(0f)
+        mutableTransitionState.value =
+            ShowOrHideOverlay(
+                overlay = Overlays.NotificationsShade,
+                fromContent = sceneDataSource.currentScene.value,
+                toContent = Overlays.NotificationsShade,
+                currentScene = sceneDataSource.currentScene.value,
+                currentOverlays = sceneDataSource.currentOverlays,
+                progress = progress,
+                isInitiatedByUserInput = true,
+                isUserInputOngoing = flowOf(true),
+                previewProgress = flowOf(0f),
+                isInPreviewStage = flowOf(false),
+            )
+        runCurrent()
+
+        progress.value = 0.5f
+        runCurrent()
+        assertMidTransition()
+
+        progress.value = 1f
+        runCurrent()
+
+        mutableTransitionState.value =
+            ObservableTransitionState.Idle(
+                sceneDataSource.currentScene.value,
+                setOf(Overlays.NotificationsShade),
+            )
+        runCurrent()
+    }
+
     /** Simulates a user-driven gradual collapse of the shade. */
-    private fun TestScope.collapseShade(
-        assertMidTransition: () -> Unit = {},
-    ) {
+    private fun TestScope.collapseShade(assertMidTransition: () -> Unit = {}) {
         val progress = MutableStateFlow(0f)
         mutableTransitionState.value =
             ObservableTransitionState.Transition(
@@ -197,10 +294,37 @@
         runCurrent()
     }
 
+    /** Simulates a user-driven gradual collapse of the dual shade (notifications). */
+    private fun TestScope.collapseDualShade(assertMidTransition: () -> Unit = {}) {
+        val progress = MutableStateFlow(0f)
+        mutableTransitionState.value =
+            ShowOrHideOverlay(
+                overlay = Overlays.NotificationsShade,
+                fromContent = Overlays.NotificationsShade,
+                toContent = Scenes.Lockscreen,
+                currentScene = Scenes.Lockscreen,
+                currentOverlays = flowOf(setOf(Overlays.NotificationsShade)),
+                progress = progress,
+                isInitiatedByUserInput = true,
+                isUserInputOngoing = flowOf(true),
+                previewProgress = flowOf(0f),
+                isInPreviewStage = flowOf(false),
+            )
+        runCurrent()
+
+        progress.value = 0.5f
+        runCurrent()
+        assertMidTransition()
+
+        progress.value = 1f
+        runCurrent()
+
+        mutableTransitionState.value = ObservableTransitionState.Idle(Scenes.Lockscreen)
+        runCurrent()
+    }
+
     /** Simulates a transition into AOD. */
-    private suspend fun TestScope.transitionIntoAod(
-        assertMidTransition: () -> Unit = {},
-    ) {
+    private suspend fun TestScope.transitionIntoAod(assertMidTransition: () -> Unit = {}) {
         val currentKeyguardState = keyguardTransitionInteractor.getCurrentState()
         keyguardTransitionRepository.sendTransitionStep(
             TransitionStep(
@@ -235,9 +359,7 @@
     }
 
     /** Simulates a transition away from AOD. */
-    private suspend fun TestScope.transitionOutOfAod(
-        assertMidTransition: () -> Unit = {},
-    ) {
+    private suspend fun TestScope.transitionOutOfAod(assertMidTransition: () -> Unit = {}) {
         keyguardTransitionRepository.sendTransitionStep(
             TransitionStep(
                 from = KeyguardState.AOD,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
index 4a7d8b0..7fe3d8d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
@@ -21,6 +21,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.ObservableTransitionState.Transition.ShowOrHideOverlay
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -38,6 +39,7 @@
 import com.android.systemui.scene.overlayKeys
 import com.android.systemui.scene.sceneContainerConfig
 import com.android.systemui.scene.sceneKeys
+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.fakeSceneDataSource
@@ -256,7 +258,7 @@
         }
 
     @Test
-    fun transitioningTo() =
+    fun transitioningTo_sceneChange() =
         testScope.runTest {
             val transitionState =
                 MutableStateFlow<ObservableTransitionState>(
@@ -293,6 +295,51 @@
         }
 
     @Test
+    fun transitioningTo_overlayChange() =
+        testScope.runTest {
+            val transitionState =
+                MutableStateFlow<ObservableTransitionState>(
+                    ObservableTransitionState.Idle(underTest.currentScene.value)
+                )
+            underTest.setTransitionState(transitionState)
+
+            val transitionTo by collectLastValue(underTest.transitioningTo)
+            assertThat(transitionTo).isNull()
+
+            underTest.showOverlay(Overlays.NotificationsShade, "reason")
+            assertThat(transitionTo).isNull()
+
+            val progress = MutableStateFlow(0f)
+            transitionState.value =
+                ShowOrHideOverlay(
+                    overlay = Overlays.NotificationsShade,
+                    fromContent = underTest.currentScene.value,
+                    toContent = Overlays.NotificationsShade,
+                    currentScene = underTest.currentScene.value,
+                    currentOverlays = underTest.currentOverlays,
+                    progress = progress,
+                    isInitiatedByUserInput = true,
+                    isUserInputOngoing = flowOf(true),
+                    previewProgress = flowOf(0f),
+                    isInPreviewStage = flowOf(false),
+                )
+            assertThat(transitionTo).isEqualTo(Overlays.NotificationsShade)
+
+            progress.value = 0.5f
+            assertThat(transitionTo).isEqualTo(Overlays.NotificationsShade)
+
+            progress.value = 1f
+            assertThat(transitionTo).isEqualTo(Overlays.NotificationsShade)
+
+            transitionState.value =
+                ObservableTransitionState.Idle(
+                    currentScene = underTest.currentScene.value,
+                    currentOverlays = setOf(Overlays.NotificationsShade),
+                )
+            assertThat(transitionTo).isNull()
+        }
+
+    @Test
     fun isTransitionUserInputOngoing_idle_false() =
         testScope.runTest {
             val transitionState =
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImplTest.kt
index 851b7b9..ba559b5 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImplTest.kt
@@ -21,6 +21,8 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.ObservableTransitionState.Transition.ShowOrHideOverlay
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -32,6 +34,7 @@
 import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.fakeSceneDataSource
 import com.android.systemui.testKosmos
@@ -125,6 +128,63 @@
 
     @Test
     @EnableSceneContainer
+    fun legacyPanelExpansion_dualShade_whenIdle_whenLocked() =
+        testScope.runTest {
+            underTest = kosmos.panelExpansionInteractorImpl
+            val panelExpansion by collectLastValue(underTest.legacyPanelExpansion)
+
+            changeScene(Scenes.Lockscreen) { assertThat(panelExpansion).isEqualTo(1f) }
+            assertThat(panelExpansion).isEqualTo(1f)
+
+            changeScene(Scenes.Bouncer) { assertThat(panelExpansion).isEqualTo(1f) }
+            assertThat(panelExpansion).isEqualTo(1f)
+
+            showOverlay(Overlays.NotificationsShade) { assertThat(panelExpansion).isEqualTo(1f) }
+            assertThat(panelExpansion).isEqualTo(1f)
+
+            showOverlay(Overlays.QuickSettingsShade) { assertThat(panelExpansion).isEqualTo(1f) }
+            assertThat(panelExpansion).isEqualTo(1f)
+
+            changeScene(Scenes.Communal) { assertThat(panelExpansion).isEqualTo(1f) }
+            assertThat(panelExpansion).isEqualTo(1f)
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun legacyPanelExpansion_dualShade_whenIdle_whenUnlocked() =
+        testScope.runTest {
+            underTest = kosmos.panelExpansionInteractorImpl
+            val unlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
+            kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
+                SuccessFingerprintAuthenticationStatus(0, true)
+            )
+            runCurrent()
+
+            assertThat(unlockStatus)
+                .isEqualTo(DeviceUnlockStatus(true, DeviceUnlockSource.Fingerprint))
+
+            val panelExpansion by collectLastValue(underTest.legacyPanelExpansion)
+
+            changeScene(Scenes.Gone) { assertThat(panelExpansion).isEqualTo(0f) }
+            assertThat(panelExpansion).isEqualTo(0f)
+
+            showOverlay(Overlays.NotificationsShade) { progress ->
+                assertThat(panelExpansion).isEqualTo(progress)
+            }
+            assertThat(panelExpansion).isEqualTo(1f)
+
+            showOverlay(Overlays.QuickSettingsShade) {
+                // Notification shade is already expanded, so moving to QS shade should also be 1f.
+                assertThat(panelExpansion).isEqualTo(1f)
+            }
+            assertThat(panelExpansion).isEqualTo(1f)
+
+            changeScene(Scenes.Communal) { assertThat(panelExpansion).isEqualTo(1f) }
+            assertThat(panelExpansion).isEqualTo(1f)
+        }
+
+    @Test
+    @EnableSceneContainer
     fun shouldHideStatusBarIconsWhenExpanded_goneScene() =
         testScope.runTest {
             underTest = kosmos.panelExpansionInteractorImpl
@@ -193,4 +253,72 @@
 
         assertThat(currentScene).isEqualTo(toScene)
     }
+
+    private fun TestScope.showOverlay(
+        toOverlay: OverlayKey,
+        assertDuringProgress: ((progress: Float) -> Unit) = {},
+    ) {
+        val currentScene by collectLastValue(sceneInteractor.currentScene)
+        val currentOverlays by collectLastValue(sceneInteractor.currentOverlays)
+        val progressFlow = MutableStateFlow(0f)
+        transitionState.value =
+            if (checkNotNull(currentOverlays).isEmpty()) {
+                ShowOrHideOverlay(
+                    overlay = toOverlay,
+                    fromContent = checkNotNull(currentScene),
+                    toContent = toOverlay,
+                    currentScene = checkNotNull(currentScene),
+                    currentOverlays = flowOf(emptySet()),
+                    progress = progressFlow,
+                    isInitiatedByUserInput = true,
+                    isUserInputOngoing = flowOf(true),
+                    previewProgress = flowOf(0f),
+                    isInPreviewStage = flowOf(false),
+                )
+            } else {
+                ObservableTransitionState.Transition.ReplaceOverlay(
+                    fromOverlay = checkNotNull(currentOverlays).first(),
+                    toOverlay = toOverlay,
+                    currentScene = checkNotNull(currentScene),
+                    currentOverlays = flowOf(emptySet()),
+                    progress = progressFlow,
+                    isInitiatedByUserInput = true,
+                    isUserInputOngoing = flowOf(true),
+                    previewProgress = flowOf(0f),
+                    isInPreviewStage = flowOf(false),
+                )
+            }
+        runCurrent()
+        assertDuringProgress(progressFlow.value)
+
+        progressFlow.value = 0.2f
+        runCurrent()
+        assertDuringProgress(progressFlow.value)
+
+        progressFlow.value = 0.6f
+        runCurrent()
+        assertDuringProgress(progressFlow.value)
+
+        progressFlow.value = 1f
+        runCurrent()
+        assertDuringProgress(progressFlow.value)
+
+        transitionState.value =
+            ObservableTransitionState.Idle(
+                currentScene = checkNotNull(currentScene),
+                currentOverlays = setOf(toOverlay),
+            )
+        if (checkNotNull(currentOverlays).isEmpty()) {
+            fakeSceneDataSource.showOverlay(toOverlay)
+        } else {
+            fakeSceneDataSource.replaceOverlay(
+                from = checkNotNull(currentOverlays).first(),
+                to = toOverlay,
+            )
+        }
+        runCurrent()
+        assertDuringProgress(progressFlow.value)
+
+        assertThat(currentOverlays).containsExactly(toOverlay)
+    }
 }
diff --git a/packages/SystemUI/res-keyguard/values-sw600dp/dimens.xml b/packages/SystemUI/res-keyguard/values-sw600dp/dimens.xml
index 6c8db91..84f7a51 100644
--- a/packages/SystemUI/res-keyguard/values-sw600dp/dimens.xml
+++ b/packages/SystemUI/res-keyguard/values-sw600dp/dimens.xml
@@ -28,4 +28,8 @@
     <!-- Overload default clock widget parameters -->
     <dimen name="widget_big_font_size">100dp</dimen>
     <dimen name="widget_label_font_size">18sp</dimen>
+
+    <!-- New keyboard shortcut helper -->
+    <dimen name="shortcut_helper_width">704dp</dimen>
+    <dimen name="shortcut_helper_height">1208dp</dimen>
 </resources>
diff --git a/packages/SystemUI/res/layout/activity_keyboard_shortcut_helper.xml b/packages/SystemUI/res/layout/activity_keyboard_shortcut_helper.xml
index 06d1bf4..a15532f 100644
--- a/packages/SystemUI/res/layout/activity_keyboard_shortcut_helper.xml
+++ b/packages/SystemUI/res/layout/activity_keyboard_shortcut_helper.xml
@@ -2,14 +2,15 @@
 <androidx.coordinatorlayout.widget.CoordinatorLayout xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:app="http://schemas.android.com/apk/res-auto"
     android:id="@+id/shortcut_helper_sheet_container"
+    android:layout_gravity="center_horizontal|bottom"
     android:layout_width="match_parent"
     android:layout_height="match_parent">
 
     <FrameLayout
         android:id="@+id/shortcut_helper_sheet"
         style="@style/ShortcutHelperBottomSheet"
-        android:layout_width="match_parent"
-        android:layout_height="match_parent"
+        android:layout_width="@dimen/shortcut_helper_width"
+        android:layout_height="@dimen/shortcut_helper_height"
         android:orientation="vertical"
         app:layout_behavior="com.google.android.material.bottomsheet.BottomSheetBehavior">
 
diff --git a/packages/SystemUI/res/values-sw600dp-land/dimens.xml b/packages/SystemUI/res/values-sw600dp-land/dimens.xml
index 2a27b47..3efe7a5 100644
--- a/packages/SystemUI/res/values-sw600dp-land/dimens.xml
+++ b/packages/SystemUI/res/values-sw600dp-land/dimens.xml
@@ -26,6 +26,10 @@
     <dimen name="keyguard_clock_top_margin">8dp</dimen>
     <dimen name="keyguard_smartspace_top_offset">0dp</dimen>
 
+    <!-- New keyboard shortcut helper -->
+    <dimen name="shortcut_helper_width">864dp</dimen>
+    <dimen name="shortcut_helper_height">728dp</dimen>
+
     <!-- QS-->
     <dimen name="qs_panel_padding_top">16dp</dimen>
     <dimen name="qs_panel_padding">24dp</dimen>
diff --git a/packages/SystemUI/res/values/dimens.xml b/packages/SystemUI/res/values/dimens.xml
index e94248d..00846cb 100644
--- a/packages/SystemUI/res/values/dimens.xml
+++ b/packages/SystemUI/res/values/dimens.xml
@@ -1005,6 +1005,10 @@
     <dimen name="ksh_app_item_minimum_height">64dp</dimen>
     <dimen name="ksh_category_separator_margin">16dp</dimen>
 
+    <!-- New keyboard shortcut helper -->
+    <dimen name="shortcut_helper_width">412dp</dimen>
+    <dimen name="shortcut_helper_height">728dp</dimen>
+
     <!-- The size of corner radius of the arrow in the onboarding toast. -->
     <dimen name="recents_onboarding_toast_arrow_corner_radius">2dp</dimen>
 
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
index 1f5878b..a327e4a 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
@@ -28,7 +28,6 @@
 import android.view.Display
 import com.android.app.tracing.FlowTracing.traceEach
 import com.android.app.tracing.traceSection
-import com.android.systemui.Flags
 import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
@@ -51,7 +50,6 @@
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.scan
-import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 
 /** Provides a [Flow] of [Display] as returned by [DisplayManager]. */
@@ -102,7 +100,7 @@
     private val displayManager: DisplayManager,
     @Background backgroundHandler: Handler,
     @Background bgApplicationScope: CoroutineScope,
-    @Background backgroundCoroutineDispatcher: CoroutineDispatcher
+    @Background backgroundCoroutineDispatcher: CoroutineDispatcher,
 ) : DisplayRepository {
     private val allDisplayEvents: Flow<DisplayEvent> =
         conflatedCallbackFlow {
@@ -139,70 +137,55 @@
     override val displayAdditionEvent: Flow<Display?> =
         allDisplayEvents.filterIsInstance<DisplayEvent.Added>().map { getDisplay(it.displayId) }
 
-    // TODO: b/345472038 - Delete after the flag is ramped up.
-    private val oldEnabledDisplays: Flow<Set<Display>> =
-        allDisplayEvents
-            .map { getDisplays() }
-            .shareIn(bgApplicationScope, started = SharingStarted.WhileSubscribed(), replay = 1)
+    // This is necessary because there might be multiple displays, and we could
+    // have missed events for those added before this process or flow started.
+    // Note it causes a binder call from the main thread (it's traced).
+    private val initialDisplays: Set<Display> =
+        traceSection("$TAG#initialDisplays") { displayManager.displays?.toSet() ?: emptySet() }
+    private val initialDisplayIds = initialDisplays.map { display -> display.displayId }.toSet()
 
     /** Propagate to the listeners only enabled displays */
     private val enabledDisplayIds: Flow<Set<Int>> =
-        if (Flags.enableEfficientDisplayRepository()) {
-                allDisplayEvents
-                    .scan(initial = emptySet()) { previousIds: Set<Int>, event: DisplayEvent ->
-                        val id = event.displayId
-                        when (event) {
-                            is DisplayEvent.Removed -> previousIds - id
-                            is DisplayEvent.Added,
-                            is DisplayEvent.Changed -> previousIds + id
-                        }
-                    }
-                    .distinctUntilChanged()
-                    .stateIn(
-                        bgApplicationScope,
-                        SharingStarted.WhileSubscribed(),
-                        // This is necessary because there might be multiple displays, and we could
-                        // have missed events for those added before this process or flow started.
-                        // Note it causes a binder call from the main thread (it's traced).
-                        getDisplays().map { display -> display.displayId }.toSet(),
-                    )
-            } else {
-                oldEnabledDisplays.map { enabledDisplaysSet ->
-                    enabledDisplaysSet.map { it.displayId }.toSet()
+        allDisplayEvents
+            .scan(initial = initialDisplayIds) { previousIds: Set<Int>, event: DisplayEvent ->
+                val id = event.displayId
+                when (event) {
+                    is DisplayEvent.Removed -> previousIds - id
+                    is DisplayEvent.Added,
+                    is DisplayEvent.Changed -> previousIds + id
                 }
             }
+            .distinctUntilChanged()
+            .stateIn(bgApplicationScope, SharingStarted.WhileSubscribed(), initialDisplayIds)
             .debugLog("enabledDisplayIds")
 
     private val defaultDisplay by lazy {
         getDisplay(Display.DEFAULT_DISPLAY) ?: error("Unable to get default display.")
     }
+
     /**
      * Represents displays that went though the [DisplayListener.onDisplayAdded] callback.
      *
      * Those are commonly the ones provided by [DisplayManager.getDisplays] by default.
      */
     private val enabledDisplays: Flow<Set<Display>> =
-        if (Flags.enableEfficientDisplayRepository()) {
-            enabledDisplayIds
-                .mapElementsLazily { displayId -> getDisplay(displayId) }
-                .onEach {
-                    if (it.isEmpty()) Log.wtf(TAG, "No enabled displays. This should never happen.")
-                }
-                .flowOn(backgroundCoroutineDispatcher)
-                .debugLog("enabledDisplays")
-                .stateIn(
-                    bgApplicationScope,
-                    started = SharingStarted.WhileSubscribed(),
-                    // This triggers a single binder call on the UI thread per process. The
-                    // alternative would be to use sharedFlows, but they are prohibited due to
-                    // performance concerns.
-                    // Ultimately, this is a trade-off between a one-time UI thread binder call and
-                    // the constant overhead of sharedFlows.
-                    initialValue = getDisplays()
-                )
-        } else {
-            oldEnabledDisplays
-        }
+        enabledDisplayIds
+            .mapElementsLazily { displayId -> getDisplay(displayId) }
+            .onEach {
+                if (it.isEmpty()) Log.wtf(TAG, "No enabled displays. This should never happen.")
+            }
+            .flowOn(backgroundCoroutineDispatcher)
+            .debugLog("enabledDisplays")
+            .stateIn(
+                bgApplicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                // This triggers a single binder call on the UI thread per process. The
+                // alternative would be to use sharedFlows, but they are prohibited due to
+                // performance concerns.
+                // Ultimately, this is a trade-off between a one-time UI thread binder call and
+                // the constant overhead of sharedFlows.
+                initialValue = initialDisplays,
+            )
 
     /**
      * Represents displays that went though the [DisplayListener.onDisplayAdded] callback.
@@ -211,10 +194,7 @@
      */
     override val displays: Flow<Set<Display>> = enabledDisplays
 
-    private fun getDisplays(): Set<Display> =
-        traceSection("$TAG#getDisplays()") { displayManager.displays?.toSet() ?: emptySet() }
-
-    private val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
+    val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
     private val ignoredDisplayIds: Flow<Set<Int>> = _ignoredDisplayIds.debugLog("ignoredDisplayIds")
 
     private fun getInitialConnectedDisplays(): Set<Int> =
@@ -271,7 +251,7 @@
                 // the flow starts being collected. This is to ensure the call to get displays (an
                 // IPC) happens in the background instead of when this object
                 // is instantiated.
-                initialValue = emptySet()
+                initialValue = emptySet(),
             )
 
     private val connectedExternalDisplayIds: Flow<Set<Int>> =
@@ -308,7 +288,7 @@
                         TAG,
                         "combining enabled=$enabledDisplaysIds, " +
                             "connectedExternalDisplayIds=$connectedExternalDisplayIds, " +
-                            "ignored=$ignoredDisplayIds"
+                            "ignored=$ignoredDisplayIds",
                     )
                 }
                 connectedExternalDisplayIds - enabledDisplaysIds - ignoredDisplayIds
@@ -382,7 +362,7 @@
             val previousSet: Set<T>,
             // Caches T values from the previousSet that were already converted to V
             val valueMap: Map<T, V>,
-            val resultSet: Set<V>
+            val resultSet: Set<V>,
         )
 
         val emptyInitialState = State(emptySet<T>(), emptyMap(), emptySet<V>())
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
index beec348..11a0543 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
@@ -144,7 +144,7 @@
                 useSinglePane,
                 onSearchQueryChanged,
                 modifier,
-                onKeyboardSettingsClicked
+                onKeyboardSettingsClicked,
             )
         }
         else -> {
@@ -159,7 +159,7 @@
     useSinglePane: @Composable () -> Boolean,
     onSearchQueryChanged: (String) -> Unit,
     modifier: Modifier,
-    onKeyboardSettingsClicked: () -> Unit
+    onKeyboardSettingsClicked: () -> Unit,
 ) {
     var selectedCategoryType by
         remember(shortcutsUiState.defaultSelectedCategory) {
@@ -183,7 +183,7 @@
             shortcutsUiState.shortcutCategories,
             selectedCategoryType,
             onCategorySelected = { selectedCategoryType = it },
-            onKeyboardSettingsClicked
+            onKeyboardSettingsClicked,
         )
     }
 }
@@ -223,14 +223,14 @@
                 searchQuery,
                 categories,
                 selectedCategoryType,
-                onCategorySelected
+                onCategorySelected,
             )
             Spacer(modifier = Modifier.weight(1f))
         }
         KeyboardSettings(
             horizontalPadding = 16.dp,
             verticalPadding = 32.dp,
-            onClick = onKeyboardSettingsClicked
+            onClick = onKeyboardSettingsClicked,
         )
     }
 }
@@ -282,11 +282,7 @@
     onClick: () -> Unit,
     shape: Shape,
 ) {
-    Surface(
-        color = MaterialTheme.colorScheme.surfaceBright,
-        shape = shape,
-        onClick = onClick,
-    ) {
+    Surface(color = MaterialTheme.colorScheme.surfaceBright, shape = shape, onClick = onClick) {
         Column {
             Row(
                 verticalAlignment = Alignment.CenterVertically,
@@ -327,7 +323,7 @@
     source: IconSource,
     modifier: Modifier = Modifier,
     contentDescription: String? = null,
-    tint: Color = LocalContentColor.current
+    tint: Color = LocalContentColor.current,
 ) {
     if (source.imageVector != null) {
         Icon(source.imageVector, contentDescription, modifier, tint)
@@ -350,7 +346,7 @@
 
 private fun getApplicationLabelForCurrentApp(
     type: ShortcutCategoryType.CurrentApp,
-    context: Context
+    context: Context,
 ): String {
     val packageManagerForUser = CentralSurfaces.getPackageManagerForUser(context, context.userId)
     return try {
@@ -358,7 +354,7 @@
             packageManagerForUser.getApplicationInfoAsUser(
                 type.packageName,
                 /* flags = */ 0,
-                context.userId
+                context.userId,
             )
         packageManagerForUser.getApplicationLabel(currentAppInfo).toString()
     } catch (e: NameNotFoundException) {
@@ -377,13 +373,13 @@
                 } else {
                     0f
                 },
-            label = "Expand icon rotation animation"
+            label = "Expand icon rotation animation",
         )
     Icon(
         modifier =
             Modifier.background(
                     color = MaterialTheme.colorScheme.surfaceContainerHigh,
-                    shape = CircleShape
+                    shape = CircleShape,
                 )
                 .graphicsLayer { rotationZ = expandIconRotationDegrees },
         imageVector = Icons.Default.ExpandMore,
@@ -393,7 +389,7 @@
             } else {
                 stringResource(R.string.shortcut_helper_content_description_expand_icon)
             },
-        tint = MaterialTheme.colorScheme.onSurface
+        tint = MaterialTheme.colorScheme.onSurface,
     )
 }
 
@@ -435,11 +431,11 @@
         Row(Modifier.fillMaxWidth()) {
             StartSidePanel(
                 onSearchQueryChanged = onSearchQueryChanged,
-                modifier = Modifier.width(200.dp),
+                modifier = Modifier.width(240.dp),
                 categories = categories,
                 onKeyboardSettingsClicked = onKeyboardSettingsClicked,
                 selectedCategory = selectedCategoryType,
-                onCategoryClicked = { onCategorySelected(it.type) }
+                onCategoryClicked = { onCategorySelected(it.type) },
             )
             Spacer(modifier = Modifier.width(24.dp))
             EndSidePanel(searchQuery, Modifier.fillMaxSize().padding(top = 8.dp), selectedCategory)
@@ -475,7 +471,7 @@
             modifier
                 .padding(vertical = 8.dp)
                 .background(MaterialTheme.colorScheme.surfaceBright, RoundedCornerShape(28.dp))
-                .padding(horizontal = horizontalPadding, vertical = 24.dp)
+                .padding(horizontal = horizontalPadding, vertical = 24.dp),
     )
 }
 
@@ -484,7 +480,7 @@
     Surface(
         modifier = Modifier.fillMaxWidth(),
         shape = RoundedCornerShape(28.dp),
-        color = MaterialTheme.colorScheme.surfaceBright
+        color = MaterialTheme.colorScheme.surfaceBright,
     ) {
         Column(Modifier.padding(24.dp)) {
             SubCategoryTitle(subCategory.label)
@@ -519,7 +515,7 @@
                 isFocused = isFocused,
                 focusColor = MaterialTheme.colorScheme.secondary,
                 padding = 8.dp,
-                cornerRadius = 16.dp
+                cornerRadius = 16.dp,
             )
     ) {
         Row(
@@ -528,21 +524,12 @@
             verticalAlignment = Alignment.CenterVertically,
         ) {
             if (shortcut.icon != null) {
-                ShortcutIcon(
-                    shortcut.icon,
-                    modifier = Modifier.size(24.dp),
-                )
+                ShortcutIcon(shortcut.icon, modifier = Modifier.size(24.dp))
             }
-            ShortcutDescriptionText(
-                searchQuery = searchQuery,
-                shortcut = shortcut,
-            )
+            ShortcutDescriptionText(searchQuery = searchQuery, shortcut = shortcut)
         }
         Spacer(modifier = Modifier.width(16.dp))
-        ShortcutKeyCombinations(
-            modifier = Modifier.weight(1f),
-            shortcut = shortcut,
-        )
+        ShortcutKeyCombinations(modifier = Modifier.weight(1f), shortcut = shortcut)
     }
 }
 
@@ -566,14 +553,11 @@
 
 @OptIn(ExperimentalLayoutApi::class)
 @Composable
-private fun ShortcutKeyCombinations(
-    modifier: Modifier = Modifier,
-    shortcut: Shortcut,
-) {
+private fun ShortcutKeyCombinations(modifier: Modifier = Modifier, shortcut: Shortcut) {
     FlowRow(
         modifier = modifier,
         verticalArrangement = Arrangement.spacedBy(8.dp),
-        horizontalArrangement = Arrangement.End
+        horizontalArrangement = Arrangement.End,
     ) {
         shortcut.commands.forEachIndexed { index, command ->
             if (index > 0) {
@@ -609,8 +593,8 @@
             Modifier.height(36.dp)
                 .background(
                     color = MaterialTheme.colorScheme.surfaceContainer,
-                    shape = RoundedCornerShape(12.dp)
-                ),
+                    shape = RoundedCornerShape(12.dp),
+                )
     ) {
         shortcutKeyContent()
     }
@@ -630,7 +614,7 @@
     Icon(
         painter = painterResource(key.drawableResId),
         contentDescription = null,
-        modifier = Modifier.align(Alignment.Center).padding(6.dp)
+        modifier = Modifier.align(Alignment.Center).padding(6.dp),
     )
 }
 
@@ -701,7 +685,7 @@
         KeyboardSettings(
             horizontalPadding = 24.dp,
             verticalPadding = 24.dp,
-            onKeyboardSettingsClicked
+            onKeyboardSettingsClicked,
         )
     }
 }
@@ -710,7 +694,7 @@
 private fun CategoriesPanelTwoPane(
     categories: List<ShortcutCategory>,
     selectedCategory: ShortcutCategoryType?,
-    onCategoryClicked: (ShortcutCategory) -> Unit
+    onCategoryClicked: (ShortcutCategory) -> Unit,
 ) {
     Column {
         categories.fastForEach {
@@ -718,7 +702,7 @@
                 label = it.label(LocalContext.current),
                 iconSource = it.icon,
                 selected = selectedCategory == it.type,
-                onClick = { onCategoryClicked(it) }
+                onClick = { onCategoryClicked(it) },
             )
         }
     }
@@ -747,7 +731,7 @@
                     isFocused = isFocused,
                     focusColor = MaterialTheme.colorScheme.secondary,
                     padding = 2.dp,
-                    cornerRadius = 33.dp
+                    cornerRadius = 33.dp,
                 ),
         shape = RoundedCornerShape(28.dp),
         color = colors.containerColor(selected).value,
@@ -758,7 +742,7 @@
                 modifier = Modifier.size(24.dp),
                 source = iconSource,
                 contentDescription = null,
-                tint = colors.iconColor(selected).value
+                tint = colors.iconColor(selected).value,
             )
             Spacer(Modifier.width(12.dp))
             Box(Modifier.weight(1f)) {
@@ -766,7 +750,7 @@
                     fontSize = 18.sp,
                     color = colors.textColor(selected).value,
                     style = MaterialTheme.typography.headlineSmall,
-                    text = label
+                    text = label,
                 )
             }
         }
@@ -777,7 +761,7 @@
     isFocused: Boolean,
     focusColor: Color,
     padding: Dp,
-    cornerRadius: Dp
+    cornerRadius: Dp,
 ): Modifier {
     if (isFocused) {
         return this.drawWithContent {
@@ -795,7 +779,7 @@
                     style = Stroke(width = 3.dp.toPx()),
                     topLeft = focusOutline.topLeft,
                     size = focusOutline.size,
-                    cornerRadius = CornerRadius(cornerRadius.toPx())
+                    cornerRadius = CornerRadius(cornerRadius.toPx()),
                 )
             }
             // Increasing Z-Index so focus outline is drawn on top of "selected" category
@@ -815,9 +799,9 @@
             Text(
                 text = stringResource(R.string.shortcut_helper_title),
                 color = MaterialTheme.colorScheme.onSurface,
-                style = MaterialTheme.typography.headlineSmall
+                style = MaterialTheme.typography.headlineSmall,
             )
-        }
+        },
     )
 }
 
@@ -852,7 +836,7 @@
         onSearch = {},
         leadingIcon = { Icon(Icons.Default.Search, contentDescription = null) },
         placeholder = { Text(text = stringResource(R.string.shortcut_helper_search_placeholder)) },
-        content = {}
+        content = {},
     )
 }
 
@@ -874,21 +858,21 @@
                         isFocused = isFocused,
                         focusColor = MaterialTheme.colorScheme.secondary,
                         padding = 8.dp,
-                        cornerRadius = 28.dp
+                        cornerRadius = 28.dp,
                     ),
-            verticalAlignment = Alignment.CenterVertically
+            verticalAlignment = Alignment.CenterVertically,
         ) {
             Text(
                 "Keyboard Settings",
                 color = MaterialTheme.colorScheme.onSurfaceVariant,
-                fontSize = 16.sp
+                fontSize = 16.sp,
             )
             Spacer(modifier = Modifier.weight(1f))
             Icon(
                 imageVector = Icons.AutoMirrored.Default.OpenInNew,
                 contentDescription = null,
                 tint = MaterialTheme.colorScheme.onSurfaceVariant,
-                modifier = Modifier.size(24.dp)
+                modifier = Modifier.size(24.dp),
             )
         }
     }
@@ -900,17 +884,15 @@
         val singlePaneFirstCategory =
             RoundedCornerShape(
                 topStart = Dimensions.SinglePaneCategoryCornerRadius,
-                topEnd = Dimensions.SinglePaneCategoryCornerRadius
+                topEnd = Dimensions.SinglePaneCategoryCornerRadius,
             )
         val singlePaneLastCategory =
             RoundedCornerShape(
                 bottomStart = Dimensions.SinglePaneCategoryCornerRadius,
-                bottomEnd = Dimensions.SinglePaneCategoryCornerRadius
+                bottomEnd = Dimensions.SinglePaneCategoryCornerRadius,
             )
         val singlePaneSingleCategory =
-            RoundedCornerShape(
-                size = Dimensions.SinglePaneCategoryCornerRadius,
-            )
+            RoundedCornerShape(size = Dimensions.SinglePaneCategoryCornerRadius)
         val singlePaneCategory = RectangleShape
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
index 2039743..799999a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
@@ -39,6 +39,7 @@
 import com.android.systemui.keyboard.shortcut.ui.viewmodel.ShortcutHelperViewModel
 import com.android.systemui.res.R
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.util.dpToPx
 import com.google.android.material.bottomsheet.BottomSheetBehavior
 import com.google.android.material.bottomsheet.BottomSheetBehavior.BottomSheetCallback
 import com.google.android.material.bottomsheet.BottomSheetBehavior.STATE_HIDDEN
@@ -51,10 +52,8 @@
  */
 class ShortcutHelperActivity
 @Inject
-constructor(
-    private val userTracker: UserTracker,
-    private val viewModel: ShortcutHelperViewModel,
-) : ComponentActivity() {
+constructor(private val userTracker: UserTracker, private val viewModel: ShortcutHelperViewModel) :
+    ComponentActivity() {
 
     private val bottomSheetContainer
         get() = requireViewById<View>(R.id.shortcut_helper_sheet_container)
@@ -69,7 +68,7 @@
         setupEdgeToEdge()
         super.onCreate(savedInstanceState)
         setContentView(R.layout.activity_keyboard_shortcut_helper)
-        setUpBottomSheetWidth()
+        setUpWidth()
         expandBottomSheet()
         setUpInsets()
         setUpPredictiveBack()
@@ -80,6 +79,13 @@
         viewModel.onViewOpened()
     }
 
+    private fun setUpWidth() {
+        // we override this because when maxWidth isn't specified, material imposes a max width
+        // constraint on bottom sheets on larger screens which is smaller than our desired width.
+        bottomSheetBehavior.maxWidth =
+            resources.getDimension(R.dimen.shortcut_helper_width).dpToPx(resources).toInt()
+    }
+
     private fun setUpComposeView() {
         requireViewById<ComposeView>(R.id.shortcut_helper_compose_container).apply {
             setContent {
@@ -102,7 +108,7 @@
         try {
             startActivityAsUser(
                 Intent(Settings.ACTION_HARD_KEYBOARD_SETTINGS),
-                userTracker.userHandle
+                userTracker.userHandle,
             )
         } catch (e: ActivityNotFoundException) {
             // From the Settings docs: In some cases, a matching Activity may not exist, so ensure
@@ -133,15 +139,6 @@
         window.setDecorFitsSystemWindows(false)
     }
 
-    private fun setUpBottomSheetWidth() {
-        val sheetScreenWidthFraction =
-            resources.getFloat(R.dimen.shortcut_helper_screen_width_fraction)
-        // maxWidth needs to be set before the sheet is drawn, otherwise the call will have no
-        // effect.
-        val screenWidth = windowManager.maximumWindowMetrics.bounds.width()
-        bottomSheetBehavior.maxWidth = (sheetScreenWidthFraction * screenWidth).toInt()
-    }
-
     private fun setUpInsets() {
         bottomSheetContainer.setOnApplyWindowInsetsListener { _, insets ->
             val safeDrawingInsets = insets.safeDrawing
@@ -153,7 +150,7 @@
             bottomSheet.updatePadding(
                 left = safeDrawingInsets.left,
                 right = safeDrawingInsets.right,
-                bottom = safeDrawingInsets.bottom
+                bottom = safeDrawingInsets.bottom,
             )
             // The bottom sheet has to be expanded only after setting up insets, otherwise there is
             // a bug and it will not use full height.
@@ -191,7 +188,7 @@
             }
         onBackPressedDispatcher.addCallback(
             owner = this,
-            onBackPressedCallback = onBackPressedCallback
+            onBackPressedCallback = onBackPressedCallback,
         )
     }
 
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/scene/domain/interactor/SceneContainerOcclusionInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
index 429b47b..667827a 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
@@ -16,13 +16,14 @@
 
 package com.android.systemui.scene.domain.interactor
 
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.ObservableTransitionState
-import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.domain.interactor.KeyguardOcclusionInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.Scenes
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
@@ -117,28 +118,30 @@
     private val ObservableTransitionState.canBeOccluded: Boolean
         get() =
             when (this) {
-                is ObservableTransitionState.Idle -> currentScene.canBeOccluded
-                is ObservableTransitionState.Transition.ChangeScene ->
-                    fromScene.canBeOccluded && toScene.canBeOccluded
-                is ObservableTransitionState.Transition.ReplaceOverlay,
-                is ObservableTransitionState.Transition.ShowOrHideOverlay ->
-                    TODO("b/359173565: Handle overlay transitions")
+                is ObservableTransitionState.Idle ->
+                    currentOverlays.all { it.canBeOccluded } && currentScene.canBeOccluded
+                is ObservableTransitionState.Transition ->
+                    // TODO(b/356596436): Should also verify currentOverlays.isEmpty(), but
+                    //  currentOverlays is a Flow and we need a state.
+                    fromContent.canBeOccluded && toContent.canBeOccluded
             }
 
     /**
-     * Whether the scene can be occluded by a "show when locked" activity. Some scenes should, on
+     * Whether the content can be occluded by a "show when locked" activity. Some content should, on
      * principle not be occlude-able because they render as if they are expanding on top of the
      * occluding activity.
      */
-    private val SceneKey.canBeOccluded: Boolean
+    private val ContentKey.canBeOccluded: Boolean
         get() =
             when (this) {
+                Overlays.NotificationsShade -> false
+                Overlays.QuickSettingsShade -> false
                 Scenes.Bouncer -> false
                 Scenes.Communal -> true
                 Scenes.Gone -> true
                 Scenes.Lockscreen -> true
                 Scenes.QuickSettings -> false
                 Scenes.Shade -> false
-                else -> error("SceneKey \"$this\" doesn't have a mapping for canBeOccluded!")
+                else -> error("ContentKey \"$this\" doesn't have a mapping for canBeOccluded!")
             }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
index 0d24adc..f20e5a5 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
@@ -120,21 +120,18 @@
             )
 
     /**
-     * The key of the scene that the UI is currently transitioning to or `null` if there is no
+     * The key of the content that the UI is currently transitioning to or `null` if there is no
      * active transition at the moment.
      *
      * This is a convenience wrapper around [transitionState], meant for flow-challenged consumers
      * like Java code.
      */
-    val transitioningTo: StateFlow<SceneKey?> =
+    val transitioningTo: StateFlow<ContentKey?> =
         transitionState
             .map { state ->
                 when (state) {
                     is ObservableTransitionState.Idle -> null
-                    is ObservableTransitionState.Transition.ChangeScene -> state.toScene
-                    is ObservableTransitionState.Transition.ShowOrHideOverlay,
-                    is ObservableTransitionState.Transition.ReplaceOverlay ->
-                        TODO("b/359173565: Handle overlay transitions")
+                    is ObservableTransitionState.Transition -> state.toContent
                 }
             }
             .stateIn(
@@ -160,15 +157,14 @@
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue = false
+                initialValue = false,
             )
 
     /** Whether the scene container is visible. */
     val isVisible: StateFlow<Boolean> =
-        combine(
-                repository.isVisible,
-                repository.isRemoteUserInputOngoing,
-            ) { isVisible, isRemoteUserInteractionOngoing ->
+        combine(repository.isVisible, repository.isRemoteUserInputOngoing) {
+                isVisible,
+                isRemoteUserInteractionOngoing ->
                 isVisibleInternal(
                     raw = isVisible,
                     isRemoteUserInputOngoing = isRemoteUserInteractionOngoing,
@@ -177,7 +173,7 @@
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue = isVisibleInternal()
+                initialValue = isVisibleInternal(),
             )
 
     /** Whether there's an ongoing remotely-initiated user interaction. */
@@ -259,10 +255,7 @@
      * The change is instantaneous and not animated; it will be observable in the next frame and
      * there will be no transition animation.
      */
-    fun snapToScene(
-        toScene: SceneKey,
-        loggingReason: String,
-    ) {
+    fun snapToScene(toScene: SceneKey, loggingReason: String) {
         val currentSceneKey = currentScene.value
         val resolvedScene =
             sceneFamilyResolvers.get()[toScene]?.let { familyResolver ->
@@ -313,15 +306,9 @@
             return
         }
 
-        logger.logOverlayChangeRequested(
-            to = overlay,
-            reason = loggingReason,
-        )
+        logger.logOverlayChangeRequested(to = overlay, reason = loggingReason)
 
-        repository.showOverlay(
-            overlay = overlay,
-            transitionKey = transitionKey,
-        )
+        repository.showOverlay(overlay = overlay, transitionKey = transitionKey)
     }
 
     /**
@@ -345,15 +332,9 @@
             return
         }
 
-        logger.logOverlayChangeRequested(
-            from = overlay,
-            reason = loggingReason,
-        )
+        logger.logOverlayChangeRequested(from = overlay, reason = loggingReason)
 
-        repository.hideOverlay(
-            overlay = overlay,
-            transitionKey = transitionKey,
-        )
+        repository.hideOverlay(overlay = overlay, transitionKey = transitionKey)
     }
 
     /**
@@ -378,17 +359,9 @@
             return
         }
 
-        logger.logOverlayChangeRequested(
-            from = from,
-            to = to,
-            reason = loggingReason,
-        )
+        logger.logOverlayChangeRequested(from = from, to = to, reason = loggingReason)
 
-        repository.replaceOverlay(
-            from = from,
-            to = to,
-            transitionKey = transitionKey,
-        )
+        repository.replaceOverlay(from = from, to = to, transitionKey = transitionKey)
     }
 
     /**
@@ -405,11 +378,7 @@
             return
         }
 
-        logger.logVisibilityChange(
-            from = wasVisible,
-            to = isVisible,
-            reason = loggingReason,
-        )
+        logger.logVisibilityChange(from = wasVisible, to = isVisible, reason = loggingReason)
         return repository.setVisible(isVisible)
     }
 
@@ -491,11 +460,7 @@
      * @param loggingReason The reason why the transition is requested, for logging purposes
      * @return `true` if the scene change is valid; `false` if it shouldn't happen
      */
-    private fun validateSceneChange(
-        from: SceneKey,
-        to: SceneKey,
-        loggingReason: String,
-    ): Boolean {
+    private fun validateSceneChange(from: SceneKey, to: SceneKey, loggingReason: String): Boolean {
         if (to !in repository.allContentKeys) {
             return false
         }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
index c451704..af1f5a7 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
@@ -52,7 +52,7 @@
     private val sceneInteractor: SceneInteractor,
     private val falsingInteractor: FalsingInteractor,
     private val powerInteractor: PowerInteractor,
-    private val shadeInteractor: ShadeInteractor,
+    shadeInteractor: ShadeInteractor,
     private val splitEdgeDetector: SplitEdgeDetector,
     private val logger: SceneLogger,
     @Assisted private val motionEventHandlerReceiver: (MotionEventHandler?) -> Unit,
@@ -212,9 +212,10 @@
                         )
                     }
                 }
+                // Overlay transitions don't use scene families, nothing to resolve.
                 is UserActionResult.ShowOverlay,
                 is UserActionResult.HideOverlay,
-                is UserActionResult.ReplaceByOverlay -> TODO("b/353679003: Support overlays")
+                is UserActionResult.ReplaceByOverlay -> null
             } ?: actionResult
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
index e276f88..cea521f 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
@@ -18,10 +18,11 @@
 
 package com.android.systemui.shade.domain.interactor
 
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.ObservableTransitionState
-import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.SysuiStatusBarStateController
 import javax.inject.Inject
@@ -55,7 +56,9 @@
             when (state) {
                 is ObservableTransitionState.Idle ->
                     flowOf(
-                        if (state.currentScene != Scenes.Gone) {
+                        if (
+                            state.currentScene != Scenes.Gone || state.currentOverlays.isNotEmpty()
+                        ) {
                             // When resting on a non-Gone scene, the panel is fully expanded.
                             1f
                         } else {
@@ -64,10 +67,10 @@
                             0f
                         }
                     )
-                is ObservableTransitionState.Transition.ChangeScene ->
+                is ObservableTransitionState.Transition ->
                     when {
-                        state.fromScene == Scenes.Gone ->
-                            if (state.toScene.isExpandable()) {
+                        state.fromContent == Scenes.Gone ->
+                            if (state.toContent.isExpandable()) {
                                 // Moving from Gone to a scene that can animate-expand has a
                                 // panel expansion that tracks with the transition.
                                 state.progress
@@ -76,8 +79,8 @@
                                 // immediately makes the panel fully expanded.
                                 flowOf(1f)
                             }
-                        state.toScene == Scenes.Gone ->
-                            if (state.fromScene.isExpandable()) {
+                        state.toContent == Scenes.Gone ->
+                            if (state.fromContent.isExpandable()) {
                                 // Moving to Gone from a scene that can animate-expand has a
                                 // panel expansion that tracks with the transition.
                                 state.progress.map { 1 - it }
@@ -88,9 +91,6 @@
                             }
                         else -> flowOf(1f)
                     }
-                is ObservableTransitionState.Transition.ShowOrHideOverlay,
-                is ObservableTransitionState.Transition.ReplaceOverlay ->
-                    TODO("b/359173565: Handle overlay transitions")
             }
         }
 
@@ -132,7 +132,13 @@
         return sceneInteractor.currentScene.value == Scenes.Lockscreen
     }
 
-    private fun SceneKey.isExpandable(): Boolean {
-        return this == Scenes.Shade || this == Scenes.QuickSettings
+    private fun ContentKey.isExpandable(): Boolean {
+        return when (this) {
+            Scenes.Shade,
+            Scenes.QuickSettings,
+            Overlays.NotificationsShade,
+            Overlays.QuickSettingsShade -> true
+            else -> false
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
index 3e42413..8d7007b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
@@ -17,6 +17,7 @@
 
 package com.android.systemui.statusbar.notification.stack.ui.viewmodel
 
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.ObservableTransitionState.Idle
 import com.android.compose.animation.scene.ObservableTransitionState.Transition
 import com.android.compose.animation.scene.ObservableTransitionState.Transition.ChangeScene
@@ -26,7 +27,7 @@
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
-import com.android.systemui.scene.shared.model.SceneFamilies
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.shade.shared.model.ShadeMode
@@ -46,7 +47,6 @@
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapNotNull
 
 /** ViewModel which represents the state of the NSSL/Controller in the world of flexiglass */
@@ -91,7 +91,7 @@
     ): Float {
         return if (fullyExpandedDuringSceneChange(change)) {
             1f
-        } else if (change.isBetween({ it == Scenes.Gone }, { it in SceneFamilies.NotifShade })) {
+        } else if (change.isBetween({ it == Scenes.Gone }, { it == Scenes.Shade })) {
             shadeExpansion
         } else if (change.isBetween({ it == Scenes.Gone }, { it == Scenes.QuickSettings })) {
             // during QS expansion, increase fraction at same rate as scrim alpha,
@@ -99,6 +99,22 @@
             (qsExpansion / EXPANSION_FOR_MAX_SCRIM_ALPHA - EXPANSION_FOR_DELAYED_STACK_FADE_IN)
                 .coerceIn(0f, 1f)
         } else {
+            // TODO(b/356596436): If notification shade overlay is open, we'll reach this point and
+            //  the expansion fraction in that case should be `shadeExpansion`.
+            0f
+        }
+    }
+
+    private fun expandFractionDuringOverlayTransition(
+        transition: Transition,
+        currentScene: SceneKey,
+        shadeExpansion: Float,
+    ): Float {
+        return if (currentScene == Scenes.Lockscreen) {
+            1f
+        } else if (transition.isTransitioningFromOrTo(Overlays.NotificationsShade)) {
+            shadeExpansion
+        } else {
             0f
         }
     }
@@ -114,18 +130,35 @@
                 shadeInteractor.shadeMode,
                 shadeInteractor.qsExpansion,
                 sceneInteractor.transitionState,
-                sceneInteractor.resolveSceneFamily(SceneFamilies.QuickSettings),
-            ) { shadeExpansion, _, qsExpansion, transitionState, _ ->
+            ) { shadeExpansion, _, qsExpansion, transitionState ->
                 when (transitionState) {
-                    is Idle -> if (expandedInScene(transitionState.currentScene)) 1f else 0f
+                    is Idle ->
+                        if (
+                            expandedInScene(transitionState.currentScene) ||
+                                Overlays.NotificationsShade in transitionState.currentOverlays
+                        ) {
+                            1f
+                        } else {
+                            0f
+                        }
                     is ChangeScene ->
                         expandFractionDuringSceneChange(
-                            transitionState,
-                            shadeExpansion,
-                            qsExpansion,
+                            change = transitionState,
+                            shadeExpansion = shadeExpansion,
+                            qsExpansion = qsExpansion,
                         )
-                    is Transition.ShowOrHideOverlay,
-                    is Transition.ReplaceOverlay -> TODO("b/359173565: Handle overlay transitions")
+                    is Transition.ShowOrHideOverlay ->
+                        expandFractionDuringOverlayTransition(
+                            transition = transitionState,
+                            currentScene = transitionState.currentScene,
+                            shadeExpansion = shadeExpansion,
+                        )
+                    is Transition.ReplaceOverlay ->
+                        expandFractionDuringOverlayTransition(
+                            transition = transitionState,
+                            currentScene = transitionState.currentScene,
+                            shadeExpansion = shadeExpansion,
+                        )
                 }
             }
             .distinctUntilChanged()
@@ -166,14 +199,14 @@
 
     fun shadeScrimShape(
         cornerRadius: Flow<Int>,
-        viewLeftOffset: Flow<Int>
+        viewLeftOffset: Flow<Int>,
     ): Flow<ShadeScrimShape?> =
         combine(shadeScrimClipping, cornerRadius, viewLeftOffset) { clipping, radius, leftOffset ->
                 if (clipping == null) return@combine null
                 ShadeScrimShape(
                     bounds = clipping.bounds.minus(leftOffset = leftOffset),
                     topRadius = radius.takeIf { clipping.rounding.isTopRounded } ?: 0,
-                    bottomRadius = radius.takeIf { clipping.rounding.isBottomRounded } ?: 0
+                    bottomRadius = radius.takeIf { clipping.rounding.isBottomRounded } ?: 0,
                 )
             }
             .dumpWhileCollecting("shadeScrimShape")
@@ -209,10 +242,10 @@
 
     /** Whether the notification stack is scrollable or not. */
     val isScrollable: Flow<Boolean> =
-        sceneInteractor.currentScene
-            .map {
-                sceneInteractor.isSceneInFamily(it, SceneFamilies.NotifShade) ||
-                    it == Scenes.Lockscreen
+        combine(sceneInteractor.currentScene, sceneInteractor.currentOverlays) {
+                currentScene,
+                currentOverlays ->
+                currentScene.showsNotifications() || currentOverlays.any { it.showsNotifications() }
             }
             .dumpWhileCollecting("isScrollable")
 
@@ -242,13 +275,20 @@
         }
     }
 
+    private fun ContentKey.showsNotifications(): Boolean {
+        return when (this) {
+            Overlays.NotificationsShade,
+            Scenes.Lockscreen,
+            Scenes.Shade -> true
+            else -> false
+        }
+    }
+
     @AssistedFactory
     interface Factory {
         fun create(): NotificationScrollViewModel
     }
 }
 
-private fun ChangeScene.isBetween(
-    a: (SceneKey) -> Boolean,
-    b: (SceneKey) -> Boolean,
-): Boolean = (a(fromScene) && b(toScene)) || (b(fromScene) && a(toScene))
+private fun ChangeScene.isBetween(a: (SceneKey) -> Boolean, b: (SceneKey) -> Boolean): Boolean =
+    (a(fromScene) && b(toScene)) || (b(fromScene) && a(toScene))
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManager.java
index 45aee5b..a658115 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManager.java
@@ -30,6 +30,7 @@
 import android.view.WindowInsets;
 
 import androidx.annotation.VisibleForTesting;
+
 import com.android.compose.animation.scene.ObservableTransitionState;
 import com.android.internal.policy.SystemBarUtils;
 import com.android.systemui.Dumpable;
@@ -69,7 +70,9 @@
     private final UnlockedScreenOffAnimationController mUnlockedScreenOffAnimationController;
 
     private boolean mIsStatusBarExpanded = false;
-    private boolean mIsIdleOnGone = true;
+    // Whether the scene container has no UI to render, i.e. is in idle state on the Gone scene and
+    // without any overlays to display.
+    private boolean mIsSceneContainerUiEmpty = true;
     private boolean mIsRemoteUserInteractionOngoing = false;
     private boolean mShouldAdjustInsets = false;
     private View mNotificationShadeWindowView;
@@ -134,7 +137,7 @@
         if (SceneContainerFlag.isEnabled()) {
             javaAdapter.alwaysCollectFlow(
                     sceneInteractor.get().getTransitionState(),
-                    this::onSceneChanged);
+                    this::onSceneContainerTransition);
             javaAdapter.alwaysCollectFlow(
                     sceneInteractor.get().isRemoteUserInteractionOngoing(),
                     this::onRemoteUserInteractionOngoingChanged);
@@ -172,11 +175,13 @@
         }
     }
 
-    private void onSceneChanged(ObservableTransitionState transitionState) {
-        boolean isIdleOnGone = transitionState.isIdle(Scenes.Gone);
-        if (isIdleOnGone != mIsIdleOnGone) {
-            mIsIdleOnGone = isIdleOnGone;
-            if (!isIdleOnGone) {
+    private void onSceneContainerTransition(ObservableTransitionState transitionState) {
+        boolean isSceneContainerUiEmpty = transitionState.isIdle(Scenes.Gone)
+                && ((ObservableTransitionState.Idle) transitionState).getCurrentOverlays()
+                .isEmpty();
+        if (isSceneContainerUiEmpty != mIsSceneContainerUiEmpty) {
+            mIsSceneContainerUiEmpty = isSceneContainerUiEmpty;
+            if (!isSceneContainerUiEmpty) {
                 // make sure our state is sensible
                 mForceCollapsedUntilLayout = false;
             }
@@ -296,7 +301,7 @@
         // underneath.
         return mIsStatusBarExpanded
                 || (SceneContainerFlag.isEnabled()
-                && (!mIsIdleOnGone || mIsRemoteUserInteractionOngoing))
+                && (!mIsSceneContainerUiEmpty || mIsRemoteUserInteractionOngoing))
                 || mPrimaryBouncerInteractor.isShowing().getValue()
                 || mAlternateBouncerInteractor.isVisibleState()
                 || mUnlockedScreenOffAnimationController.isAnimationPlaying();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerTest.java b/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerTest.java
index c451c32..400b3b3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayerTest.java
@@ -174,6 +174,7 @@
         mMenuAnimationController = mMenuView.getMenuAnimationController();
 
         doNothing().when(mSpyContext).startActivity(any());
+        doNothing().when(mSpyContext).startActivityAsUser(any(), any());
         when(mSpyContext.getPackageManager()).thenReturn(mMockPackageManager);
 
         mLastAccessibilityButtonTargets =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
index 76539d7..633efd8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
@@ -63,21 +63,27 @@
     private val defaultDisplay =
         display(type = TYPE_INTERNAL, id = DEFAULT_DISPLAY, state = Display.STATE_ON)
 
-    private lateinit var displayRepository: DisplayRepositoryImpl
+    // This is Lazy as displays could be set before the instance is created, and we want to verify
+    // that the initial state (soon after construction) contains the expected ones set in every
+    // test.
+    private val displayRepository: DisplayRepositoryImpl by lazy {
+        DisplayRepositoryImpl(
+                displayManager,
+                testHandler,
+                TestScope(UnconfinedTestDispatcher()),
+                UnconfinedTestDispatcher(),
+            )
+            .also {
+                verify(displayManager, never()).registerDisplayListener(any(), any())
+                // It needs to be called, just once, for the initial value.
+                verify(displayManager).getDisplays()
+            }
+    }
 
     @Before
     fun setup() {
         setDisplays(listOf(defaultDisplay))
         setAllDisplaysIncludingDisabled(DEFAULT_DISPLAY)
-        displayRepository =
-            DisplayRepositoryImpl(
-                displayManager,
-                testHandler,
-                TestScope(UnconfinedTestDispatcher()),
-                UnconfinedTestDispatcher()
-            )
-        verify(displayManager, never()).registerDisplayListener(any(), any())
-        verify(displayManager, never()).getDisplays(any())
     }
 
     @Test
@@ -502,7 +508,7 @@
             .registerDisplayListener(
                 connectedDisplayListener.capture(),
                 eq(testHandler),
-                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED)
+                eq(DisplayManager.EVENT_FLAG_DISPLAY_CONNECTION_CHANGED),
             )
         return flowValue
     }
@@ -522,7 +528,7 @@
                     DisplayManager.EVENT_FLAG_DISPLAY_ADDED or
                         DisplayManager.EVENT_FLAG_DISPLAY_CHANGED or
                         DisplayManager.EVENT_FLAG_DISPLAY_REMOVED
-                )
+                ),
             )
     }
 
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/crashrecovery/TEST_MAPPING b/services/core/java/com/android/server/crashrecovery/TEST_MAPPING
index 615db34..537fb325 100644
--- a/services/core/java/com/android/server/crashrecovery/TEST_MAPPING
+++ b/services/core/java/com/android/server/crashrecovery/TEST_MAPPING
@@ -1,4 +1,9 @@
 {
+  "presubmit": [
+    {
+      "name": "CrashRecoveryModuleTests"
+    }
+  ],
   "postsubmit": [
     {
       "name": "FrameworksMockingServicesTests",
@@ -7,9 +12,6 @@
           "include-filter": "com.android.server.RescuePartyTest"
         }
       ]
-    },
-    {
-      "name": "CrashRecoveryModuleTests"
     }
   ]
 }
\ No newline at end of file
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/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/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt
index 095c819..3753b23 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/open/OpenTrampolineActivityTest.kt
@@ -23,7 +23,6 @@
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.flicker.subject.region.RegionSubject
-import androidx.test.filters.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.server.wm.flicker.activityembedding.ActivityEmbeddingTestBase
 import com.android.server.wm.flicker.helpers.ActivityEmbeddingAppHelper
@@ -43,7 +42,6 @@
  *
  * To run this test: `atest FlickerTestsActivityEmbedding:OpenTrampolineActivityTest`
  */
-@FlakyTest(bugId = 341209752)
 @RequiresDevice
 @RunWith(Parameterized::class)
 @Parameterized.UseParametersRunnerFactory(FlickerParametersRunnerFactory::class)