diff --git a/cmds/bootanimation/Android.bp b/cmds/bootanimation/Android.bp
index 3534624..f80ccf7 100644
--- a/cmds/bootanimation/Android.bp
+++ b/cmds/bootanimation/Android.bp
@@ -7,6 +7,18 @@
     default_applicable_licenses: ["frameworks_base_license"],
 }
 
+aconfig_declarations {
+    name: "bootanimation_flags",
+    package: "com.android.graphics.bootanimation.flags",
+    container: "system",
+    srcs: ["bootanimation_flags.aconfig"],
+}
+
+cc_aconfig_library {
+    name: "libbootanimationflags",
+    aconfig_declarations: "bootanimation_flags",
+}
+
 cc_defaults {
     name: "bootanimation_defaults",
 
@@ -28,6 +40,10 @@
         "liblog",
         "libutils",
     ],
+
+    static_libs: [
+        "libbootanimationflags",
+    ],
 }
 
 // bootanimation executable
diff --git a/cmds/bootanimation/BootAnimation.cpp b/cmds/bootanimation/BootAnimation.cpp
index 87c9fa4..14e2387 100644
--- a/cmds/bootanimation/BootAnimation.cpp
+++ b/cmds/bootanimation/BootAnimation.cpp
@@ -15,6 +15,7 @@
  */
 
 #define LOG_NDEBUG 0
+
 #define LOG_TAG "BootAnimation"
 #define ATRACE_TAG ATRACE_TAG_GRAPHICS
 
@@ -61,6 +62,8 @@
 
 #include "BootAnimation.h"
 
+#include <com_android_graphics_bootanimation_flags.h>
+
 #define ANIM_PATH_MAX 255
 #define STR(x)   #x
 #define STRTO(x) STR(x)
@@ -448,19 +451,21 @@
                     auto token = SurfaceComposerClient::getPhysicalDisplayToken(
                         event.header.displayId);
 
-                    if (token != mBootAnimation->mDisplayToken) {
+                    auto firstDisplay = mBootAnimation->mDisplays.front();
+                    if (token != firstDisplay.displayToken) {
                         // ignore hotplug of a secondary display
                         continue;
                     }
 
                     DisplayMode displayMode;
                     const status_t error = SurfaceComposerClient::getActiveDisplayMode(
-                        mBootAnimation->mDisplayToken, &displayMode);
+                                               firstDisplay.displayToken, &displayMode);
                     if (error != NO_ERROR) {
                         SLOGE("Can't get active display mode.");
                     }
                     mBootAnimation->resizeSurface(displayMode.resolution.getWidth(),
-                        displayMode.resolution.getHeight());
+                                                  displayMode.resolution.getHeight(),
+                                                  firstDisplay);
                 }
             }
         } while (numEvents > 0);
@@ -506,91 +511,106 @@
 status_t BootAnimation::readyToRun() {
     ATRACE_CALL();
     mAssets.addDefaultAssets();
+    return initDisplaysAndSurfaces();
+}
 
-    const std::vector<PhysicalDisplayId> ids = SurfaceComposerClient::getPhysicalDisplayIds();
-    if (ids.empty()) {
-        SLOGE("Failed to get ID for any displays\n");
+status_t BootAnimation::initDisplaysAndSurfaces() {
+    std::vector<PhysicalDisplayId> displayIds = SurfaceComposerClient::getPhysicalDisplayIds();
+    if (displayIds.empty()) {
+        SLOGE("Failed to get ID for any displays");
         return NAME_NOT_FOUND;
     }
 
-    mDisplayToken = SurfaceComposerClient::getPhysicalDisplayToken(ids.front());
-    if (mDisplayToken == nullptr) {
-        return NAME_NOT_FOUND;
+    // If Multi-Display isn't explicitly enabled, ignore all displays after the first one
+    if (!com::android::graphics::bootanimation::flags::multidisplay()) {
+        displayIds.erase(displayIds.begin() + 1, displayIds.end());
     }
 
-    DisplayMode displayMode;
-    const status_t error =
-            SurfaceComposerClient::getActiveDisplayMode(mDisplayToken, &displayMode);
-    if (error != NO_ERROR) {
-        return error;
+    for (const auto id : displayIds) {
+        if (const auto token = SurfaceComposerClient::getPhysicalDisplayToken(id)) {
+            mDisplays.push_back({.displayToken = token});
+        } else {
+            SLOGE("Failed to get display token for a display");
+            SLOGE("Failed to get display token for display %" PRIu64, id.value);
+            return NAME_NOT_FOUND;
+        }
     }
 
+    // Initialize EGL
+    mEgl = eglGetDisplay(EGL_DEFAULT_DISPLAY);
+    eglInitialize(mEgl, nullptr, nullptr);
+    EGLConfig config = getEglConfig(mEgl);
+    EGLint contextAttributes[] = {EGL_CONTEXT_CLIENT_VERSION, 2, EGL_NONE};
+    mEglContext = eglCreateContext(mEgl, config, nullptr, contextAttributes);
+
     mMaxWidth = android::base::GetIntProperty("ro.surface_flinger.max_graphics_width", 0);
     mMaxHeight = android::base::GetIntProperty("ro.surface_flinger.max_graphics_height", 0);
-    ui::Size resolution = displayMode.resolution;
-    resolution = limitSurfaceSize(resolution.width, resolution.height);
-    // create the native surface
-    sp<SurfaceControl> control = session()->createSurface(String8("BootAnimation"),
-            resolution.getWidth(), resolution.getHeight(), PIXEL_FORMAT_RGB_565,
-            ISurfaceComposerClient::eOpaque);
 
-    SurfaceComposerClient::Transaction t;
-    t.setDisplayLayerStack(mDisplayToken, ui::DEFAULT_LAYER_STACK);
-    t.setLayerStack(control, ui::DEFAULT_LAYER_STACK);
+    for (size_t displayIdx = 0; displayIdx < mDisplays.size(); displayIdx++) {
+        auto& display = mDisplays[displayIdx];
+        DisplayMode displayMode;
+        const status_t error =
+                SurfaceComposerClient::getActiveDisplayMode(display.displayToken, &displayMode);
+        if (error != NO_ERROR) {
+            return error;
+        }
+        ui::Size resolution = displayMode.resolution;
+        // Clamp each surface to max size
+        resolution = limitSurfaceSize(resolution.width, resolution.height);
+        // Create the native surface
+        display.surfaceControl =
+                session()->createSurface(String8("BootAnimation"), resolution.width,
+                                         resolution.height, PIXEL_FORMAT_RGB_565,
+                                         ISurfaceComposerClient::eOpaque);
+        // Attach surface to layerstack, and associate layerstack with physical display
+        configureDisplayAndLayerStack(display, ui::LayerStack::fromValue(displayIdx));
+        display.surface = display.surfaceControl->getSurface();
+        display.eglSurface = eglCreateWindowSurface(mEgl, config, display.surface.get(), nullptr);
 
-    t.setLayer(control, 0x40000000)
-        .apply();
+        EGLint w, h;
+        eglQuerySurface(mEgl, display.eglSurface, EGL_WIDTH, &w);
+        eglQuerySurface(mEgl, display.eglSurface, EGL_HEIGHT, &h);
+        if (eglMakeCurrent(mEgl, display.eglSurface, display.eglSurface,
+                           mEglContext) == EGL_FALSE) {
+            return NO_INIT;
+        }
+        display.initWidth = display.width = w;
+        display.initHeight = display.height = h;
+        mTargetInset = -1;
 
-    sp<Surface> s = control->getSurface();
+        // Rotate the boot animation according to the value specified in the sysprop
+        // ro.bootanim.set_orientation_<display_id>. Four values are supported: ORIENTATION_0,
+        // ORIENTATION_90, ORIENTATION_180 and ORIENTATION_270.
+        // If the value isn't specified or is ORIENTATION_0, nothing will be changed.
+        // This is needed to support boot animation in orientations different from the natural
+        // device orientation. For example, on tablets that may want to keep natural orientation
+        // portrait for applications compatibility and to have the boot animation in landscape.
+        rotateAwayFromNaturalOrientationIfNeeded(display);
 
-    // initialize opengl and egl
-    EGLDisplay display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
-    eglInitialize(display, nullptr, nullptr);
-    EGLConfig config = getEglConfig(display);
-    EGLSurface surface = eglCreateWindowSurface(display, config, s.get(), nullptr);
-    // Initialize egl context with client version number 2.0.
-    EGLint contextAttributes[] = {EGL_CONTEXT_CLIENT_VERSION, 2, EGL_NONE};
-    EGLContext context = eglCreateContext(display, config, nullptr, contextAttributes);
-    EGLint w, h;
-    eglQuerySurface(display, surface, EGL_WIDTH, &w);
-    eglQuerySurface(display, surface, EGL_HEIGHT, &h);
-
-    if (eglMakeCurrent(display, surface, surface, context) == EGL_FALSE) {
-        return NO_INIT;
-    }
-
-    mDisplay = display;
-    mContext = context;
-    mSurface = surface;
-    mInitWidth = mWidth = w;
-    mInitHeight = mHeight = h;
-    mFlingerSurfaceControl = control;
-    mFlingerSurface = s;
-    mTargetInset = -1;
-
-    // Rotate the boot animation according to the value specified in the sysprop
-    // ro.bootanim.set_orientation_<display_id>. Four values are supported: ORIENTATION_0,
-    // ORIENTATION_90, ORIENTATION_180 and ORIENTATION_270.
-    // If the value isn't specified or is ORIENTATION_0, nothing will be changed.
-    // This is needed to support having boot animation in orientations different from the natural
-    // device orientation. For example, on tablets that may want to keep natural orientation
-    // portrait for applications compatibility and to have the boot animation in landscape.
-    rotateAwayFromNaturalOrientationIfNeeded();
-
-    projectSceneToWindow();
+        projectSceneToWindow(display);
+    } // end iteration over all display tokens
 
     // Register a display event receiver
     mDisplayEventReceiver = std::make_unique<DisplayEventReceiver>();
     status_t status = mDisplayEventReceiver->initCheck();
     SLOGE_IF(status != NO_ERROR, "Initialization of DisplayEventReceiver failed with status: %d",
-            status);
+             status);
     mLooper->addFd(mDisplayEventReceiver->getFd(), 0, Looper::EVENT_INPUT,
-            new DisplayEventCallback(this), nullptr);
+                   new DisplayEventCallback(this), nullptr);
 
     return NO_ERROR;
 }
 
-void BootAnimation::rotateAwayFromNaturalOrientationIfNeeded() {
+void BootAnimation::configureDisplayAndLayerStack(const Display& display,
+                                                  ui::LayerStack layerStack) {
+    SurfaceComposerClient::Transaction t;
+    t.setDisplayLayerStack(display.displayToken, layerStack);
+    t.setLayerStack(display.surfaceControl, layerStack)
+     .setLayer(display.surfaceControl, 0x40000000)
+     .apply();
+}
+
+void BootAnimation::rotateAwayFromNaturalOrientationIfNeeded(Display& display) {
     ATRACE_CALL();
     const auto orientation = parseOrientationProperty();
 
@@ -600,16 +620,16 @@
     }
 
     if (orientation == ui::ROTATION_90 || orientation == ui::ROTATION_270) {
-        std::swap(mWidth, mHeight);
-        std::swap(mInitWidth, mInitHeight);
-        mFlingerSurfaceControl->updateDefaultBufferSize(mWidth, mHeight);
+        std::swap(display.width, display.height);
+        std::swap(display.initWidth, display.initHeight);
+        display.surfaceControl->updateDefaultBufferSize(display.width, display.height);
     }
 
-    Rect displayRect(0, 0, mWidth, mHeight);
-    Rect layerStackRect(0, 0, mWidth, mHeight);
+    Rect displayRect(0, 0, display.width, display.height);
+    Rect layerStackRect(0, 0, display.width, display.height);
 
     SurfaceComposerClient::Transaction t;
-    t.setDisplayProjection(mDisplayToken, orientation, layerStackRect, displayRect);
+    t.setDisplayProjection(display.displayToken, orientation, layerStackRect, displayRect);
     t.apply();
 }
 
@@ -640,38 +660,37 @@
     return ui::ROTATION_0;
 }
 
-void BootAnimation::projectSceneToWindow() {
+void BootAnimation::projectSceneToWindow(const Display& display) {
     ATRACE_CALL();
-    glViewport(0, 0, mWidth, mHeight);
-    glScissor(0, 0, mWidth, mHeight);
+    glViewport(0, 0, display.width, display.height);
+    glScissor(0, 0, display.width, display.height);
 }
 
-void BootAnimation::resizeSurface(int newWidth, int newHeight) {
+void BootAnimation::resizeSurface(int newWidth, int newHeight, Display& display) {
     ATRACE_CALL();
     // We assume this function is called on the animation thread.
-    if (newWidth == mWidth && newHeight == mHeight) {
+    if (newWidth == display.width && newHeight == display.height) {
         return;
     }
-    SLOGV("Resizing the boot animation surface to %d %d", newWidth, newHeight);
 
-    eglMakeCurrent(mDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
-    eglDestroySurface(mDisplay, mSurface);
+    eglMakeCurrent(mEgl, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+    eglDestroySurface(mEgl, display.eglSurface);
 
     const auto limitedSize = limitSurfaceSize(newWidth, newHeight);
-    mWidth = limitedSize.width;
-    mHeight = limitedSize.height;
+    display.width = limitedSize.width;
+    display.height = limitedSize.height;
 
-    mFlingerSurfaceControl->updateDefaultBufferSize(mWidth, mHeight);
-    EGLConfig config = getEglConfig(mDisplay);
-    EGLSurface surface = eglCreateWindowSurface(mDisplay, config, mFlingerSurface.get(), nullptr);
-    if (eglMakeCurrent(mDisplay, surface, surface, mContext) == EGL_FALSE) {
-        SLOGE("Can't make the new surface current. Error %d", eglGetError());
+    display.surfaceControl->updateDefaultBufferSize(display.width, display.height);
+    EGLConfig config = getEglConfig(mEgl);
+    EGLSurface eglSurface = eglCreateWindowSurface(mEgl, config, display.surface.get(), nullptr);
+    if (eglMakeCurrent(mEgl, eglSurface, eglSurface, mEglContext) == EGL_FALSE) {
+        SLOGE("Can't make the new eglSurface current. Error %d", eglGetError());
         return;
     }
 
-    projectSceneToWindow();
+    projectSceneToWindow(display);
 
-    mSurface = surface;
+    display.eglSurface = eglSurface;
 }
 
 bool BootAnimation::preloadAnimation() {
@@ -801,24 +820,26 @@
     // animation.
     if (mZipFileName.empty()) {
         ALOGD("No animation file");
-        result = android();
+        result = android(mDisplays.front());
     } else {
         result = movie();
     }
 
     mCallbacks->shutdown();
-    eglMakeCurrent(mDisplay, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
-    eglDestroyContext(mDisplay, mContext);
-    eglDestroySurface(mDisplay, mSurface);
-    mFlingerSurface.clear();
-    mFlingerSurfaceControl.clear();
-    eglTerminate(mDisplay);
+    eglMakeCurrent(mEgl, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);
+    eglDestroyContext(mEgl, mEglContext);
+    for (auto& display : mDisplays) {
+        eglDestroySurface(mEgl, display.eglSurface);
+        display.surface.clear();
+        display.surfaceControl.clear();
+    }
+    eglTerminate(mEgl);
     eglReleaseThread();
     IPCThreadState::self()->stopProcess();
     return result;
 }
 
-bool BootAnimation::android() {
+bool BootAnimation::android(const Display& display) {
     ATRACE_CALL();
     glActiveTexture(GL_TEXTURE0);
 
@@ -836,7 +857,7 @@
 
     glClearColor(0,0,0,1);
     glClear(GL_COLOR_BUFFER_BIT);
-    eglSwapBuffers(mDisplay, mSurface);
+    eglSwapBuffers(mEgl, display.eglSurface);
 
     // Blend state
     glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
@@ -844,11 +865,11 @@
     const nsecs_t startTime = systemTime();
     do {
         processDisplayEvents();
-        const GLint xc = (mWidth  - mAndroid[0].w) / 2;
-        const GLint yc = (mHeight - mAndroid[0].h) / 2;
+        const GLint xc = (display.width  - mAndroid[0].w) / 2;
+        const GLint yc = (display.height - mAndroid[0].h) / 2;
         const Rect updateRect(xc, yc, xc + mAndroid[0].w, yc + mAndroid[0].h);
-        glScissor(updateRect.left, mHeight - updateRect.bottom, updateRect.width(),
-                updateRect.height());
+        glScissor(updateRect.left, display.height - updateRect.bottom, updateRect.width(),
+                  updateRect.height());
 
         nsecs_t now = systemTime();
         double time = now - startTime;
@@ -862,14 +883,14 @@
         glEnable(GL_SCISSOR_TEST);
         glDisable(GL_BLEND);
         glBindTexture(GL_TEXTURE_2D, mAndroid[1].name);
-        drawTexturedQuad(x,                 yc, mAndroid[1].w, mAndroid[1].h);
-        drawTexturedQuad(x + mAndroid[1].w, yc, mAndroid[1].w, mAndroid[1].h);
+        drawTexturedQuad(x,                 yc, mAndroid[1].w, mAndroid[1].h, display);
+        drawTexturedQuad(x + mAndroid[1].w, yc, mAndroid[1].w, mAndroid[1].h, display);
 
         glEnable(GL_BLEND);
         glBindTexture(GL_TEXTURE_2D, mAndroid[0].name);
-        drawTexturedQuad(xc, yc, mAndroid[0].w, mAndroid[0].h);
+        drawTexturedQuad(xc, yc, mAndroid[0].w, mAndroid[0].h, display);
 
-        EGLBoolean res = eglSwapBuffers(mDisplay, mSurface);
+        EGLBoolean res = eglSwapBuffers(mEgl, display.eglSurface);
         if (res == EGL_FALSE)
             break;
 
@@ -888,7 +909,7 @@
 
 void BootAnimation::checkExit() {
     ATRACE_CALL();
-    // Allow surface flinger to gracefully request shutdown
+    // Allow SurfaceFlinger to gracefully request shutdown
     char value[PROPERTY_VALUE_MAX];
     property_get(EXIT_PROP_NAME, value, "0");
     int exitnow = atoi(value);
@@ -1034,7 +1055,8 @@
     return status;
 }
 
-void BootAnimation::drawText(const char* str, const Font& font, bool bold, int* x, int* y) {
+void BootAnimation::drawText(const char* str, const Font& font, bool bold,
+                             int* x, int* y, const Display& display) {
     ATRACE_CALL();
     glEnable(GL_BLEND);  // Allow us to draw on top of the animation
     glBindTexture(GL_TEXTURE_2D, font.texture.name);
@@ -1045,14 +1067,14 @@
     const int strWidth = font.char_width * len;
 
     if (*x == TEXT_CENTER_VALUE) {
-        *x = (mWidth - strWidth) / 2;
+        *x = (display.width - strWidth) / 2;
     } else if (*x < 0) {
-        *x = mWidth + *x - strWidth;
+        *x = display.width + *x - strWidth;
     }
     if (*y == TEXT_CENTER_VALUE) {
-        *y = (mHeight - font.char_height) / 2;
+        *y = (display.height - font.char_height) / 2;
     } else if (*y < 0) {
-        *y = mHeight + *y - font.char_height;
+        *y = display.height + *y - font.char_height;
     }
 
     for (int i = 0; i < len; i++) {
@@ -1072,7 +1094,7 @@
         float v1 = v0 + 1.0f / FONT_NUM_ROWS / 2;
         float u1 = u0 + 1.0f / FONT_NUM_COLS;
         glUniform4f(mTextCropAreaLocation, u0, v0, u1, v1);
-        drawTexturedQuad(*x, *y, font.char_width, font.char_height);
+        drawTexturedQuad(*x, *y, font.char_width, font.char_height, display);
 
         *x += font.char_width;
     }
@@ -1082,7 +1104,8 @@
 }
 
 // We render 12 or 24 hour time.
-void BootAnimation::drawClock(const Font& font, const int xPos, const int yPos) {
+void BootAnimation::drawClock(const Font& font, const int xPos, const int yPos,
+                              const Display& display) {
     ATRACE_CALL();
     static constexpr char TIME_FORMAT_12[] = "%l:%M";
     static constexpr char TIME_FORMAT_24[] = "%H:%M";
@@ -1105,10 +1128,11 @@
     char* out = timeBuff[0] == ' ' ? &timeBuff[1] : &timeBuff[0];
     int x = xPos;
     int y = yPos;
-    drawText(out, font, false, &x, &y);
+    drawText(out, font, false, &x, &y, display);
 }
 
-void BootAnimation::drawProgress(int percent, const Font& font, const int xPos, const int yPos) {
+void BootAnimation::drawProgress(int percent, const Font& font, const int xPos, const int yPos,
+                                 const Display& display) {
     ATRACE_CALL();
     static constexpr int PERCENT_LENGTH = 5;
 
@@ -1118,7 +1142,7 @@
     sprintf(percentBuff, "%d;", percent);
     int x = xPos;
     int y = yPos;
-    drawText(percentBuff, font, false, &x, &y);
+    drawText(percentBuff, font, false, &x, &y, display);
 }
 
 bool BootAnimation::parseAnimationDesc(Animation& animation)  {
@@ -1247,8 +1271,7 @@
 
 bool BootAnimation::preloadZip(Animation& animation) {
     ATRACE_CALL();
-    // read all the data structures
-    const size_t pcount = animation.parts.size();
+    const size_t numParts = animation.parts.size();
     void *cookie = nullptr;
     ZipFileRO* zip = animation.zip;
     if (!zip->startIteration(&cookie)) {
@@ -1284,8 +1307,8 @@
                 continue;
             }
 
-            for (size_t j = 0; j < pcount; j++) {
-                if (path.string() == animation.parts[j].path.c_str()) {
+            for (size_t partIdx = 0; partIdx < numParts; partIdx++) {
+                if (path.string() == animation.parts[partIdx].path.c_str()) {
                     uint16_t method;
                     // supports only stored png files
                     if (zip->getEntryInfo(entry, &method, nullptr, nullptr, nullptr, nullptr,
@@ -1293,7 +1316,7 @@
                         if (method == ZipFileRO::kCompressStored) {
                             FileMap* map = zip->createEntryFileMap(entry);
                             if (map) {
-                                Animation::Part& part(animation.parts.editItemAt(j));
+                                Animation::Part& part(animation.parts.editItemAt(partIdx));
                                 if (leaf == "audio.wav") {
                                     // a part may have at most one audio file
                                     part.audioData = (uint8_t *)map->getDataPtr();
@@ -1324,7 +1347,8 @@
     // If there is trimData present, override the positioning defaults.
     for (Animation::Part& part : animation.parts) {
         const char* trimDataStr = part.trimData.c_str();
-        for (size_t frameIdx = 0; frameIdx < part.frames.size(); frameIdx++) {
+        const size_t numFramesInPart = part.frames.size();
+        for (size_t frameIdxInPart = 0; frameIdxInPart < numFramesInPart; frameIdxInPart++) {
             const char* endl = strstr(trimDataStr, "\n");
             // No more trimData for this part.
             if (endl == nullptr) {
@@ -1335,7 +1359,7 @@
             trimDataStr = ++endl;
             int width = 0, height = 0, x = 0, y = 0;
             if (sscanf(lineStr, "%dx%d+%d+%d", &width, &height, &x, &y) == 4) {
-                Animation::Frame& frame(part.frames.editItemAt(frameIdx));
+                Animation::Frame& frame(part.frames.editItemAt(frameIdxInPart));
                 frame.trimWidth = width;
                 frame.trimHeight = height;
                 frame.trimX = x;
@@ -1458,13 +1482,15 @@
     return b1 + ( x - a1 ) * ( b2 - b1 ) / ( a2 - a1 );
 }
 
-void BootAnimation::drawTexturedQuad(float xStart, float yStart, float width, float height) {
+void BootAnimation::drawTexturedQuad(float xStart, float yStart,
+                                     float width, float height,
+                                     const Display& display) {
     ATRACE_CALL();
     // Map coordinates from screen space to world space.
-    float x0 = mapLinear(xStart, 0, mWidth, -1, 1);
-    float y0 = mapLinear(yStart, 0, mHeight, -1, 1);
-    float x1 = mapLinear(xStart + width, 0, mWidth, -1, 1);
-    float y1 = mapLinear(yStart + height, 0, mHeight, -1, 1);
+    float x0 = mapLinear(xStart, 0, display.width, -1, 1);
+    float y0 = mapLinear(yStart, 0, display.height, -1, 1);
+    float x1 = mapLinear(xStart + width, 0, display.width, -1, 1);
+    float y1 = mapLinear(yStart + height, 0, display.height, -1, 1);
     // Update quad vertex positions.
     quadPositions[0] = x0;
     quadPositions[1] = y0;
@@ -1511,7 +1537,7 @@
 
 bool BootAnimation::playAnimation(const Animation& animation) {
     ATRACE_CALL();
-    const size_t pcount = animation.parts.size();
+    const size_t numParts = animation.parts.size();
     nsecs_t frameDuration = s2ns(1) / animation.fps;
 
     SLOGD("%sAnimationShownTiming start time: %" PRId64 "ms", mShuttingDown ? "Shutdown" : "Boot",
@@ -1521,9 +1547,9 @@
     int lastDisplayedProgress = 0;
     int colorTransitionStart = animation.colorTransitionStart;
     int colorTransitionEnd = animation.colorTransitionEnd;
-    for (size_t i=0 ; i<pcount ; i++) {
-        const Animation::Part& part(animation.parts[i]);
-        const size_t fcount = part.frames.size();
+    for (size_t partIdx = 0; partIdx < numParts; partIdx++) {
+        const Animation::Part& part(animation.parts[partIdx]);
+        const size_t numFramesInPart = part.frames.size();
         glBindTexture(GL_TEXTURE_2D, 0);
 
         // Handle animation package
@@ -1535,7 +1561,9 @@
         }
 
         // process the part not only while the count allows but also if already fading
-        for (int r=0 ; !part.count || r<part.count || fadedFramesCount > 0 ; r++) {
+        for (int frameIdx = 0;
+             !part.count || frameIdx < part.count || fadedFramesCount > 0;
+             frameIdx++) {
             if (shouldStopPlayingPart(part, fadedFramesCount, lastDisplayedProgress)) break;
 
             // It's possible that the sysprops were not loaded yet at this boot phase.
@@ -1550,12 +1578,12 @@
                     const int transitionLength = colorTransitionEnd - colorTransitionStart;
                     if (part.postDynamicColoring) {
                         colorTransitionStart = 0;
-                        colorTransitionEnd = fmin(transitionLength, fcount - 1);
+                        colorTransitionEnd = fmin(transitionLength, numFramesInPart - 1);
                     }
                 }
             }
 
-            mCallbacks->playPart(i, part, r);
+            mCallbacks->playPart(partIdx, part, frameIdx);
 
             glClearColor(
                     part.backgroundColor[0],
@@ -1569,11 +1597,10 @@
 
             // For the last animation, if we have progress indicator from
             // the system, display it.
-            int currentProgress = android::base::GetIntProperty(PROGRESS_PROP_NAME, 0);
-            bool displayProgress = animation.progressEnabled &&
-                (i == (pcount -1)) && currentProgress != 0;
+            const bool displayProgress = animation.progressEnabled && (partIdx == (numParts - 1)) &&
+                    android::base::GetIntProperty(PROGRESS_PROP_NAME, 0) != 0;
 
-            for (size_t j=0 ; j<fcount ; j++) {
+            for (size_t frameIdxInPart = 0; frameIdxInPart < numFramesInPart; frameIdxInPart++) {
                 if (shouldStopPlayingPart(part, fadedFramesCount, lastDisplayedProgress)) break;
 
                 // Color progress is
@@ -1584,21 +1611,15 @@
                 // - 1 for parts that come after.
                 float colorProgress = part.useDynamicColoring
                     ? fmin(fmax(
-                        ((float)j - colorTransitionStart) /
+                        (static_cast<float>(frameIdxInPart) - colorTransitionStart) /
                             fmax(colorTransitionEnd - colorTransitionStart, 1.0f), 0.0f), 1.0f)
                     : (part.postDynamicColoring ? 1 : 0);
-
                 processDisplayEvents();
 
-                const double ratio_w = static_cast<double>(mWidth) / mInitWidth;
-                const double ratio_h = static_cast<double>(mHeight) / mInitHeight;
-                const int animationX = (mWidth - animation.width * ratio_w) / 2;
-                const int animationY = (mHeight - animation.height * ratio_h) / 2;
-
-                const Animation::Frame& frame(part.frames[j]);
+                const Animation::Frame& frame(part.frames[frameIdxInPart]);
                 nsecs_t lastFrame = systemTime();
 
-                if (r > 0) {
+                if (frameIdx > 0) {
                     glBindTexture(GL_TEXTURE_2D, frame.tid);
                 } else {
                     if (part.count != 1) {
@@ -1611,17 +1632,6 @@
                     initTexture(frame.map, &w, &h, false /* don't premultiply alpha */);
                 }
 
-                const int trimWidth = frame.trimWidth * ratio_w;
-                const int trimHeight = frame.trimHeight * ratio_h;
-                const int trimX = frame.trimX * ratio_w;
-                const int trimY = frame.trimY * ratio_h;
-                const int xc = animationX + trimX;
-                const int yc = animationY + trimY;
-                glClear(GL_COLOR_BUFFER_BIT);
-                // specify the y center as ceiling((mHeight - frame.trimHeight) / 2)
-                // which is equivalent to mHeight - (yc + frame.trimHeight)
-                const int frameDrawY = mHeight - (yc + trimHeight);
-
                 float fade = 0;
                 // if the part hasn't been stopped yet then continue fading if necessary
                 if (exitPending() && part.hasFadingPhase()) {
@@ -1630,40 +1640,66 @@
                         fadedFramesCount = MAX_FADED_FRAMES_COUNT; // no more fading
                     }
                 }
-                glUseProgram(mImageShader);
-                glUniform1i(mImageTextureLocation, 0);
-                glUniform1f(mImageFadeLocation, fade);
-                if (animation.dynamicColoringEnabled) {
-                    glUniform1f(mImageColorProgressLocation, colorProgress);
-                }
-                glEnable(GL_BLEND);
-                drawTexturedQuad(xc, frameDrawY, trimWidth, trimHeight);
-                glDisable(GL_BLEND);
 
-                if (mClockEnabled && mTimeIsAccurate && validClock(part)) {
-                    drawClock(animation.clockFont, part.clockPosX, part.clockPosY);
-                }
+                // Draw the current frame's texture on every physical display that is enabled.
+                for (const auto& display : mDisplays) {
+                    eglMakeCurrent(mEgl, display.eglSurface, display.eglSurface, mEglContext);
 
-                if (displayProgress) {
-                    int newProgress = android::base::GetIntProperty(PROGRESS_PROP_NAME, 0);
-                    // In case the new progress jumped suddenly, still show an
-                    // increment of 1.
-                    if (lastDisplayedProgress != 100) {
-                      // Artificially sleep 1/10th a second to slow down the animation.
-                      usleep(100000);
-                      if (lastDisplayedProgress < newProgress) {
-                        lastDisplayedProgress++;
-                      }
+                    const double ratioW =
+                            static_cast<double>(display.width) / display.initWidth;
+                    const double ratioH =
+                            static_cast<double>(display.height) / display.initHeight;
+                    const int animationX = (display.width - animation.width * ratioW) / 2;
+                    const int animationY = (display.height - animation.height * ratioH) / 2;
+
+                    const int trimWidth = frame.trimWidth * ratioW;
+                    const int trimHeight = frame.trimHeight * ratioH;
+                    const int trimX = frame.trimX * ratioW;
+                    const int trimY = frame.trimY * ratioH;
+                    const int xc = animationX + trimX;
+                    const int yc = animationY + trimY;
+                    projectSceneToWindow(display);
+                    handleViewport(frameDuration, display);
+                    glClear(GL_COLOR_BUFFER_BIT);
+                    // specify the y center as ceiling((height - frame.trimHeight) / 2)
+                    // which is equivalent to height - (yc + frame.trimHeight)
+                    const int frameDrawY = display.height - (yc + trimHeight);
+
+                    glUseProgram(mImageShader);
+                    glUniform1i(mImageTextureLocation, 0);
+                    glUniform1f(mImageFadeLocation, fade);
+                    if (animation.dynamicColoringEnabled) {
+                        glUniform1f(mImageColorProgressLocation, colorProgress);
                     }
-                    // Put the progress percentage right below the animation.
-                    int posY = animation.height / 3;
-                    int posX = TEXT_CENTER_VALUE;
-                    drawProgress(lastDisplayedProgress, animation.progressFont, posX, posY);
+                    glEnable(GL_BLEND);
+                    drawTexturedQuad(xc, frameDrawY, trimWidth, trimHeight, display);
+                    glDisable(GL_BLEND);
+
+                    if (mClockEnabled && mTimeIsAccurate && validClock(part)) {
+                        drawClock(animation.clockFont, part.clockPosX, part.clockPosY, display);
+                    }
+
+                    if (displayProgress) {
+                        int newProgress = android::base::GetIntProperty(PROGRESS_PROP_NAME, 0);
+                        // In case the new progress jumped suddenly, still show an
+                        // increment of 1.
+                        if (lastDisplayedProgress != 100) {
+                          // Artificially sleep 1/10th a second to slow down the animation.
+                          usleep(100000);
+                          if (lastDisplayedProgress < newProgress) {
+                            lastDisplayedProgress++;
+                          }
+                        }
+                        // Put the progress percentage right below the animation.
+                        int posY = animation.height / 3;
+                        int posX = TEXT_CENTER_VALUE;
+                        drawProgress(lastDisplayedProgress,
+                            animation.progressFont, posX, posY, display);
+                    }
+
+                    eglSwapBuffers(mEgl, display.eglSurface);
                 }
 
-                handleViewport(frameDuration);
-
-                eglSwapBuffers(mDisplay, mSurface);
 
                 nsecs_t now = systemTime();
                 nsecs_t delay = frameDuration - (now - lastFrame);
@@ -1709,9 +1745,7 @@
     // Free textures created for looping parts now that the animation is done.
     for (const Animation::Part& part : animation.parts) {
         if (part.count != 1) {
-            const size_t fcount = part.frames.size();
-            for (size_t j = 0; j < fcount; j++) {
-                const Animation::Frame& frame(part.frames[j]);
+            for (const auto& frame : part.frames) {
                 glDeleteTextures(1, &frame.tid);
             }
         }
@@ -1730,16 +1764,17 @@
     mLooper->pollOnce(0);
 }
 
-void BootAnimation::handleViewport(nsecs_t timestep) {
+void BootAnimation::handleViewport(nsecs_t timestep, const Display& display) {
     ATRACE_CALL();
-    if (mShuttingDown || !mFlingerSurfaceControl || mTargetInset == 0) {
+    if (mShuttingDown || !display.surfaceControl || mTargetInset == 0) {
         return;
     }
     if (mTargetInset < 0) {
         // Poll the amount for the top display inset. This will return -1 until persistent properties
         // have been loaded.
-        mTargetInset = android::base::GetIntProperty("persist.sys.displayinset.top",
-                -1 /* default */, -1 /* min */, mHeight / 2 /* max */);
+        mTargetInset =
+                android::base::GetIntProperty("persist.sys.displayinset.top", -1 /* default */,
+                                              -1 /* min */, display.height / 2 /* max */);
     }
     if (mTargetInset <= 0) {
         return;
@@ -1751,19 +1786,27 @@
         int interpolatedInset = (cosf((fraction + 1) * M_PI) / 2.0f + 0.5f) * mTargetInset;
 
         SurfaceComposerClient::Transaction()
-                .setCrop(mFlingerSurfaceControl, Rect(0, interpolatedInset, mWidth, mHeight))
+                .setCrop(display.surfaceControl,
+                         Rect(0, interpolatedInset, display.width, display.height))
                 .apply();
     } else {
         // At the end of the animation, we switch to the viewport that DisplayManager will apply
         // later. This changes the coordinate system, and means we must move the surface up by
         // the inset amount.
-        Rect layerStackRect(0, 0, mWidth, mHeight - mTargetInset);
-        Rect displayRect(0, mTargetInset, mWidth, mHeight);
-
+        Rect layerStackRect(0, 0,
+                display.width,
+                display.height - mTargetInset);
+        Rect displayRect(0, mTargetInset,
+                display.width,
+                display.height);
         SurfaceComposerClient::Transaction t;
-        t.setPosition(mFlingerSurfaceControl, 0, -mTargetInset)
-                .setCrop(mFlingerSurfaceControl, Rect(0, mTargetInset, mWidth, mHeight));
-        t.setDisplayProjection(mDisplayToken, ui::ROTATION_0, layerStackRect, displayRect);
+        t.setPosition(display.surfaceControl, 0, -mTargetInset)
+                .setCrop(display.surfaceControl,
+                         Rect(0, mTargetInset,
+                              display.width,
+                              display.height));
+        t.setDisplayProjection(display.displayToken,
+                ui::ROTATION_0, layerStackRect, displayRect);
         t.apply();
 
         mTargetInset = mCurrentInset = 0;
diff --git a/cmds/bootanimation/BootAnimation.h b/cmds/bootanimation/BootAnimation.h
index 8683b71..0a05746 100644
--- a/cmds/bootanimation/BootAnimation.h
+++ b/cmds/bootanimation/BootAnimation.h
@@ -31,6 +31,7 @@
 #include <binder/IBinder.h>
 
 #include <ui/Rotation.h>
+#include <ui/LayerStack.h>
 
 #include <EGL/egl.h>
 #include <GLES2/gl2.h>
@@ -119,6 +120,18 @@
         float endColors[4][3];   // End colors of dynamic color transition.
     };
 
+    // Collects all attributes that must be tracked per physical display.
+    struct Display {
+        int width;
+        int height;
+        int initWidth;
+        int initHeight;
+        EGLDisplay  eglSurface;
+        sp<IBinder> displayToken;
+        sp<SurfaceControl> surfaceControl;
+        sp<Surface> surface;
+    };
+
     // All callbacks will be called from this class's internal thread.
     class Callbacks : public RefBase {
     public:
@@ -181,14 +194,18 @@
         bool premultiplyAlpha = true);
     status_t initFont(Font* font, const char* fallback);
     void initShaders();
-    bool android();
+    bool android(const Display& display);
+    status_t initDisplaysAndSurfaces();
     bool movie();
-    void drawText(const char* str, const Font& font, bool bold, int* x, int* y);
-    void drawClock(const Font& font, const int xPos, const int yPos);
-    void drawProgress(int percent, const Font& font, const int xPos, const int yPos);
+    void drawText(const char* str, const Font& font, bool bold,
+                  int* x, int* y, const Display& display);
+    void drawClock(const Font& font, const int xPos, const int yPos, const Display& display);
+    void drawProgress(int percent, const Font& font,
+                      const int xPos, const int yPos, const Display& display);
     void fadeFrame(int frameLeft, int frameBottom, int frameWidth, int frameHeight,
                    const Animation::Part& part, int fadedFramesCount);
-    void drawTexturedQuad(float xStart, float yStart, float width, float height);
+    void drawTexturedQuad(float xStart, float yStart,
+                          float width, float height, const Display& display);
     bool validClock(const Animation::Part& part);
     Animation* loadAnimation(const String8&);
     bool playAnimation(const Animation&);
@@ -200,36 +217,31 @@
     bool preloadAnimation();
     EGLConfig getEglConfig(const EGLDisplay&);
     ui::Size limitSurfaceSize(int width, int height) const;
-    void resizeSurface(int newWidth, int newHeight);
-    void projectSceneToWindow();
-    void rotateAwayFromNaturalOrientationIfNeeded();
+    void resizeSurface(int newWidth, int newHeight, Display& display);
+    void projectSceneToWindow(const Display& display);
+    void rotateAwayFromNaturalOrientationIfNeeded(Display& display);
     ui::Rotation parseOrientationProperty();
+    void configureDisplayAndLayerStack(const Display& display, ui::LayerStack layerStack);
 
     bool shouldStopPlayingPart(const Animation::Part& part, int fadedFramesCount,
                                int lastDisplayedProgress);
     void checkExit();
 
-    void handleViewport(nsecs_t timestep);
+    void handleViewport(nsecs_t timestep, const Display& display);
     void initDynamicColors();
 
     sp<SurfaceComposerClient>       mSession;
     AssetManager mAssets;
     Texture     mAndroid[2];
-    int         mWidth;
-    int         mHeight;
-    int         mInitWidth;
-    int         mInitHeight;
     int         mMaxWidth = 0;
     int         mMaxHeight = 0;
     int         mCurrentInset;
     int         mTargetInset;
     bool        mUseNpotTextures = false;
-    EGLDisplay  mDisplay;
-    EGLDisplay  mContext;
-    EGLDisplay  mSurface;
-    sp<IBinder> mDisplayToken;
-    sp<SurfaceControl> mFlingerSurfaceControl;
-    sp<Surface> mFlingerSurface;
+    EGLDisplay  mEgl;
+    EGLDisplay  mEglContext;
+    // Per-Display Attributes (to support multi-display)
+    std::vector<Display> mDisplays;
     bool        mClockEnabled;
     bool        mTimeIsAccurate;
     bool        mTimeFormat12Hour;
diff --git a/cmds/bootanimation/bootanimation_flags.aconfig b/cmds/bootanimation/bootanimation_flags.aconfig
new file mode 100644
index 0000000..04837b9
--- /dev/null
+++ b/cmds/bootanimation/bootanimation_flags.aconfig
@@ -0,0 +1,12 @@
+package: "com.android.graphics.bootanimation.flags"
+container: "system"
+
+flag {
+  name: "multidisplay"
+  namespace: "bootanimation"
+  description: "Enable boot animation on multiple displays (e.g. foldables)"
+  bug: "335406617"
+  is_fixed_read_only: true
+}
+
+
diff --git a/cmds/idmap2/include/idmap2/Idmap.h b/cmds/idmap2/include/idmap2/Idmap.h
index e86f814..b0ba019 100644
--- a/cmds/idmap2/include/idmap2/Idmap.h
+++ b/cmds/idmap2/include/idmap2/Idmap.h
@@ -21,18 +21,19 @@
  * header                     := magic version target_crc overlay_crc fulfilled_policies
  *                               enforce_overlayable target_path overlay_path overlay_name
  *                               debug_info
- * data                       := data_header target_entry* target_inline_entry*
-                                 target_inline_entry_value* config* overlay_entry* string_pool
+ * data                       := data_header target_entries target_inline_entries
+                                 target_inline_entry_value* config* overlay_entries string_pool
  * data_header                := target_entry_count target_inline_entry_count
                                  target_inline_entry_value_count config_count overlay_entry_count
  *                               string_pool_index
- * target_entry               := target_id overlay_id
- * target_inline_entry        := target_id start_value_index value_count
+ * target_entries             := target_id* overlay_id*
+ * target_inline_entries      := target_id* target_inline_value_header*
+ * target_inline_value_header := start_value_index value_count
  * target_inline_entry_value  := config_index Res_value::size padding(1) Res_value::type
  *                               Res_value::value
  * config                     := target_id Res_value::size padding(1) Res_value::type
  *                               Res_value::value
- * overlay_entry              := overlay_id target_id
+ * overlay_entries            := overlay_id* target_id*
  *
  * debug_info                       := string
  * enforce_overlayable              := <uint32_t>
diff --git a/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp b/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp
index 8976924..00ef0c7 100644
--- a/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp
+++ b/cmds/idmap2/libidmap2/BinaryStreamVisitor.cpp
@@ -66,43 +66,57 @@
 void BinaryStreamVisitor::visit(const IdmapData& data) {
   for (const auto& target_entry : data.GetTargetEntries()) {
     Write32(target_entry.target_id);
+  }
+  for (const auto& target_entry : data.GetTargetEntries()) {
     Write32(target_entry.overlay_id);
   }
 
-  static constexpr uint16_t kValueSize = 8U;
-  std::vector<std::pair<ConfigDescription, TargetValue>> target_values;
-  target_values.reserve(data.GetHeader()->GetTargetInlineEntryValueCount());
-  for (const auto& target_entry : data.GetTargetInlineEntries()) {
-    Write32(target_entry.target_id);
-    Write32(target_values.size());
-    Write32(target_entry.values.size());
-    target_values.insert(
-        target_values.end(), target_entry.values.begin(), target_entry.values.end());
+  uint32_t current_inline_entry_values_count = 0;
+  for (const auto& target_inline_entry : data.GetTargetInlineEntries()) {
+    Write32(target_inline_entry.target_id);
+  }
+  for (const auto& target_inline_entry : data.GetTargetInlineEntries()) {
+    Write32(current_inline_entry_values_count);
+    Write32(target_inline_entry.values.size());
+    current_inline_entry_values_count += target_inline_entry.values.size();
   }
 
   std::vector<ConfigDescription> configs;
   configs.reserve(data.GetHeader()->GetConfigCount());
-  for (const auto& target_entry_value : target_values) {
-    auto config_it = find(configs.begin(), configs.end(), target_entry_value.first);
-    if (config_it != configs.end()) {
-      Write32(config_it - configs.begin());
-    } else {
-      Write32(configs.size());
-      configs.push_back(target_entry_value.first);
+  for (const auto& target_entry : data.GetTargetInlineEntries()) {
+    for (const auto& target_entry_value : target_entry.values) {
+      auto config_it = std::find(configs.begin(), configs.end(), target_entry_value.first);
+      if (config_it != configs.end()) {
+        Write32(config_it - configs.begin());
+      } else {
+        Write32(configs.size());
+        configs.push_back(target_entry_value.first);
+      }
+      // We're writing a Res_value entry here, and the first 3 bytes of that are
+      // sizeof() + a padding 0 byte
+      static constexpr decltype(android::Res_value::size) kSize = sizeof(android::Res_value);
+      Write16(kSize);
+      Write8(0U);
+      Write8(target_entry_value.second.data_type);
+      Write32(target_entry_value.second.data_value);
     }
-    Write16(kValueSize);
-    Write8(0U);  // padding
-    Write8(target_entry_value.second.data_type);
-    Write32(target_entry_value.second.data_value);
   }
 
-  for( auto& cd : configs) {
-    cd.swapHtoD();
-    stream_.write(reinterpret_cast<char*>(&cd), sizeof(cd));
+  if (!configs.empty()) {
+    stream_.write(reinterpret_cast<const char*>(&configs.front()),
+                  sizeof(configs.front()) * configs.size());
+    if (configs.size() >= 100) {
+      // Let's write a message to future us so that they know when to replace the linear search
+      // in `configs` vector with something more efficient.
+      LOG(WARNING) << "Idmap got " << configs.size()
+                   << " configurations, time to fix the bruteforce search";
+    }
   }
 
   for (const auto& overlay_entry : data.GetOverlayEntries()) {
     Write32(overlay_entry.overlay_id);
+  }
+  for (const auto& overlay_entry : data.GetOverlayEntries()) {
     Write32(overlay_entry.target_id);
   }
 
diff --git a/cmds/idmap2/libidmap2/Idmap.cpp b/cmds/idmap2/libidmap2/Idmap.cpp
index 12d9dd9..7680109 100644
--- a/cmds/idmap2/libidmap2/Idmap.cpp
+++ b/cmds/idmap2/libidmap2/Idmap.cpp
@@ -204,73 +204,91 @@
   }
 
   // Read the mapping of target resource id to overlay resource value.
+  data->target_entries_.resize(data->header_->GetTargetEntryCount());
   for (size_t i = 0; i < data->header_->GetTargetEntryCount(); i++) {
-    TargetEntry target_entry{};
-    if (!Read32(stream, &target_entry.target_id) || !Read32(stream, &target_entry.overlay_id)) {
+    if (!Read32(stream, &data->target_entries_[i].target_id)) {
       return nullptr;
     }
-    data->target_entries_.emplace_back(target_entry);
+  }
+  for (size_t i = 0; i < data->header_->GetTargetEntryCount(); i++) {
+    if (!Read32(stream, &data->target_entries_[i].overlay_id)) {
+      return nullptr;
+    }
   }
 
   // Read the mapping of target resource id to inline overlay values.
-  std::vector<std::tuple<TargetInlineEntry, uint32_t, uint32_t>> target_inline_entries;
+  struct TargetInlineEntryHeader {
+    ResourceId target_id;
+    uint32_t values_offset;
+    uint32_t values_count;
+  };
+  std::vector<TargetInlineEntryHeader> target_inline_entries(
+      data->header_->GetTargetInlineEntryCount());
   for (size_t i = 0; i < data->header_->GetTargetInlineEntryCount(); i++) {
-    TargetInlineEntry target_entry{};
-    uint32_t entry_offset;
-    uint32_t entry_count;
-    if (!Read32(stream, &target_entry.target_id) || !Read32(stream, &entry_offset)
-        || !Read32(stream, &entry_count)) {
+    if (!Read32(stream, &target_inline_entries[i].target_id)) {
       return nullptr;
     }
-    target_inline_entries.emplace_back(target_entry, entry_offset, entry_count);
+  }
+  for (size_t i = 0; i < data->header_->GetTargetInlineEntryCount(); i++) {
+    if (!Read32(stream, &target_inline_entries[i].values_offset) ||
+        !Read32(stream, &target_inline_entries[i].values_count)) {
+      return nullptr;
+    }
   }
 
   // Read the inline overlay resource values
-  std::vector<std::pair<uint32_t, TargetValue>> target_values;
-  uint8_t unused1;
-  uint16_t unused2;
-  for (size_t i = 0; i < data->header_->GetTargetInlineEntryValueCount(); i++) {
+  struct TargetValueHeader {
     uint32_t config_index;
-    if (!Read32(stream, &config_index)) {
+    DataType data_type;
+    DataValue data_value;
+  };
+  std::vector<TargetValueHeader> target_values(data->header_->GetTargetInlineEntryValueCount());
+  for (size_t i = 0; i < data->header_->GetTargetInlineEntryValueCount(); i++) {
+    auto& value = target_values[i];
+    if (!Read32(stream, &value.config_index)) {
       return nullptr;
     }
-    TargetValue value;
-    if (!Read16(stream, &unused2)
-        || !Read8(stream, &unused1)
-        || !Read8(stream, &value.data_type)
-        || !Read32(stream, &value.data_value)) {
+    // skip the padding
+    stream.seekg(3, std::ios::cur);
+    if (!Read8(stream, &value.data_type) || !Read32(stream, &value.data_value)) {
       return nullptr;
     }
-    target_values.emplace_back(config_index, value);
   }
 
   // Read the configurations
-  std::vector<ConfigDescription> configurations;
-  for (size_t i = 0; i < data->header_->GetConfigCount(); i++) {
-    ConfigDescription cd;
-    if (!stream.read(reinterpret_cast<char*>(&cd), sizeof(ConfigDescription))) {
+  std::vector<ConfigDescription> configurations(data->header_->GetConfigCount());
+  if (!configurations.empty()) {
+    if (!stream.read(reinterpret_cast<char*>(&configurations.front()),
+                     sizeof(configurations.front()) * configurations.size())) {
       return nullptr;
     }
-    configurations.emplace_back(cd);
   }
 
   // Construct complete target inline entries
-  for (auto [target_entry, entry_offset, entry_count] : target_inline_entries) {
-    for(size_t i = 0; i < entry_count; i++) {
-      const auto& target_value = target_values[entry_offset + i];
-      const auto& config = configurations[target_value.first];
-      target_entry.values[config] = target_value.second;
+  data->target_inline_entries_.reserve(target_inline_entries.size());
+  for (auto&& entry_header : target_inline_entries) {
+    TargetInlineEntry& entry = data->target_inline_entries_.emplace_back();
+    entry.target_id = entry_header.target_id;
+    for (size_t i = 0; i < entry_header.values_count; i++) {
+      const auto& value_header = target_values[entry_header.values_offset + i];
+      const auto& config = configurations[value_header.config_index];
+      auto& value = entry.values[config];
+      value.data_type = value_header.data_type;
+      value.data_value = value_header.data_value;
     }
-    data->target_inline_entries_.emplace_back(target_entry);
   }
 
   // Read the mapping of overlay resource id to target resource id.
+  data->overlay_entries_.resize(data->header_->GetOverlayEntryCount());
   for (size_t i = 0; i < data->header_->GetOverlayEntryCount(); i++) {
-    OverlayEntry overlay_entry{};
-    if (!Read32(stream, &overlay_entry.overlay_id) || !Read32(stream, &overlay_entry.target_id)) {
+    if (!Read32(stream, &data->overlay_entries_[i].overlay_id)) {
       return nullptr;
     }
-    data->overlay_entries_.emplace_back(overlay_entry);
+  }
+  for (size_t i = 0; i < data->header_->GetOverlayEntryCount(); i++) {
+    if (!Read32(stream, &data->overlay_entries_[i].target_id)) {
+      return nullptr;
+    }
   }
 
   // Read raw string pool bytes.
@@ -320,7 +338,7 @@
   std::unique_ptr<IdmapData> data(new IdmapData());
   data->string_pool_data_ = std::string(resource_mapping.GetStringPoolData());
   uint32_t inline_value_count = 0;
-  std::set<std::string> config_set;
+  std::set<std::string_view> config_set;
   for (const auto& mapping : resource_mapping.GetTargetToOverlayMap()) {
     if (auto overlay_resource = std::get_if<ResourceId>(&mapping.second)) {
       data->target_entries_.push_back({mapping.first, *overlay_resource});
@@ -329,7 +347,9 @@
       for (const auto& [config, value] : std::get<ConfigMap>(mapping.second)) {
         config_set.insert(config);
         ConfigDescription cd;
-        ConfigDescription::Parse(config, &cd);
+        if (!ConfigDescription::Parse(config, &cd)) {
+          return Error("failed to parse configuration string '%s'", config.c_str());
+        }
         values[cd] = value;
         inline_value_count++;
       }
diff --git a/cmds/idmap2/tests/IdmapTests.cpp b/cmds/idmap2/tests/IdmapTests.cpp
index c85619c..1b656e8 100644
--- a/cmds/idmap2/tests/IdmapTests.cpp
+++ b/cmds/idmap2/tests/IdmapTests.cpp
@@ -68,7 +68,7 @@
   std::unique_ptr<const IdmapHeader> header = IdmapHeader::FromBinaryStream(stream);
   ASSERT_THAT(header, NotNull());
   ASSERT_EQ(header->GetMagic(), 0x504d4449U);
-  ASSERT_EQ(header->GetVersion(), 0x09U);
+  ASSERT_EQ(header->GetVersion(), 10);
   ASSERT_EQ(header->GetTargetCrc(), 0x1234U);
   ASSERT_EQ(header->GetOverlayCrc(), 0x5678U);
   ASSERT_EQ(header->GetFulfilledPolicies(), 0x11);
@@ -143,7 +143,7 @@
 
   ASSERT_THAT(idmap->GetHeader(), NotNull());
   ASSERT_EQ(idmap->GetHeader()->GetMagic(), 0x504d4449U);
-  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 0x09U);
+  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 10);
   ASSERT_EQ(idmap->GetHeader()->GetTargetCrc(), 0x1234U);
   ASSERT_EQ(idmap->GetHeader()->GetOverlayCrc(), 0x5678U);
   ASSERT_EQ(idmap->GetHeader()->GetFulfilledPolicies(), kIdmapRawDataPolicies);
@@ -204,7 +204,7 @@
 
   ASSERT_THAT(idmap->GetHeader(), NotNull());
   ASSERT_EQ(idmap->GetHeader()->GetMagic(), 0x504d4449U);
-  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 0x09U);
+  ASSERT_EQ(idmap->GetHeader()->GetVersion(), 10);
   ASSERT_EQ(idmap->GetHeader()->GetTargetCrc(), android::idmap2::TestConstants::TARGET_CRC);
   ASSERT_EQ(idmap->GetHeader()->GetOverlayCrc(), android::idmap2::TestConstants::OVERLAY_CRC);
   ASSERT_EQ(idmap->GetHeader()->GetFulfilledPolicies(), PolicyFlags::PUBLIC);
diff --git a/cmds/idmap2/tests/RawPrintVisitorTests.cpp b/cmds/idmap2/tests/RawPrintVisitorTests.cpp
index 68164e2..7fae1c6 100644
--- a/cmds/idmap2/tests/RawPrintVisitorTests.cpp
+++ b/cmds/idmap2/tests/RawPrintVisitorTests.cpp
@@ -64,7 +64,7 @@
   (*idmap)->accept(&visitor);
 
   ASSERT_CONTAINS_REGEX(ADDRESS "504d4449  magic\n", stream.str());
-  ASSERT_CONTAINS_REGEX(ADDRESS "00000009  version\n", stream.str());
+  ASSERT_CONTAINS_REGEX(ADDRESS "0000000a  version\n", stream.str());
   ASSERT_CONTAINS_REGEX(
       StringPrintf(ADDRESS "%s  target crc\n", android::idmap2::TestConstants::TARGET_CRC_STRING),
       stream.str());
@@ -113,7 +113,7 @@
   (*idmap)->accept(&visitor);
 
   ASSERT_CONTAINS_REGEX(ADDRESS "504d4449  magic\n", stream.str());
-  ASSERT_CONTAINS_REGEX(ADDRESS "00000009  version\n", stream.str());
+  ASSERT_CONTAINS_REGEX(ADDRESS "0000000a  version\n", stream.str());
   ASSERT_CONTAINS_REGEX(ADDRESS "00001234  target crc\n", stream.str());
   ASSERT_CONTAINS_REGEX(ADDRESS "00005678  overlay crc\n", stream.str());
   ASSERT_CONTAINS_REGEX(ADDRESS "00000011  fulfilled policies: public|signature\n", stream.str());
diff --git a/cmds/idmap2/tests/TestHelpers.h b/cmds/idmap2/tests/TestHelpers.h
index bf01c32..2b4ebd1 100644
--- a/cmds/idmap2/tests/TestHelpers.h
+++ b/cmds/idmap2/tests/TestHelpers.h
@@ -34,7 +34,7 @@
     0x49, 0x44, 0x4d, 0x50,
 
     // 0x4: version
-    0x09, 0x00, 0x00, 0x00,
+    0x0a, 0x00, 0x00, 0x00,
 
     // 0x8: target crc
     0x34, 0x12, 0x00, 0x00,
@@ -95,19 +95,15 @@
     // TARGET ENTRIES
     // 0x6c: target id (0x7f020000)
     0x00, 0x00, 0x02, 0x7f,
-
-    // 0x70: overlay_id (0x7f020000)
-    0x00, 0x00, 0x02, 0x7f,
-
-    // 0x74: target id (0x7f030000)
+    // 0x70: target id (0x7f030000)
     0x00, 0x00, 0x03, 0x7f,
-
-    // 0x78: overlay_id (0x7f030000)
-    0x00, 0x00, 0x03, 0x7f,
-
-    // 0x7c: target id (0x7f030002)
+    // 0x74: target id (0x7f030002)
     0x02, 0x00, 0x03, 0x7f,
 
+    // 0x78: overlay_id (0x7f020000)
+    0x00, 0x00, 0x02, 0x7f,
+    // 0x7c: overlay_id (0x7f030000)
+    0x00, 0x00, 0x03, 0x7f,
     // 0x80: overlay_id (0x7f030001)
     0x01, 0x00, 0x03, 0x7f,
 
@@ -178,16 +174,20 @@
     // 0xe1: padding
     0x00, 0x00, 0x00,
 
-
     // OVERLAY ENTRIES
-    // 0xe4: 0x7f020000 -> 0x7f020000
-    0x00, 0x00, 0x02, 0x7f, 0x00, 0x00, 0x02, 0x7f,
+    // 0xe4: 0x7f020000 -> ...
+    0x00, 0x00, 0x02, 0x7f,
+    // 0xe8: 0x7f030000 -> ...
+    0x00, 0x00, 0x03, 0x7f,
+    // 0xec: 0x7f030001 -> ...
+    0x01, 0x00, 0x03, 0x7f,
 
-    // 0xec: 0x7f030000 -> 0x7f030000
-    0x00, 0x00, 0x03, 0x7f, 0x00, 0x00, 0x03, 0x7f,
-
-    // 0xf4: 0x7f030001 -> 0x7f030002
-    0x01, 0x00, 0x03, 0x7f, 0x02, 0x00, 0x03, 0x7f,
+    // 0xf0: ... -> 0x7f020000
+    0x00, 0x00, 0x02, 0x7f,
+    // 0xf4: ... -> 0x7f030000
+    0x00, 0x00, 0x03, 0x7f,
+    // 0xf8: ... -> 0x7f030002
+    0x02, 0x00, 0x03, 0x7f,
 
     // 0xfc: string pool
     // string length,
diff --git a/core/api/current.txt b/core/api/current.txt
index ca45d1b..7a07d58 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -9834,6 +9834,7 @@
   public final class AssociationInfo implements android.os.Parcelable {
     method public int describeContents();
     method @Nullable public android.companion.AssociatedDevice getAssociatedDevice();
+    method @FlaggedApi("android.companion.association_device_icon") @Nullable public android.graphics.drawable.Icon getDeviceIcon();
     method @Nullable public android.net.MacAddress getDeviceMacAddress();
     method @Nullable public String getDeviceProfile();
     method @Nullable public CharSequence getDisplayName();
@@ -9847,6 +9848,7 @@
 
   public final class AssociationRequest implements android.os.Parcelable {
     method public int describeContents();
+    method @FlaggedApi("android.companion.association_device_icon") @Nullable public android.graphics.drawable.Icon getDeviceIcon();
     method @Nullable public String getDeviceProfile();
     method @Nullable public CharSequence getDisplayName();
     method public boolean isForceConfirmation();
@@ -9866,6 +9868,7 @@
     ctor public AssociationRequest.Builder();
     method @NonNull public android.companion.AssociationRequest.Builder addDeviceFilter(@Nullable android.companion.DeviceFilter<?>);
     method @NonNull public android.companion.AssociationRequest build();
+    method @FlaggedApi("android.companion.association_device_icon") @NonNull @RequiresPermission(android.Manifest.permission.REQUEST_COMPANION_SELF_MANAGED) public android.companion.AssociationRequest.Builder setDeviceIcon(@NonNull android.graphics.drawable.Icon);
     method @NonNull public android.companion.AssociationRequest.Builder setDeviceProfile(@NonNull String);
     method @NonNull public android.companion.AssociationRequest.Builder setDisplayName(@NonNull CharSequence);
     method @NonNull @RequiresPermission(android.Manifest.permission.REQUEST_COMPANION_SELF_MANAGED) public android.companion.AssociationRequest.Builder setForceConfirmation(boolean);
@@ -22601,6 +22604,7 @@
     method public void sendEvent(int, int, @Nullable byte[]) throws android.media.MediaCasException;
     method public void setEventListener(@Nullable android.media.MediaCas.EventListener, @Nullable android.os.Handler);
     method public void setPrivateData(@NonNull byte[]) throws android.media.MediaCasException;
+    method @FlaggedApi("com.android.media.flags.update_client_profile_priority") public boolean updateResourcePriority(int, int);
     field public static final int PLUGIN_STATUS_PHYSICAL_MODULE_CHANGED = 0; // 0x0
     field public static final int PLUGIN_STATUS_SESSION_NUMBER_CHANGED = 1; // 0x1
     field public static final int SCRAMBLING_MODE_AES128 = 9; // 0x9
@@ -49254,6 +49258,7 @@
     field public static final String ARG_PROTOCOL = "android.arg.protocol";
     field public static final String ARG_QUANTITY = "android.arg.quantity";
     field public static final String ARG_QUERY_STRING = "android.arg.query_string";
+    field @FlaggedApi("com.android.text.flags.tts_span_duration") public static final String ARG_SECONDS = "android.arg.seconds";
     field public static final String ARG_TEXT = "android.arg.text";
     field public static final String ARG_UNIT = "android.arg.unit";
     field public static final String ARG_USERNAME = "android.arg.username";
@@ -49290,6 +49295,7 @@
     field public static final String TYPE_DATE = "android.type.date";
     field public static final String TYPE_DECIMAL = "android.type.decimal";
     field public static final String TYPE_DIGITS = "android.type.digits";
+    field @FlaggedApi("com.android.text.flags.tts_span_duration") public static final String TYPE_DURATION = "android.type.duration";
     field public static final String TYPE_ELECTRONIC = "android.type.electronic";
     field public static final String TYPE_FRACTION = "android.type.fraction";
     field public static final String TYPE_MEASURE = "android.type.measure";
@@ -49349,6 +49355,13 @@
     method public android.text.style.TtsSpan.DigitsBuilder setDigits(String);
   }
 
+  @FlaggedApi("com.android.text.flags.tts_span_duration") public static class TtsSpan.DurationBuilder extends android.text.style.TtsSpan.SemioticClassBuilder<android.text.style.TtsSpan.DurationBuilder> {
+    ctor @FlaggedApi("com.android.text.flags.tts_span_duration") public TtsSpan.DurationBuilder();
+    method @FlaggedApi("com.android.text.flags.tts_span_duration") @NonNull public android.text.style.TtsSpan.DurationBuilder setHours(int);
+    method @FlaggedApi("com.android.text.flags.tts_span_duration") @NonNull public android.text.style.TtsSpan.DurationBuilder setMinutes(int);
+    method @FlaggedApi("com.android.text.flags.tts_span_duration") @NonNull public android.text.style.TtsSpan.DurationBuilder setSeconds(int);
+  }
+
   public static class TtsSpan.ElectronicBuilder extends android.text.style.TtsSpan.SemioticClassBuilder<android.text.style.TtsSpan.ElectronicBuilder> {
     ctor public TtsSpan.ElectronicBuilder();
     method public android.text.style.TtsSpan.ElectronicBuilder setDomain(String);
@@ -49431,6 +49444,7 @@
     ctor public TtsSpan.TimeBuilder(int, int);
     method public android.text.style.TtsSpan.TimeBuilder setHours(int);
     method public android.text.style.TtsSpan.TimeBuilder setMinutes(int);
+    method @FlaggedApi("com.android.text.flags.tts_span_duration") @NonNull public android.text.style.TtsSpan.TimeBuilder setSeconds(int);
   }
 
   public static class TtsSpan.VerbatimBuilder extends android.text.style.TtsSpan.SemioticClassBuilder<android.text.style.TtsSpan.VerbatimBuilder> {
diff --git a/core/api/module-lib-current.txt b/core/api/module-lib-current.txt
index 287e787..4d1a423 100644
--- a/core/api/module-lib-current.txt
+++ b/core/api/module-lib-current.txt
@@ -321,7 +321,7 @@
 package android.net.wifi {
 
   public final class WifiMigration {
-    method @FlaggedApi("android.net.wifi.flags.legacy_keystore_to_wifi_blobstore_migration_read_only") public static int migrateLegacyKeystoreToWifiBlobstore();
+    method @FlaggedApi("android.net.wifi.flags.legacy_keystore_to_wifi_blobstore_migration_read_only") public static void migrateLegacyKeystoreToWifiBlobstore(@NonNull java.util.concurrent.Executor, @NonNull java.util.function.IntConsumer);
     field @FlaggedApi("android.net.wifi.flags.legacy_keystore_to_wifi_blobstore_migration_read_only") public static final int KEYSTORE_MIGRATION_FAILURE_ENCOUNTERED_EXCEPTION = 2; // 0x2
     field @FlaggedApi("android.net.wifi.flags.legacy_keystore_to_wifi_blobstore_migration_read_only") public static final int KEYSTORE_MIGRATION_SUCCESS_MIGRATION_COMPLETE = 0; // 0x0
     field @FlaggedApi("android.net.wifi.flags.legacy_keystore_to_wifi_blobstore_migration_read_only") public static final int KEYSTORE_MIGRATION_SUCCESS_MIGRATION_NOT_NEEDED = 1; // 0x1
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 76e9ca0..5e4485c 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -914,6 +914,7 @@
     ctor public AssociationInfo.Builder(@NonNull android.companion.AssociationInfo);
     method @NonNull public android.companion.AssociationInfo build();
     method @NonNull public android.companion.AssociationInfo.Builder setAssociatedDevice(@Nullable android.companion.AssociatedDevice);
+    method @FlaggedApi("android.companion.association_device_icon") @NonNull public android.companion.AssociationInfo.Builder setDeviceIcon(@Nullable android.graphics.drawable.Icon);
     method @NonNull public android.companion.AssociationInfo.Builder setDeviceMacAddress(@Nullable android.net.MacAddress);
     method @NonNull public android.companion.AssociationInfo.Builder setDeviceProfile(@Nullable String);
     method @NonNull public android.companion.AssociationInfo.Builder setDisplayName(@Nullable CharSequence);
diff --git a/core/java/Android.bp b/core/java/Android.bp
index 2fa418a..1265de1 100644
--- a/core/java/Android.bp
+++ b/core/java/Android.bp
@@ -647,7 +647,6 @@
 
 java_library {
     name: "protolog-lib",
-    platform_apis: true,
     srcs: [
         "com/android/internal/protolog/ProtoLogImpl.java",
         "com/android/internal/protolog/ProtoLogViewerConfigReader.java",
diff --git a/core/java/android/accessibilityservice/AccessibilityServiceInfo.java b/core/java/android/accessibilityservice/AccessibilityServiceInfo.java
index 8bb2857..5bc7de9 100644
--- a/core/java/android/accessibilityservice/AccessibilityServiceInfo.java
+++ b/core/java/android/accessibilityservice/AccessibilityServiceInfo.java
@@ -84,19 +84,25 @@
  * @attr ref android.R.styleable#AccessibilityService_accessibilityEventTypes
  * @attr ref android.R.styleable#AccessibilityService_accessibilityFeedbackType
  * @attr ref android.R.styleable#AccessibilityService_accessibilityFlags
+ * @attr ref android.R.styleable#AccessibilityService_animatedImageDrawable
+ * @attr ref android.R.styleable#AccessibilityService_canControlMagnification
+ * @attr ref android.R.styleable#AccessibilityService_canPerformGestures
  * @attr ref android.R.styleable#AccessibilityService_canRequestFilterKeyEvents
+ * @attr ref android.R.styleable#AccessibilityService_canRequestFingerprintGestures
  * @attr ref android.R.styleable#AccessibilityService_canRequestTouchExplorationMode
  * @attr ref android.R.styleable#AccessibilityService_canRetrieveWindowContent
- * @attr ref android.R.styleable#AccessibilityService_intro
+ * @attr ref android.R.styleable#AccessibilityService_canTakeScreenshot
  * @attr ref android.R.styleable#AccessibilityService_description
- * @attr ref android.R.styleable#AccessibilityService_summary
+ * @attr ref android.R.styleable#AccessibilityService_htmlDescription
+ * @attr ref android.R.styleable#AccessibilityService_interactiveUiTimeout
+ * @attr ref android.R.styleable#AccessibilityService_intro
+ * @attr ref android.R.styleable#AccessibilityService_isAccessibilityTool
+ * @attr ref android.R.styleable#AccessibilityService_nonInteractiveUiTimeout
  * @attr ref android.R.styleable#AccessibilityService_notificationTimeout
  * @attr ref android.R.styleable#AccessibilityService_packageNames
  * @attr ref android.R.styleable#AccessibilityService_settingsActivity
+ * @attr ref android.R.styleable#AccessibilityService_summary
  * @attr ref android.R.styleable#AccessibilityService_tileService
- * @attr ref android.R.styleable#AccessibilityService_nonInteractiveUiTimeout
- * @attr ref android.R.styleable#AccessibilityService_interactiveUiTimeout
- * @attr ref android.R.styleable#AccessibilityService_canTakeScreenshot
  * @see AccessibilityService
  * @see android.view.accessibility.AccessibilityEvent
  * @see android.view.accessibility.AccessibilityManager
diff --git a/core/java/android/app/WallpaperManager.java b/core/java/android/app/WallpaperManager.java
index c1c96ea..014e4660 100644
--- a/core/java/android/app/WallpaperManager.java
+++ b/core/java/android/app/WallpaperManager.java
@@ -285,6 +285,14 @@
     public static final String COMMAND_UNFREEZE = "android.wallpaper.unfreeze";
 
     /**
+     * Command for {@link #sendWallpaperCommand}: in sendWallpaperCommand put extra to this command
+     * to give the bounds of space between the bottom of notifications and the top of shortcuts
+     * @hide
+     */
+    public static final String COMMAND_LOCKSCREEN_LAYOUT_CHANGED =
+            "android.wallpaper.lockscreen_layout_changed";
+
+    /**
      * Extra passed back from setWallpaper() giving the new wallpaper's assigned ID.
      * @hide
      */
diff --git a/core/java/android/app/wearable/flags.aconfig b/core/java/android/app/wearable/flags.aconfig
index b68bafe..534f461 100644
--- a/core/java/android/app/wearable/flags.aconfig
+++ b/core/java/android/app/wearable/flags.aconfig
@@ -38,4 +38,12 @@
     namespace: "machine_learning"
     description: "This flag enables the APIs related to hotword in WearableSensingManager and WearableSensingService."
     bug: "310055381"
+}
+
+flag {
+    name: "enable_concurrent_wearable_connections"
+    is_exported: true
+    namespace: "machine_learning"
+    description: "This flag enables the APIs for providing multiple concurrent connections to the WearableSensingService."
+    bug: "358133158"
 }
\ No newline at end of file
diff --git a/core/java/android/companion/AssociationInfo.java b/core/java/android/companion/AssociationInfo.java
index b4b96e2..7f30d7c 100644
--- a/core/java/android/companion/AssociationInfo.java
+++ b/core/java/android/companion/AssociationInfo.java
@@ -22,6 +22,7 @@
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
 import android.annotation.UserIdInt;
+import android.graphics.drawable.Icon;
 import android.net.MacAddress;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -86,6 +87,11 @@
     private final int mSystemDataSyncFlags;
 
     /**
+     * A device icon displayed on a selfManaged association dialog.
+     */
+    private final Icon mDeviceIcon;
+
+    /**
      * Creates a new Association.
      *
      * @hide
@@ -95,7 +101,7 @@
             @Nullable CharSequence displayName, @Nullable String deviceProfile,
             @Nullable AssociatedDevice associatedDevice, boolean selfManaged,
             boolean notifyOnDeviceNearby, boolean revoked, boolean pending, long timeApprovedMs,
-            long lastTimeConnectedMs, int systemDataSyncFlags) {
+            long lastTimeConnectedMs, int systemDataSyncFlags, @Nullable Icon deviceIcon) {
         if (id <= 0) {
             throw new IllegalArgumentException("Association ID should be greater than 0");
         }
@@ -119,6 +125,7 @@
         mTimeApprovedMs = timeApprovedMs;
         mLastTimeConnectedMs = lastTimeConnectedMs;
         mSystemDataSyncFlags = systemDataSyncFlags;
+        mDeviceIcon = deviceIcon;
     }
 
     /**
@@ -278,6 +285,20 @@
     }
 
     /**
+     * Get the device icon of the associated device. The device icon represents the device type.
+     *
+     * @return the device icon, or {@code null} if no device icon is has been set for the
+     * associated device.
+     *
+     * @see AssociationRequest.Builder#setDeviceIcon(Icon)
+     */
+    @FlaggedApi(Flags.FLAG_ASSOCIATION_DEVICE_ICON)
+    @Nullable
+    public Icon getDeviceIcon() {
+        return mDeviceIcon;
+    }
+
+    /**
      * Utility method for checking if the association represents a device with the given MAC
      * address.
      *
@@ -370,14 +391,16 @@
                 && Objects.equals(mDisplayName, that.mDisplayName)
                 && Objects.equals(mDeviceProfile, that.mDeviceProfile)
                 && Objects.equals(mAssociatedDevice, that.mAssociatedDevice)
-                && mSystemDataSyncFlags == that.mSystemDataSyncFlags;
+                && mSystemDataSyncFlags == that.mSystemDataSyncFlags
+                && (mDeviceIcon == null ? that.mDeviceIcon == null
+                : mDeviceIcon.sameAs(that.mDeviceIcon));
     }
 
     @Override
     public int hashCode() {
         return Objects.hash(mId, mUserId, mPackageName, mTag, mDeviceMacAddress, mDisplayName,
                 mDeviceProfile, mAssociatedDevice, mSelfManaged, mNotifyOnDeviceNearby, mRevoked,
-                mPending, mTimeApprovedMs, mLastTimeConnectedMs, mSystemDataSyncFlags);
+                mPending, mTimeApprovedMs, mLastTimeConnectedMs, mSystemDataSyncFlags, mDeviceIcon);
     }
 
     @Override
@@ -402,6 +425,12 @@
         dest.writeLong(mTimeApprovedMs);
         dest.writeLong(mLastTimeConnectedMs);
         dest.writeInt(mSystemDataSyncFlags);
+        if (mDeviceIcon != null) {
+            dest.writeInt(1);
+            mDeviceIcon.writeToParcel(dest, flags);
+        } else {
+            dest.writeInt(0);
+        }
     }
 
     private AssociationInfo(@NonNull Parcel in) {
@@ -420,6 +449,11 @@
         mTimeApprovedMs = in.readLong();
         mLastTimeConnectedMs = in.readLong();
         mSystemDataSyncFlags = in.readInt();
+        if (in.readInt() == 1) {
+            mDeviceIcon = Icon.CREATOR.createFromParcel(in);
+        } else {
+            mDeviceIcon = null;
+        }
     }
 
     @NonNull
@@ -459,6 +493,7 @@
         private long mTimeApprovedMs;
         private long mLastTimeConnectedMs;
         private int mSystemDataSyncFlags;
+        private Icon mDeviceIcon;
 
         /** @hide */
         @TestApi
@@ -486,6 +521,7 @@
             mTimeApprovedMs = info.mTimeApprovedMs;
             mLastTimeConnectedMs = info.mLastTimeConnectedMs;
             mSystemDataSyncFlags = info.mSystemDataSyncFlags;
+            mDeviceIcon = info.mDeviceIcon;
         }
 
         /**
@@ -510,6 +546,7 @@
             mTimeApprovedMs = info.mTimeApprovedMs;
             mLastTimeConnectedMs = info.mLastTimeConnectedMs;
             mSystemDataSyncFlags = info.mSystemDataSyncFlags;
+            mDeviceIcon = info.mDeviceIcon;
         }
 
         /** @hide */
@@ -625,6 +662,16 @@
         /** @hide */
         @TestApi
         @NonNull
+        @SuppressLint("MissingGetterMatchingBuilder")
+        @FlaggedApi(Flags.FLAG_ASSOCIATION_DEVICE_ICON)
+        public Builder setDeviceIcon(@Nullable Icon deviceIcon) {
+            mDeviceIcon = deviceIcon;
+            return this;
+        }
+
+        /** @hide */
+        @TestApi
+        @NonNull
         public AssociationInfo build() {
             if (mId <= 0) {
                 throw new IllegalArgumentException("Association ID should be greater than 0");
@@ -648,7 +695,8 @@
                     mPending,
                     mTimeApprovedMs,
                     mLastTimeConnectedMs,
-                    mSystemDataSyncFlags
+                    mSystemDataSyncFlags,
+                    mDeviceIcon
             );
         }
     }
diff --git a/core/java/android/companion/AssociationRequest.java b/core/java/android/companion/AssociationRequest.java
index 2e969f8..41a6791 100644
--- a/core/java/android/companion/AssociationRequest.java
+++ b/core/java/android/companion/AssociationRequest.java
@@ -23,12 +23,14 @@
 import static java.util.Objects.requireNonNull;
 
 import android.Manifest;
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.StringDef;
 import android.annotation.UserIdInt;
 import android.compat.annotation.UnsupportedAppUsage;
+import android.graphics.drawable.Icon;
 import android.os.Build;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -234,6 +236,13 @@
     private boolean mSkipPrompt;
 
     /**
+     * The device icon displayed in selfManaged association dialog.
+     * @hide
+     */
+    @Nullable
+    private Icon mDeviceIcon;
+
+    /**
      * Creates a new AssociationRequest.
      *
      * @param singleDevice
@@ -258,15 +267,16 @@
             @Nullable @DeviceProfile String deviceProfile,
             @Nullable CharSequence displayName,
             boolean selfManaged,
-            boolean forceConfirmation) {
+            boolean forceConfirmation,
+            @Nullable Icon deviceIcon) {
         mSingleDevice = singleDevice;
         mDeviceFilters = requireNonNull(deviceFilters);
         mDeviceProfile = deviceProfile;
         mDisplayName = displayName;
         mSelfManaged = selfManaged;
         mForceConfirmation = forceConfirmation;
-
         mCreationTime = System.currentTimeMillis();
+        mDeviceIcon = deviceIcon;
     }
 
     /**
@@ -318,6 +328,19 @@
         return mSingleDevice;
     }
 
+    /**
+     * Get the device icon of the self-managed association request.
+     *
+     * @return the device icon, or {@code null} if no device icon has been set.
+     *
+     * @see Builder#setDeviceIcon(Icon)
+     */
+    @FlaggedApi(Flags.FLAG_ASSOCIATION_DEVICE_ICON)
+    @Nullable
+    public Icon getDeviceIcon() {
+        return mDeviceIcon;
+    }
+
     /** @hide */
     public void setPackageName(@NonNull String packageName) {
         mPackageName = packageName;
@@ -365,6 +388,7 @@
         private CharSequence mDisplayName;
         private boolean mSelfManaged = false;
         private boolean mForceConfirmation = false;
+        private Icon mDeviceIcon = null;
 
         public Builder() {}
 
@@ -450,6 +474,23 @@
             return this;
         }
 
+        /**
+         * Set the device icon for the self-managed device and this icon will be
+         * displayed in the self-managed association dialog.
+         *
+         * @throws IllegalArgumentException if the icon is not exactly 24dp by 24dp
+         * or if it is {@link Icon#TYPE_URI} or {@link Icon#TYPE_URI_ADAPTIVE_BITMAP}.
+         * @see #setSelfManaged(boolean)
+         */
+        @NonNull
+        @RequiresPermission(REQUEST_COMPANION_SELF_MANAGED)
+        @FlaggedApi(Flags.FLAG_ASSOCIATION_DEVICE_ICON)
+        public Builder setDeviceIcon(@NonNull Icon deviceIcon) {
+            checkNotUsed();
+            mDeviceIcon = requireNonNull(deviceIcon);
+            return this;
+        }
+
         /** @inheritDoc */
         @NonNull
         @Override
@@ -460,7 +501,7 @@
                         + "provide the display name of the device");
             }
             return new AssociationRequest(mSingleDevice, emptyIfNull(mDeviceFilters),
-                    mDeviceProfile, mDisplayName, mSelfManaged, mForceConfirmation);
+                    mDeviceProfile, mDisplayName, mSelfManaged, mForceConfirmation, mDeviceIcon);
         }
     }
 
@@ -561,7 +602,9 @@
                 && Objects.equals(mDeviceProfilePrivilegesDescription,
                         that.mDeviceProfilePrivilegesDescription)
                 && mCreationTime == that.mCreationTime
-                && mSkipPrompt == that.mSkipPrompt;
+                && mSkipPrompt == that.mSkipPrompt
+                && (mDeviceIcon == null ? that.mDeviceIcon == null
+                : mDeviceIcon.sameAs(that.mDeviceIcon));
     }
 
     @Override
@@ -579,6 +622,8 @@
         _hash = 31 * _hash + Objects.hashCode(mDeviceProfilePrivilegesDescription);
         _hash = 31 * _hash + Long.hashCode(mCreationTime);
         _hash = 31 * _hash + Boolean.hashCode(mSkipPrompt);
+        _hash = 31 * _hash + Objects.hashCode(mDeviceIcon);
+
         return _hash;
     }
 
@@ -606,6 +651,12 @@
             dest.writeString8(mDeviceProfilePrivilegesDescription);
         }
         dest.writeLong(mCreationTime);
+        if (mDeviceIcon != null) {
+            dest.writeInt(1);
+            mDeviceIcon.writeToParcel(dest, flags);
+        } else {
+            dest.writeInt(0);
+        }
     }
 
     @Override
@@ -650,6 +701,11 @@
         this.mDeviceProfilePrivilegesDescription = deviceProfilePrivilegesDescription;
         this.mCreationTime = creationTime;
         this.mSkipPrompt = skipPrompt;
+        if (in.readInt() == 1) {
+            mDeviceIcon = Icon.CREATOR.createFromParcel(in);
+        } else {
+            mDeviceIcon = null;
+        }
     }
 
     @NonNull
diff --git a/core/java/android/companion/CompanionDeviceManager.java b/core/java/android/companion/CompanionDeviceManager.java
index 1cdf3b1..dfad6de 100644
--- a/core/java/android/companion/CompanionDeviceManager.java
+++ b/core/java/android/companion/CompanionDeviceManager.java
@@ -20,6 +20,8 @@
 import static android.Manifest.permission.REQUEST_COMPANION_PROFILE_AUTOMOTIVE_PROJECTION;
 import static android.Manifest.permission.REQUEST_COMPANION_PROFILE_COMPUTER;
 import static android.Manifest.permission.REQUEST_COMPANION_PROFILE_WATCH;
+import static android.graphics.drawable.Icon.TYPE_URI;
+import static android.graphics.drawable.Icon.TYPE_URI_ADAPTIVE_BITMAP;
 
 
 import android.annotation.CallbackExecutor;
@@ -49,6 +51,11 @@
 import android.content.Intent;
 import android.content.IntentSender;
 import android.content.pm.PackageManager;
+import android.graphics.Bitmap;
+import android.graphics.drawable.BitmapDrawable;
+import android.graphics.drawable.Drawable;
+import android.graphics.drawable.Icon;
+import android.graphics.drawable.VectorDrawable;
 import android.net.MacAddress;
 import android.os.Binder;
 import android.os.Handler;
@@ -535,6 +542,13 @@
         Objects.requireNonNull(executor, "Executor cannot be null");
         Objects.requireNonNull(callback, "Callback cannot be null");
 
+        final Icon deviceIcon = request.getDeviceIcon();
+
+        if (deviceIcon != null && !isValidIcon(deviceIcon, mContext)) {
+            throw new IllegalArgumentException("The size of the device icon must be 24dp x 24dp to"
+                    + "ensure proper display");
+        }
+
         try {
             mService.associate(request, new AssociationRequestCallbackProxy(executor, callback),
                     mContext.getOpPackageName(), mContext.getUserId());
@@ -2027,4 +2041,34 @@
             }
         }
     }
+
+    private boolean isValidIcon(Icon icon, Context context) {
+        if (icon.getType() == TYPE_URI_ADAPTIVE_BITMAP || icon.getType() == TYPE_URI) {
+            throw new IllegalArgumentException("The URI based Icon is not supported.");
+        }
+        Drawable drawable = icon.loadDrawable(context);
+        float density = context.getResources().getDisplayMetrics().density;
+
+        if (drawable instanceof BitmapDrawable) {
+            Bitmap bitmap = ((BitmapDrawable) drawable).getBitmap();
+
+            float widthDp = bitmap.getWidth() / density;
+            float heightDp = bitmap.getHeight() / density;
+
+            if (widthDp != 24 || heightDp != 24) {
+                return false;
+            }
+        } else if (drawable instanceof VectorDrawable) {
+            VectorDrawable vectorDrawable = (VectorDrawable) drawable;
+            float widthDp = vectorDrawable.getIntrinsicWidth() / density;
+            float heightDp = vectorDrawable.getIntrinsicHeight() / density;
+
+            if (widthDp != 24 || heightDp != 24) {
+                return false;
+            }
+        } else {
+            throw new IllegalArgumentException("The format of the device icon is unsupported.");
+        }
+        return true;
+    }
 }
diff --git a/core/java/android/companion/flags.aconfig b/core/java/android/companion/flags.aconfig
index 93d62cf..2539a12 100644
--- a/core/java/android/companion/flags.aconfig
+++ b/core/java/android/companion/flags.aconfig
@@ -55,3 +55,11 @@
     description: "Enable association failure code API"
     bug: "331459560"
 }
+
+flag {
+    name: "association_device_icon"
+    is_exported: true
+    namespace: "companion"
+    description: "Enable set device icon API"
+    bug: "341057668"
+}
diff --git a/core/java/android/content/pm/multiuser.aconfig b/core/java/android/content/pm/multiuser.aconfig
index fa26837..6be7cc4 100644
--- a/core/java/android/content/pm/multiuser.aconfig
+++ b/core/java/android/content/pm/multiuser.aconfig
@@ -248,6 +248,50 @@
 }
 
 flag {
+    name: "cache_profile_parent_read_only"
+    namespace: "multiuser"
+    description: "Cache getProfileParent to avoid unnecessary binder calls"
+    bug: "350417399"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+  }
+  is_fixed_read_only: true
+}
+
+flag {
+    name: "cache_profile_ids_read_only"
+    namespace: "multiuser"
+    description: "Cache getProfileIds to avoid unnecessary binder calls"
+    bug: "350421409"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+  }
+  is_fixed_read_only: true
+}
+
+flag {
+    name: "cache_profile_type_read_only"
+    namespace: "multiuser"
+    description: "Cache getProfileType to avoid unnecessary binder calls"
+    bug: "350417403"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+  }
+  is_fixed_read_only: true
+}
+
+flag {
+    name: "cache_profiles_read_only"
+    namespace: "multiuser"
+    description: "Cache getProfiles to avoid unnecessary binder calls"
+    bug: "350419395"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+  }
+  is_fixed_read_only: true
+}
+
+flag {
     name: "cache_quiet_mode_state"
     namespace: "multiuser"
     description: "Optimise quiet mode state retrieval"
diff --git a/core/java/android/hardware/devicestate/DeviceStateInfo.java b/core/java/android/hardware/devicestate/DeviceStateInfo.java
index 28561ec..fd6f0b8 100644
--- a/core/java/android/hardware/devicestate/DeviceStateInfo.java
+++ b/core/java/android/hardware/devicestate/DeviceStateInfo.java
@@ -28,7 +28,6 @@
 import java.util.Objects;
 import java.util.concurrent.Executor;
 
-
 /**
  * Information about the state of the device.
  *
@@ -63,11 +62,13 @@
      * ignoring any override requests made through a call to {@link DeviceStateManager#requestState(
      * DeviceStateRequest, Executor, DeviceStateRequest.Callback)}.
      */
+    @NonNull
     public final DeviceState baseState;
 
     /**
      * The state of the device.
      */
+    @NonNull
     public final DeviceState currentState;
 
     /**
@@ -78,8 +79,9 @@
      */
     // Using the specific types to avoid virtual method calls in binder transactions
     @SuppressWarnings("NonApiType")
-    public DeviceStateInfo(@NonNull ArrayList<DeviceState> supportedStates, DeviceState baseState,
-            DeviceState state) {
+    public DeviceStateInfo(@NonNull ArrayList<DeviceState> supportedStates,
+            @NonNull DeviceState baseState,
+            @NonNull DeviceState state) {
         this.supportedStates = supportedStates;
         this.baseState = baseState;
         this.currentState = state;
@@ -97,7 +99,7 @@
     public boolean equals(@Nullable Object other) {
         if (this == other) return true;
         if (other == null || (getClass() != other.getClass())) return false;
-        DeviceStateInfo that = (DeviceStateInfo) other;
+        final DeviceStateInfo that = (DeviceStateInfo) other;
         return baseState.equals(that.baseState)
                 &&  currentState.equals(that.currentState)
                 && Objects.equals(supportedStates, that.supportedStates);
@@ -126,8 +128,16 @@
         return diff;
     }
 
+    // Parcelable implementation
+
     @Override
-    public void writeToParcel(Parcel dest, int flags) {
+    public int describeContents() {
+        return 0;
+    }
+
+    /** Writes to Parcel. */
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
         dest.writeInt(supportedStates.size());
         for (int i = 0; i < supportedStates.size(); i++) {
             dest.writeTypedObject(supportedStates.get(i).getConfiguration(), flags);
@@ -137,28 +147,27 @@
         dest.writeTypedObject(currentState.getConfiguration(), flags);
     }
 
-    @Override
-    public int describeContents() {
-        return 0;
+    /** Reads from Parcel. */
+    private DeviceStateInfo(@NonNull Parcel in) {
+        final int numberOfSupportedStates = in.readInt();
+        final ArrayList<DeviceState> supportedStates = new ArrayList<>(numberOfSupportedStates);
+        for (int i = 0; i < numberOfSupportedStates; i++) {
+            final DeviceState.Configuration configuration =
+                    Objects.requireNonNull(in.readTypedObject(DeviceState.Configuration.CREATOR));
+            supportedStates.add(i, new DeviceState(configuration));
+        }
+        this.supportedStates = supportedStates;
+
+        this.baseState = new DeviceState(
+                Objects.requireNonNull(in.readTypedObject(DeviceState.Configuration.CREATOR)));
+        this.currentState = new DeviceState(
+                Objects.requireNonNull(in.readTypedObject(DeviceState.Configuration.CREATOR)));
     }
 
     public static final @NonNull Creator<DeviceStateInfo> CREATOR = new Creator<>() {
         @Override
-        public DeviceStateInfo createFromParcel(Parcel source) {
-            final int numberOfSupportedStates = source.readInt();
-            final ArrayList<DeviceState> supportedStates = new ArrayList<>(numberOfSupportedStates);
-            for (int i = 0; i < numberOfSupportedStates; i++) {
-                DeviceState.Configuration configuration = source.readTypedObject(
-                        DeviceState.Configuration.CREATOR);
-                supportedStates.add(i, new DeviceState(configuration));
-            }
-
-            final DeviceState baseState = new DeviceState(
-                    source.readTypedObject(DeviceState.Configuration.CREATOR));
-            final DeviceState currentState = new DeviceState(
-                    source.readTypedObject(DeviceState.Configuration.CREATOR));
-
-            return new DeviceStateInfo(supportedStates, baseState, currentState);
+        public DeviceStateInfo createFromParcel(@NonNull Parcel in) {
+            return new DeviceStateInfo(in);
         }
 
         @Override
diff --git a/core/java/android/hardware/devicestate/DeviceStateManagerGlobal.java b/core/java/android/hardware/devicestate/DeviceStateManagerGlobal.java
index 0c84019..1845827 100644
--- a/core/java/android/hardware/devicestate/DeviceStateManagerGlobal.java
+++ b/core/java/android/hardware/devicestate/DeviceStateManagerGlobal.java
@@ -46,6 +46,7 @@
  */
 @VisibleForTesting(visibility = Visibility.PACKAGE)
 public final class DeviceStateManagerGlobal {
+    @Nullable
     private static DeviceStateManagerGlobal sInstance;
     private static final String TAG = "DeviceStateManagerGlobal";
     private static final boolean DEBUG = Build.IS_DEBUGGABLE;
@@ -58,7 +59,7 @@
     public static DeviceStateManagerGlobal getInstance() {
         synchronized (DeviceStateManagerGlobal.class) {
             if (sInstance == null) {
-                IBinder b = ServiceManager.getService(Context.DEVICE_STATE_SERVICE);
+                final IBinder b = ServiceManager.getService(Context.DEVICE_STATE_SERVICE);
                 if (b != null) {
                     sInstance = new DeviceStateManagerGlobal(IDeviceStateManager
                             .Stub.asInterface(b));
@@ -94,6 +95,7 @@
      *
      * @see DeviceStateManager#getSupportedDeviceStates()
      */
+    @NonNull
     public List<DeviceState> getSupportedDeviceStates() {
         synchronized (mLock) {
             final DeviceStateInfo currentInfo;
@@ -126,8 +128,8 @@
             conditional = true)
     public void requestState(@NonNull DeviceStateRequest request,
             @Nullable Executor executor, @Nullable DeviceStateRequest.Callback callback) {
-        DeviceStateRequestWrapper requestWrapper = new DeviceStateRequestWrapper(request, callback,
-                executor);
+        final DeviceStateRequestWrapper requestWrapper =
+                new DeviceStateRequestWrapper(request, callback, executor);
         synchronized (mLock) {
             if (findRequestTokenLocked(request) != null) {
                 // This request has already been submitted.
@@ -136,7 +138,7 @@
             // Add the request wrapper to the mRequests array before requesting the state as the
             // callback could be triggered immediately if the mDeviceStateManager IBinder is in the
             // same process as this instance.
-            IBinder token = new Binder();
+            final IBinder token = new Binder();
             mRequests.put(token, requestWrapper);
 
             try {
@@ -176,8 +178,8 @@
     @RequiresPermission(android.Manifest.permission.CONTROL_DEVICE_STATE)
     public void requestBaseStateOverride(@NonNull DeviceStateRequest request,
             @Nullable Executor executor, @Nullable DeviceStateRequest.Callback callback) {
-        DeviceStateRequestWrapper requestWrapper = new DeviceStateRequestWrapper(request, callback,
-                executor);
+        final DeviceStateRequestWrapper requestWrapper =
+                new DeviceStateRequestWrapper(request, callback, executor);
         synchronized (mLock) {
             if (findRequestTokenLocked(request) != null) {
                 // This request has already been submitted.
@@ -186,7 +188,7 @@
             // Add the request wrapper to the mRequests array before requesting the state as the
             // callback could be triggered immediately if the mDeviceStateManager IBinder is in the
             // same process as this instance.
-            IBinder token = new Binder();
+            final IBinder token = new Binder();
             mRequests.put(token, requestWrapper);
 
             try {
@@ -225,7 +227,7 @@
     public void registerDeviceStateCallback(@NonNull DeviceStateCallback callback,
             @NonNull Executor executor) {
         synchronized (mLock) {
-            int index = findCallbackLocked(callback);
+            final int index = findCallbackLocked(callback);
             if (index != -1) {
                 // This callback is already registered.
                 return;
@@ -233,7 +235,8 @@
             // Add the callback wrapper to the mCallbacks array after registering the callback as
             // the callback could be triggered immediately if the mDeviceStateManager IBinder is in
             // the same process as this instance.
-            DeviceStateCallbackWrapper wrapper = new DeviceStateCallbackWrapper(callback, executor);
+            final DeviceStateCallbackWrapper wrapper =
+                    new DeviceStateCallbackWrapper(callback, executor);
             mCallbacks.add(wrapper);
 
             if (mLastReceivedInfo != null) {
@@ -253,7 +256,7 @@
     @VisibleForTesting(visibility = Visibility.PACKAGE)
     public void unregisterDeviceStateCallback(@NonNull DeviceStateCallback callback) {
         synchronized (mLock) {
-            int indexToRemove = findCallbackLocked(callback);
+            final int indexToRemove = findCallbackLocked(callback);
             if (indexToRemove != -1) {
                 mCallbacks.remove(indexToRemove);
             }
@@ -277,14 +280,16 @@
     }
 
     private void registerCallbackIfNeededLocked() {
-        if (mCallback == null) {
-            mCallback = new DeviceStateManagerCallback();
-            try {
-                mDeviceStateManager.registerCallback(mCallback);
-            } catch (RemoteException ex) {
-                mCallback = null;
-                throw ex.rethrowFromSystemServer();
-            }
+        if (mCallback != null) {
+            return;
+        }
+
+        mCallback = new DeviceStateManagerCallback();
+        try {
+            mDeviceStateManager.registerCallback(mCallback);
+        } catch (RemoteException ex) {
+            mCallback = null;
+            throw ex.rethrowFromSystemServer();
         }
     }
 
@@ -298,6 +303,7 @@
     }
 
     @Nullable
+    @GuardedBy("mLock")
     private IBinder findRequestTokenLocked(@NonNull DeviceStateRequest request) {
         for (int i = 0; i < mRequests.size(); i++) {
             if (mRequests.valueAt(i).mRequest.equals(request)) {
@@ -309,8 +315,8 @@
 
     /** Handles a call from the server that the device state info has changed. */
     private void handleDeviceStateInfoChanged(@NonNull DeviceStateInfo info) {
-        ArrayList<DeviceStateCallbackWrapper> callbacks;
-        DeviceStateInfo oldInfo;
+        final ArrayList<DeviceStateCallbackWrapper> callbacks;
+        final DeviceStateInfo oldInfo;
         synchronized (mLock) {
             oldInfo = mLastReceivedInfo;
             mLastReceivedInfo = info;
@@ -335,8 +341,8 @@
      * Handles a call from the server that a request for the supplied {@code token} has become
      * active.
      */
-    private void handleRequestActive(IBinder token) {
-        DeviceStateRequestWrapper request;
+    private void handleRequestActive(@NonNull IBinder token) {
+        final DeviceStateRequestWrapper request;
         synchronized (mLock) {
             request = mRequests.get(token);
         }
@@ -349,8 +355,8 @@
      * Handles a call from the server that a request for the supplied {@code token} has become
      * canceled.
      */
-    private void handleRequestCanceled(IBinder token) {
-        DeviceStateRequestWrapper request;
+    private void handleRequestCanceled(@NonNull IBinder token) {
+        final DeviceStateRequestWrapper request;
         synchronized (mLock) {
             request = mRequests.remove(token);
         }
@@ -361,17 +367,17 @@
 
     private final class DeviceStateManagerCallback extends IDeviceStateManagerCallback.Stub {
         @Override
-        public void onDeviceStateInfoChanged(DeviceStateInfo info) {
+        public void onDeviceStateInfoChanged(@NonNull DeviceStateInfo info) {
             handleDeviceStateInfoChanged(info);
         }
 
         @Override
-        public void onRequestActive(IBinder token) {
+        public void onRequestActive(@NonNull IBinder token) {
             handleRequestActive(token);
         }
 
         @Override
-        public void onRequestCanceled(IBinder token) {
+        public void onRequestCanceled(@NonNull IBinder token) {
             handleRequestCanceled(token);
         }
     }
@@ -388,7 +394,8 @@
             mExecutor = executor;
         }
 
-        void notifySupportedDeviceStatesChanged(List<DeviceState> newSupportedDeviceStates) {
+        void notifySupportedDeviceStatesChanged(
+                @NonNull List<DeviceState> newSupportedDeviceStates) {
             mExecutor.execute(() ->
                     mDeviceStateCallback.onSupportedStatesChanged(newSupportedDeviceStates));
         }
@@ -398,7 +405,7 @@
                     () -> mDeviceStateCallback.onDeviceStateChanged(newDeviceState));
         }
 
-        private void execute(String traceName, Runnable r) {
+        private void execute(@NonNull String traceName, @NonNull Runnable r) {
             mExecutor.execute(() -> {
                 if (DEBUG) {
                     Trace.beginSection(
@@ -416,6 +423,7 @@
     }
 
     private static final class DeviceStateRequestWrapper {
+        @NonNull
         private final DeviceStateRequest mRequest;
         @Nullable
         private final DeviceStateRequest.Callback mCallback;
diff --git a/core/java/android/os/SystemClock.java b/core/java/android/os/SystemClock.java
index dfc591b..8ef62e3 100644
--- a/core/java/android/os/SystemClock.java
+++ b/core/java/android/os/SystemClock.java
@@ -62,7 +62,7 @@
  *     sleep (CPU off, display dark, device waiting for external input),
  *     but is not affected by clock scaling, idle, or other power saving
  *     mechanisms.  This is the basis for most interval timing
- *     such as {@link Thread#sleep(long) Thread.sleep(millls)},
+ *     such as {@link Thread#sleep(long) Thread.sleep(millis)},
  *     {@link Object#wait(long) Object.wait(millis)}, and
  *     {@link System#nanoTime System.nanoTime()}.  This clock is guaranteed
  *     to be monotonic, and is suitable for interval timing when the
diff --git a/core/java/android/permission/flags.aconfig b/core/java/android/permission/flags.aconfig
index bca5bcc..feeb339 100644
--- a/core/java/android/permission/flags.aconfig
+++ b/core/java/android/permission/flags.aconfig
@@ -206,17 +206,6 @@
 }
 
 flag {
-    name: "apex_signature_permission_allowlist_enabled"
-    is_fixed_read_only: true
-    namespace: "permissions"
-    description: "Enable reading signature permission allowlist from APEXes"
-    bug: "308573169"
-    metadata {
-        purpose: PURPOSE_BUGFIX
-    }
-}
-
-flag {
     name: "check_op_validate_package"
     namespace: "permissions"
     description: "Validate package/uid match in checkOp similar to noteOp"
diff --git a/core/java/android/provider/ContactsContract.java b/core/java/android/provider/ContactsContract.java
index 98d58d0..5ecf361 100644
--- a/core/java/android/provider/ContactsContract.java
+++ b/core/java/android/provider/ContactsContract.java
@@ -3209,7 +3209,11 @@
                     return new DefaultAccountAndState(DEFAULT_ACCOUNT_STATE_NOT_SET, null);
                 }
 
-                private static boolean isCloudOrSimAccount(@DefaultAccountState int state) {
+                /**
+                 *
+                 * @hide
+                 */
+                public static boolean isCloudOrSimAccount(@DefaultAccountState int state) {
                     return state == DEFAULT_ACCOUNT_STATE_CLOUD
                             || state == DEFAULT_ACCOUNT_STATE_SIM;
                 }
@@ -3287,23 +3291,20 @@
                 Bundle response = nullSafeCall(resolver, ContactsContract.AUTHORITY_URI,
                         QUERY_DEFAULT_ACCOUNT_FOR_NEW_CONTACTS_METHOD, null, null);
 
-                int defaultContactsAccountState = response.getInt(KEY_DEFAULT_ACCOUNT_STATE, -1);
-                if (defaultContactsAccountState
-                        == DefaultAccountAndState.DEFAULT_ACCOUNT_STATE_CLOUD) {
+                int defaultAccountState = response.getInt(KEY_DEFAULT_ACCOUNT_STATE, -1);
+                if (DefaultAccountAndState.isCloudOrSimAccount(defaultAccountState)) {
                     String accountName = response.getString(Settings.ACCOUNT_NAME);
                     String accountType = response.getString(Settings.ACCOUNT_TYPE);
                     if (TextUtils.isEmpty(accountName) || TextUtils.isEmpty(accountType)) {
                         throw new IllegalStateException(
                                 "account name and type cannot be null or empty");
                     }
-                    return new DefaultAccountAndState(
-                            DefaultAccountAndState.DEFAULT_ACCOUNT_STATE_CLOUD,
+                    return new DefaultAccountAndState(defaultAccountState,
                             new Account(accountName, accountType));
-                } else if (defaultContactsAccountState
-                        == DefaultAccountAndState.DEFAULT_ACCOUNT_STATE_LOCAL
-                        || defaultContactsAccountState
+                } else if (defaultAccountState == DefaultAccountAndState.DEFAULT_ACCOUNT_STATE_LOCAL
+                        || defaultAccountState
                         == DefaultAccountAndState.DEFAULT_ACCOUNT_STATE_NOT_SET) {
-                    return new DefaultAccountAndState(defaultContactsAccountState, /*cloudAccount=*/
+                    return new DefaultAccountAndState(defaultAccountState, /*account=*/
                             null);
                 } else {
                     throw new IllegalStateException("Invalid default account state");
@@ -3348,12 +3349,11 @@
                 Bundle extras = new Bundle();
 
                 extras.putInt(KEY_DEFAULT_ACCOUNT_STATE, defaultAccountAndState.getState());
-                if (defaultAccountAndState.getState()
-                        == DefaultAccountAndState.DEFAULT_ACCOUNT_STATE_CLOUD) {
-                    Account cloudAccount = defaultAccountAndState.getAccount();
-                    assert cloudAccount != null;
-                    extras.putString(Settings.ACCOUNT_NAME, cloudAccount.name);
-                    extras.putString(Settings.ACCOUNT_TYPE, cloudAccount.type);
+                if (DefaultAccountAndState.isCloudOrSimAccount(defaultAccountAndState.getState())) {
+                    Account account = defaultAccountAndState.getAccount();
+                    assert account != null;
+                    extras.putString(Settings.ACCOUNT_NAME, account.name);
+                    extras.putString(Settings.ACCOUNT_TYPE, account.type);
                 }
                 nullSafeCall(resolver, ContactsContract.AUTHORITY_URI,
                         SET_DEFAULT_ACCOUNT_FOR_NEW_CONTACTS_METHOD, null, extras);
diff --git a/core/java/android/text/flags/flags.aconfig b/core/java/android/text/flags/flags.aconfig
index ec5b488..09b2201 100644
--- a/core/java/android/text/flags/flags.aconfig
+++ b/core/java/android/text/flags/flags.aconfig
@@ -188,3 +188,10 @@
     purpose: PURPOSE_BUGFIX
   }
 }
+
+flag {
+  name: "tts_span_duration"
+  namespace: "text"
+  description: "Feature flag for adding a TYPE_DURATION to TtsSpan"
+  bug: "337103893"
+}
diff --git a/core/java/android/text/style/TtsSpan.java b/core/java/android/text/style/TtsSpan.java
index f9a1a0d..b7b8f0b1 100644
--- a/core/java/android/text/style/TtsSpan.java
+++ b/core/java/android/text/style/TtsSpan.java
@@ -16,6 +16,10 @@
 
 package android.text.style;
 
+import static com.android.text.flags.Flags.FLAG_TTS_SPAN_DURATION;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
 import android.os.Parcel;
 import android.os.PersistableBundle;
 import android.text.ParcelableSpan;
@@ -112,6 +116,16 @@
     public static final String TYPE_TIME = "android.type.time";
 
     /**
+     * The text associated with this span is a duration, consisting of a number of
+     * hours, minutes, and seconds specified with {@link #ARG_HOURS},
+     * {@link #ARG_MINUTES}, and {@link #ARG_SECONDS}. This is different from {@link #TYPE_TIME}.
+     * This should be used to convey an interval of time, while {@link #TYPE_TIME} should be used to
+     * convey a particular moment in time, such as a clock time.
+     */
+    @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+    public static final String TYPE_DURATION = "android.type.duration";
+
+    /**
      * The text associated with this span is a date. At least one of the
      * arguments {@link #ARG_MONTH} and {@link #ARG_YEAR} has to be provided.
      * The argument {@link #ARG_DAY} is optional if {@link #ARG_MONTH} is set.
@@ -302,13 +316,21 @@
     public static final String ARG_HOURS = "android.arg.hours";
 
     /**
-     * Argument used to specify the minutes of a time. The hours should be
+     * Argument used to specify the minutes of a time. The minutes should be
      * provided as an integer in the range from 0 up to and including 59.
      * Can be used with {@link #TYPE_TIME}.
      */
     public static final String ARG_MINUTES = "android.arg.minutes";
 
     /**
+     * Argument used to specify the seconds of a time or duration. The seconds should be
+     * provided as an integer in the range from 0 up to and including 59.
+     * Can be used with {@link #TYPE_TIME} or {@link #TYPE_DURATION}.
+     */
+    @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+    public static final String ARG_SECONDS = "android.arg.seconds";
+
+    /**
      * Argument used to specify the weekday of a date. The value should be
      * provided as an integer and can be any of {@link #WEEKDAY_SUNDAY},
      * {@link #WEEKDAY_MONDAY}, {@link #WEEKDAY_TUESDAY},
@@ -1132,9 +1154,70 @@
         public TimeBuilder setMinutes(int minutes) {
             return setIntArgument(TtsSpan.ARG_MINUTES, minutes);
         }
+
+        /**
+         * Sets the {@link #ARG_SECONDS} argument.
+         * @param seconds The value to be set for seconds.
+         * @return This instance.
+         */
+        @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+        @NonNull
+        public TimeBuilder setSeconds(int seconds) {
+            return setIntArgument(TtsSpan.ARG_SECONDS, seconds);
+        }
     }
 
     /**
+     * A builder for TtsSpans of type {@link #TYPE_DURATION}.
+     */
+    @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+    public static class DurationBuilder
+            extends SemioticClassBuilder<DurationBuilder> {
+
+        /**
+         * Creates a builder for a TtsSpan of type {@link #TYPE_DURATION}.
+         */
+        @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+        public DurationBuilder() {
+            super(TtsSpan.TYPE_DURATION);
+        }
+
+        /**
+         * Sets the {@link #ARG_HOURS} argument.
+         * @param hours The value to be set for hours.
+         * @return This instance.
+         */
+        @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+        @NonNull
+        public DurationBuilder setHours(int hours) {
+            return setIntArgument(TtsSpan.ARG_HOURS, hours);
+        }
+
+        /**
+         * Sets the {@link #ARG_MINUTES} argument.
+         * @param minutes The value to be set for minutes.
+         * @return This instance.
+         */
+        @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+        @NonNull
+        public DurationBuilder setMinutes(int minutes) {
+            return setIntArgument(TtsSpan.ARG_MINUTES, minutes);
+        }
+
+        /**
+         * Sets the {@link #ARG_SECONDS} argument.
+         * @param seconds The value to be set for seconds.
+         * @return This instance.
+         */
+        @FlaggedApi(FLAG_TTS_SPAN_DURATION)
+        @NonNull
+        public DurationBuilder setSeconds(int seconds) {
+            return setIntArgument(TtsSpan.ARG_SECONDS, seconds);
+        }
+    }
+
+
+    /**
      * A builder for TtsSpans of type {@link #TYPE_DATE}.
      */
     public static class DateBuilder
diff --git a/core/java/android/view/InsetsState.java b/core/java/android/view/InsetsState.java
index 5c41516..67050e0 100644
--- a/core/java/android/view/InsetsState.java
+++ b/core/java/android/view/InsetsState.java
@@ -38,7 +38,7 @@
 import static android.view.WindowManager.LayoutParams.SOFT_INPUT_MASK_ADJUST;
 import static android.view.WindowManager.LayoutParams.TYPE_SYSTEM_ERROR;
 import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
-import static android.window.flags.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION_ALWAYS;
+import static android.window.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION_ALWAYS;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 0ca442d..5afc7b2 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -125,7 +125,7 @@
 import static android.view.flags.Flags.toolkitSetFrameRateReadOnly;
 import static android.view.inputmethod.InputMethodEditorTraceProto.InputMethodClientsTraceProto.ClientSideProto.IME_FOCUS_CONTROLLER;
 import static android.view.inputmethod.InputMethodEditorTraceProto.InputMethodClientsTraceProto.ClientSideProto.INSETS_CONTROLLER;
-import static android.window.flags.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
+import static android.window.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PACKAGE;
 import static com.android.text.flags.Flags.disableHandwritingInitiatorForIme;
diff --git a/core/java/android/view/accessibility/flags/accessibility_flags.aconfig b/core/java/android/view/accessibility/flags/accessibility_flags.aconfig
index b9e9750..69cbb9b 100644
--- a/core/java/android/view/accessibility/flags/accessibility_flags.aconfig
+++ b/core/java/android/view/accessibility/flags/accessibility_flags.aconfig
@@ -11,6 +11,13 @@
 }
 
 flag {
+    name: "a11y_is_required_api"
+    namespace: "accessibility"
+    description: "Adds an API to indicate whether a form field (or similar element) is required."
+    bug: "362784403"
+}
+
+flag {
     name: "a11y_overlay_callbacks"
     is_exported: true
     namespace: "accessibility"
@@ -210,3 +217,13 @@
     description: "Feature flag for declaring system pinch zoom opt-out apis"
     bug: "315089687"
 }
+
+flag {
+    name: "warning_use_default_dialog_type"
+    namespace: "accessibility"
+    description: "Uses the default type for the A11yService warning dialog, instead of SYSTEM_ALERT_DIALOG"
+    bug: "336719951"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/core/java/android/view/flags/view_flags.aconfig b/core/java/android/view/flags/view_flags.aconfig
index f570a9a..5115b13 100644
--- a/core/java/android/view/flags/view_flags.aconfig
+++ b/core/java/android/view/flags/view_flags.aconfig
@@ -97,4 +97,15 @@
     description: "Disable Draw Wakelock starting U."
     bug: "331698645"
     is_fixed_read_only: true
-}
\ No newline at end of file
+}
+
+flag {
+  name: "calculate_bounds_in_parent_from_bounds_in_screen"
+  namespace: "accessibility"
+  description: "Calculate bounds in parent of each node in ViewStructure from its bounds set relative to screen and its parent's"
+  bug: "366131857"
+  is_fixed_read_only: true
+  metadata {
+      purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/core/java/android/view/inputmethod/InputMethodInfo.java b/core/java/android/view/inputmethod/InputMethodInfo.java
index 806a593..0a83bdc 100644
--- a/core/java/android/view/inputmethod/InputMethodInfo.java
+++ b/core/java/android/view/inputmethod/InputMethodInfo.java
@@ -860,7 +860,7 @@
      * <p>e.g.<pre><code>startActivity(createStylusHandwritingSettingsActivityIntent());</code>
      * </pre></p>
      *
-     * @attr ref R.styleable#InputMethod_stylusHandwritingSettingsActivity
+     * @attr ref android.R.styleable#InputMethod_stylusHandwritingSettingsActivity
      * @see #getSettingsActivity()
      * @see #supportsStylusHandwriting()
      */
@@ -888,8 +888,8 @@
      * the IME language settings activity.</p>
      * <p>e.g.<pre><code>startActivity(createImeLanguageSettingsActivityIntent());</code></pre></p>
      *
-     * @attr ref R.styleable#InputMethod_languageSettingsActivity
-     * @attr ref R.styleable#InputMethod_settingsActivity
+     * @attr ref android.R.styleable#InputMethod_languageSettingsActivity
+     * @attr ref android.R.styleable#InputMethod_settingsActivity
      */
     @FlaggedApi(android.view.inputmethod.Flags.FLAG_IME_SWITCHER_REVAMP_API)
     @Nullable
diff --git a/core/java/android/widget/ListView.java b/core/java/android/widget/ListView.java
index 510a92d..3f611c7 100644
--- a/core/java/android/widget/ListView.java
+++ b/core/java/android/widget/ListView.java
@@ -3636,7 +3636,7 @@
      * Returns the drawable that will be drawn between each item in the list.
      *
      * @return the current drawable drawn between list elements
-     * @attr ref R.styleable#ListView_divider
+     * @attr ref android.R.styleable#ListView_divider
      */
     @InspectableProperty
     @Nullable
@@ -3651,7 +3651,7 @@
      * height, you should also call {@link #setDividerHeight(int)}.
      *
      * @param divider the drawable to use
-     * @attr ref R.styleable#ListView_divider
+     * @attr ref android.R.styleable#ListView_divider
      */
     public void setDivider(@Nullable Drawable divider) {
         if (divider != null) {
diff --git a/core/java/android/widget/RemoteViews.java b/core/java/android/widget/RemoteViews.java
index 06820cd..d7b5211 100644
--- a/core/java/android/widget/RemoteViews.java
+++ b/core/java/android/widget/RemoteViews.java
@@ -46,6 +46,7 @@
 import android.app.PendingIntent;
 import android.app.RemoteInput;
 import android.appwidget.AppWidgetHostView;
+import android.appwidget.flags.Flags;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.ComponentName;
 import android.content.Context;
@@ -4119,6 +4120,71 @@
         public void visitUris(@NonNull Consumer<Uri> visitor) {
             mNestedViews.visitUris(visitor);
         }
+
+        @Override
+        public boolean canWriteToProto() {
+            return true;
+        }
+
+        @Override
+        public void writeToProto(ProtoOutputStream out, Context context, Resources appResources) {
+            if (!Flags.remoteViewsProto()) return;
+            final long token = out.start(RemoteViewsProto.Action.VIEW_GROUP_ADD_ACTION);
+            out.write(RemoteViewsProto.ViewGroupAddAction.VIEW_ID,
+                    appResources.getResourceName(mViewId));
+            out.write(RemoteViewsProto.ViewGroupAddAction.INDEX, mIndex);
+            out.write(RemoteViewsProto.ViewGroupAddAction.STABLE_ID, mStableId);
+            long rvToken = out.start(RemoteViewsProto.ViewGroupAddAction.NESTED_VIEWS);
+            mNestedViews.writePreviewToProto(context, out);
+            out.end(rvToken);
+            out.end(token);
+        }
+    }
+
+    private PendingResources<Action> createViewGroupActionAddFromProto(ProtoInputStream in)
+            throws Exception {
+        final LongSparseArray<Object> values = new LongSparseArray<>();
+
+        final long token = in.start(RemoteViewsProto.Action.VIEW_GROUP_ADD_ACTION);
+        while (in.nextField() != NO_MORE_FIELDS) {
+            switch (in.getFieldNumber()) {
+                case (int) RemoteViewsProto.ViewGroupAddAction.VIEW_ID:
+                    values.put(RemoteViewsProto.ViewGroupAddAction.VIEW_ID,
+                            in.readString(RemoteViewsProto.ViewGroupAddAction.VIEW_ID));
+                    break;
+                case (int) RemoteViewsProto.ViewGroupAddAction.NESTED_VIEWS:
+                    final long nvToken = in.start(RemoteViewsProto.ViewGroupAddAction.NESTED_VIEWS);
+                    values.put(RemoteViewsProto.ViewGroupAddAction.NESTED_VIEWS,
+                            createFromProto(in));
+                    in.end(nvToken);
+                    break;
+                case (int) RemoteViewsProto.ViewGroupAddAction.INDEX:
+                    values.put(RemoteViewsProto.ViewGroupAddAction.INDEX,
+                            in.readInt(RemoteViewsProto.ViewGroupAddAction.INDEX));
+                    break;
+                case (int) RemoteViewsProto.ViewGroupAddAction.STABLE_ID:
+                    values.put(RemoteViewsProto.ViewGroupAddAction.STABLE_ID,
+                            in.readInt(RemoteViewsProto.ViewGroupAddAction.STABLE_ID));
+                    break;
+                default:
+                    Log.w(LOG_TAG, "Unhandled field while reading RemoteViews proto!\n"
+                            + ProtoUtils.currentFieldToString(in));
+            }
+        }
+        in.end(token);
+
+        checkContainsKeys(values, new long[]{RemoteViewsProto.ViewGroupAddAction.VIEW_ID,
+                RemoteViewsProto.ViewGroupAddAction.NESTED_VIEWS});
+
+        return (context, resources, rootData, depth) -> {
+            int viewId = getAsIdentifier(resources, values,
+                    RemoteViewsProto.ViewGroupAddAction.VIEW_ID);
+            return new ViewGroupActionAdd(viewId, ((PendingResources<RemoteViews>) values.get(
+                    RemoteViewsProto.ViewGroupAddAction.NESTED_VIEWS)).create(context, resources,
+                    rootData, depth),
+                    (int) values.get(RemoteViewsProto.ViewGroupAddAction.INDEX, 0),
+                    (int) values.get(RemoteViewsProto.ViewGroupAddAction.STABLE_ID, 0));
+        };
     }
 
     /**
@@ -4241,6 +4307,60 @@
         public int mergeBehavior() {
             return MERGE_APPEND;
         }
+
+        @Override
+        public boolean canWriteToProto() {
+            return true;
+        }
+
+        @Override
+        public void writeToProto(ProtoOutputStream out, Context context, Resources appResources) {
+            final long token = out.start(RemoteViewsProto.Action.VIEW_GROUP_REMOVE_ACTION);
+            out.write(RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID,
+                    appResources.getResourceName(mViewId));
+            if (mViewIdToKeep != REMOVE_ALL_VIEWS_ID) {
+                out.write(RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID_TO_KEEP,
+                        appResources.getResourceName(mViewIdToKeep));
+            }
+            out.end(token);
+        }
+
+        public static PendingResources<Action> createFromProto(ProtoInputStream in)
+                throws Exception {
+            final LongSparseArray<Object> values = new LongSparseArray<>();
+
+            final long token = in.start(RemoteViewsProto.Action.VIEW_GROUP_REMOVE_ACTION);
+            while (in.nextField() != NO_MORE_FIELDS) {
+                switch (in.getFieldNumber()) {
+                    case (int) RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID:
+                        values.put(RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID,
+                                in.readString(RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID));
+                        break;
+                    case (int) RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID_TO_KEEP:
+                        values.put(RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID_TO_KEEP,
+                                in.readString(
+                                        RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID_TO_KEEP));
+                        break;
+                    default:
+                        Log.w(LOG_TAG, "Unhandled field while reading RemoteViews proto!\n"
+                                + ProtoUtils.currentFieldToString(in));
+                }
+            }
+            in.end(token);
+
+            checkContainsKeys(values, new long[]{RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID});
+
+            return (context, resources, rootData, depth) -> {
+                int viewId = getAsIdentifier(resources, values,
+                        RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID);
+                int viewIdToKeep = (values.indexOfKey(
+                        RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID_TO_KEEP) >= 0)
+                        ? getAsIdentifier(resources, values,
+                        RemoteViewsProto.ViewGroupRemoveAction.VIEW_ID_TO_KEEP)
+                        : REMOVE_ALL_VIEWS_ID;
+                return new ViewGroupActionRemove(viewId, viewIdToKeep);
+            };
+        }
     }
 
     /**
@@ -4497,6 +4617,200 @@
                 visitIconUri(mI4, visitor);
             }
         }
+
+        @Override
+        public boolean canWriteToProto() {
+            return true;
+        }
+
+        @Override
+        public void writeToProto(ProtoOutputStream out, Context context,
+                Resources appResources) { // rebase
+            final long token = out.start(RemoteViewsProto.Action.TEXT_VIEW_DRAWABLE_ACTION);
+            out.write(RemoteViewsProto.TextViewDrawableAction.VIEW_ID,
+                    appResources.getResourceName(mViewId));
+            out.write(RemoteViewsProto.TextViewDrawableAction.IS_RELATIVE, mIsRelative);
+            if (mUseIcons) {
+                long iconsToken = out.start(RemoteViewsProto.TextViewDrawableAction.ICONS);
+                if (mI1 != null) {
+                    writeIconToProto(out, appResources, mI1,
+                            RemoteViewsProto.TextViewDrawableAction.Icons.ONE);
+                }
+                if (mI2 != null) {
+                    writeIconToProto(out, appResources, mI2,
+                            RemoteViewsProto.TextViewDrawableAction.Icons.TWO);
+                }
+                if (mI3 != null) {
+                    writeIconToProto(out, appResources, mI3,
+                            RemoteViewsProto.TextViewDrawableAction.Icons.THREE);
+                }
+                if (mI4 != null) {
+                    writeIconToProto(out, appResources, mI4,
+                            RemoteViewsProto.TextViewDrawableAction.Icons.FOUR);
+                }
+                out.end(iconsToken);
+            } else {
+                long resourcesToken = out.start(RemoteViewsProto.TextViewDrawableAction.RESOURCES);
+                if (mD1 != 0) {
+                    out.write(RemoteViewsProto.TextViewDrawableAction.Resources.ONE,
+                            appResources.getResourceName(mD1));
+                }
+                if (mD2 != 0) {
+                    out.write(RemoteViewsProto.TextViewDrawableAction.Resources.TWO,
+                            appResources.getResourceName(mD2));
+                }
+                if (mD3 != 0) {
+                    out.write(RemoteViewsProto.TextViewDrawableAction.Resources.THREE,
+                            appResources.getResourceName(mD3));
+                }
+                if (mD4 != 0) {
+                    out.write(RemoteViewsProto.TextViewDrawableAction.Resources.FOUR,
+                            appResources.getResourceName(mD4));
+                }
+                out.end(resourcesToken);
+            }
+            out.end(token);
+        }
+
+        public static PendingResources<Action> createFromProto(ProtoInputStream in)
+                throws Exception {
+            final LongSparseArray<Object> values = new LongSparseArray<>();
+
+            values.put(RemoteViewsProto.TextViewDrawableAction.ICONS,
+                    new SparseArray<PendingResources<Icon>>());
+            values.put(RemoteViewsProto.TextViewDrawableAction.RESOURCES,
+                    new SparseArray<String>());
+            final long token = in.start(RemoteViewsProto.Action.TEXT_VIEW_DRAWABLE_ACTION);
+            while (in.nextField() != NO_MORE_FIELDS) {
+                switch (in.getFieldNumber()) {
+                    case (int) RemoteViewsProto.TextViewDrawableAction.VIEW_ID:
+                        values.put(RemoteViewsProto.TextViewDrawableAction.VIEW_ID,
+                                in.readString(RemoteViewsProto.TextViewDrawableAction.VIEW_ID));
+                        break;
+                    case (int) RemoteViewsProto.TextViewDrawableAction.IS_RELATIVE:
+                        values.put(RemoteViewsProto.TextViewDrawableAction.IS_RELATIVE,
+                                in.readBoolean(
+                                        RemoteViewsProto.TextViewDrawableAction.IS_RELATIVE));
+                        break;
+                    case (int) RemoteViewsProto.TextViewDrawableAction.RESOURCES:
+                        final long resourcesToken = in.start(
+                                RemoteViewsProto.TextViewDrawableAction.RESOURCES);
+                        while (in.nextField() != NO_MORE_FIELDS) {
+                            switch (in.getFieldNumber()) {
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Resources.ONE:
+                                    ((SparseArray<String>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.RESOURCES)).put(
+                                            1, in.readString(
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction.Resources.ONE));
+                                    break;
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Resources.TWO:
+                                    ((SparseArray<String>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.RESOURCES)).put(
+                                            2, in.readString(
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction.Resources.TWO));
+                                    break;
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Resources.THREE:
+                                    ((SparseArray<String>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.RESOURCES)).put(
+                                            3, in.readString(
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction
+                                                            .Resources.THREE));
+                                    break;
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Resources.FOUR:
+                                    ((SparseArray<String>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.RESOURCES)).put(
+                                            4, in.readString(
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction
+                                                            .Resources.FOUR));
+                                    break;
+                                default:
+                                    Log.w(LOG_TAG,
+                                            "Unhandled field while reading RemoteViews proto!\n"
+                                                    + ProtoUtils.currentFieldToString(in));
+                            }
+                        }
+                        in.end(resourcesToken);
+                        break;
+                    case (int) RemoteViewsProto.TextViewDrawableAction.ICONS:
+                        final long iconsToken = in.start(
+                                RemoteViewsProto.TextViewDrawableAction.ICONS);
+                        while (in.nextField() != NO_MORE_FIELDS) {
+                            switch (in.getFieldNumber()) {
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Icons.ONE:
+                                    ((SparseArray<PendingResources<Icon>>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.ICONS)).put(1,
+                                            createIconFromProto(in,
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction.Icons.ONE));
+                                    break;
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Icons.TWO:
+                                    ((SparseArray<PendingResources<Icon>>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.ICONS)).put(2,
+                                            createIconFromProto(in,
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction.Icons.TWO));
+                                    break;
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Icons.THREE:
+                                    ((SparseArray<PendingResources<Icon>>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.ICONS)).put(3,
+                                            createIconFromProto(in,
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction.Icons.THREE));
+                                    break;
+                                case (int) RemoteViewsProto.TextViewDrawableAction.Icons.FOUR:
+                                    ((SparseArray<PendingResources<Icon>>) values.get(
+                                            RemoteViewsProto.TextViewDrawableAction.ICONS)).put(4,
+                                            createIconFromProto(in,
+                                                    RemoteViewsProto
+                                                            .TextViewDrawableAction.Icons.FOUR));
+                                    break;
+                                default:
+                                    Log.w(LOG_TAG,
+                                            "Unhandled field while reading RemoteViews proto!\n"
+                                                    + ProtoUtils.currentFieldToString(in));
+                            }
+                        }
+                        in.end(iconsToken);
+                        break;
+                    default:
+                        Log.w(LOG_TAG, "Unhandled field while reading RemoteViews proto!\n"
+                                + ProtoUtils.currentFieldToString(in));
+                }
+            }
+            in.end(token);
+
+            checkContainsKeys(values, new long[]{RemoteViewsProto.TextViewDrawableAction.VIEW_ID});
+
+            return (context, resources, rootData, depth) -> {
+                int viewId = getAsIdentifier(resources, values,
+                        RemoteViewsProto.TextViewDrawableAction.VIEW_ID);
+                SparseArray<PendingResources<Icon>> icons =
+                        (SparseArray<PendingResources<Icon>>) values.get(
+                                RemoteViewsProto.TextViewDrawableAction.ICONS);
+                SparseArray<String> resArray = (SparseArray<String>) values.get(
+                        RemoteViewsProto.TextViewDrawableAction.RESOURCES);
+                boolean isRelative = (boolean) values.get(
+                        RemoteViewsProto.TextViewDrawableAction.IS_RELATIVE, false);
+                if (icons.size() > 0) {
+                    return new TextViewDrawableAction(viewId, isRelative,
+                            icons.get(1).create(context, resources, rootData, depth),
+                            icons.get(2).create(context, resources, rootData, depth),
+                            icons.get(3).create(context, resources, rootData, depth),
+                            icons.get(4).create(context, resources, rootData, depth));
+                } else {
+                    int first = resArray.contains(1) ? getAsIdentifier(resources, resArray, 1) : 0;
+                    int second = resArray.contains(2) ? getAsIdentifier(resources, resArray, 2) : 0;
+                    int third = resArray.contains(3) ? getAsIdentifier(resources, resArray, 3) : 0;
+                    int fourth = resArray.contains(4) ? getAsIdentifier(resources, resArray, 4) : 0;
+                    return new TextViewDrawableAction(viewId, isRelative, first, second, third,
+                            fourth);
+                }
+            };
+        }
     }
 
     /**
@@ -4535,6 +4849,58 @@
         public int getActionTag() {
             return TEXT_VIEW_SIZE_ACTION_TAG;
         }
+
+        @Override
+        public boolean canWriteToProto() {
+            return true;
+        }
+
+        @Override
+        public void writeToProto(ProtoOutputStream out, Context context, Resources appResources) {
+            final long token = out.start(RemoteViewsProto.Action.TEXT_VIEW_SIZE_ACTION);
+            out.write(RemoteViewsProto.TextViewSizeAction.VIEW_ID,
+                    appResources.getResourceName(mViewId));
+            out.write(RemoteViewsProto.TextViewSizeAction.UNITS, mUnits);
+            out.write(RemoteViewsProto.TextViewSizeAction.SIZE, mSize);
+            out.end(token);
+        }
+
+        public static PendingResources<Action> createFromProto(ProtoInputStream in)
+                throws Exception {
+            final LongSparseArray<Object> values = new LongSparseArray<>();
+
+            final long token = in.start(RemoteViewsProto.Action.TEXT_VIEW_SIZE_ACTION);
+            while (in.nextField() != NO_MORE_FIELDS) {
+                switch (in.getFieldNumber()) {
+                    case (int) RemoteViewsProto.TextViewSizeAction.VIEW_ID:
+                        values.put(RemoteViewsProto.TextViewSizeAction.VIEW_ID,
+                                in.readString(RemoteViewsProto.TextViewSizeAction.VIEW_ID));
+                        break;
+                    case (int) RemoteViewsProto.TextViewSizeAction.UNITS:
+                        values.put(RemoteViewsProto.TextViewSizeAction.UNITS,
+                                in.readInt(RemoteViewsProto.TextViewSizeAction.UNITS));
+                        break;
+                    case (int) RemoteViewsProto.TextViewSizeAction.SIZE:
+                        values.put(RemoteViewsProto.TextViewSizeAction.SIZE,
+                                in.readFloat(RemoteViewsProto.TextViewSizeAction.SIZE));
+                        break;
+                    default:
+                        Log.w(LOG_TAG, "Unhandled field while reading RemoteViews proto!\n"
+                                + ProtoUtils.currentFieldToString(in));
+                }
+            }
+            in.end(token);
+
+            checkContainsKeys(values, new long[]{RemoteViewsProto.TextViewSizeAction.VIEW_ID});
+
+            return (context, resources, rootData, depth) -> {
+                int viewId = getAsIdentifier(resources, values,
+                        RemoteViewsProto.TextViewSizeAction.VIEW_ID);
+                return new TextViewSizeAction(viewId,
+                        (int) values.get(RemoteViewsProto.TextViewSizeAction.UNITS, 0),
+                        (float) values.get(RemoteViewsProto.TextViewSizeAction.SIZE, 0));
+            };
+        }
     }
 
     /**
@@ -4579,6 +4945,70 @@
         public int getActionTag() {
             return VIEW_PADDING_ACTION_TAG;
         }
+
+        @Override
+        public boolean canWriteToProto() {
+            return true;
+        }
+
+        @Override
+        public void writeToProto(ProtoOutputStream out, Context context, Resources appResources) {
+            final long token = out.start(RemoteViewsProto.Action.VIEW_PADDING_ACTION);
+            out.write(RemoteViewsProto.ViewPaddingAction.VIEW_ID,
+                    appResources.getResourceName(mViewId));
+            out.write(RemoteViewsProto.ViewPaddingAction.LEFT, mLeft);
+            out.write(RemoteViewsProto.ViewPaddingAction.RIGHT, mRight);
+            out.write(RemoteViewsProto.ViewPaddingAction.TOP, mTop);
+            out.write(RemoteViewsProto.ViewPaddingAction.BOTTOM, mBottom);
+            out.end(token);
+        }
+
+        public static PendingResources<Action> createFromProto(ProtoInputStream in)
+                throws Exception {
+            final LongSparseArray<Object> values = new LongSparseArray<>();
+
+            final long token = in.start(RemoteViewsProto.Action.VIEW_PADDING_ACTION);
+            while (in.nextField() != NO_MORE_FIELDS) {
+                switch (in.getFieldNumber()) {
+                    case (int) RemoteViewsProto.ViewPaddingAction.VIEW_ID:
+                        values.put(RemoteViewsProto.ViewPaddingAction.VIEW_ID,
+                                in.readString(RemoteViewsProto.ViewPaddingAction.VIEW_ID));
+                        break;
+                    case (int) RemoteViewsProto.ViewPaddingAction.LEFT:
+                        values.put(RemoteViewsProto.ViewPaddingAction.LEFT,
+                                in.readInt(RemoteViewsProto.ViewPaddingAction.LEFT));
+                        break;
+                    case (int) RemoteViewsProto.ViewPaddingAction.RIGHT:
+                        values.put(RemoteViewsProto.ViewPaddingAction.RIGHT,
+                                in.readInt(RemoteViewsProto.ViewPaddingAction.RIGHT));
+                        break;
+                    case (int) RemoteViewsProto.ViewPaddingAction.TOP:
+                        values.put(RemoteViewsProto.ViewPaddingAction.TOP,
+                                in.readInt(RemoteViewsProto.ViewPaddingAction.TOP));
+                        break;
+                    case (int) RemoteViewsProto.ViewPaddingAction.BOTTOM:
+                        values.put(RemoteViewsProto.ViewPaddingAction.BOTTOM,
+                                in.readInt(RemoteViewsProto.ViewPaddingAction.BOTTOM));
+                        break;
+                    default:
+                        Log.w(LOG_TAG, "Unhandled field while reading RemoteViews proto!\n"
+                                + ProtoUtils.currentFieldToString(in));
+                }
+            }
+            in.end(token);
+
+            checkContainsKeys(values, new long[]{RemoteViewsProto.ViewPaddingAction.VIEW_ID});
+
+            return (context, resources, rootData, depth) -> {
+                int viewId = getAsIdentifier(resources, values,
+                        RemoteViewsProto.ViewPaddingAction.VIEW_ID);
+                return new ViewPaddingAction(viewId,
+                        (int) values.get(RemoteViewsProto.ViewPaddingAction.LEFT, 0),
+                        (int) values.get(RemoteViewsProto.ViewPaddingAction.TOP, 0),
+                        (int) values.get(RemoteViewsProto.ViewPaddingAction.RIGHT, 0),
+                        (int) values.get(RemoteViewsProto.ViewPaddingAction.BOTTOM, 0));
+            };
+        }
     }
 
     /**
@@ -5241,6 +5671,69 @@
         public int getActionTag() {
             return SET_VIEW_OUTLINE_RADIUS_TAG;
         }
+
+        @Override
+        public boolean canWriteToProto() {
+            return true;
+        }
+
+        @Override
+        public void writeToProto(ProtoOutputStream out, Context context, Resources appResources) {
+            final long token = out.start(
+                    RemoteViewsProto.Action.SET_VIEW_OUTLINE_PREFERRED_RADIUS_ACTION);
+            out.write(RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VIEW_ID,
+                    appResources.getResourceName(mViewId));
+            out.write(RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE_TYPE, mValueType);
+            out.write(RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE, mValue);
+            out.end(token);
+        }
+
+        public static PendingResources<Action> createFromProto(ProtoInputStream in)
+                throws Exception {
+            final LongSparseArray<Object> values = new LongSparseArray<>();
+
+            final long token = in.start(
+                    RemoteViewsProto.Action.SET_VIEW_OUTLINE_PREFERRED_RADIUS_ACTION);
+            while (in.nextField() != NO_MORE_FIELDS) {
+                switch (in.getFieldNumber()) {
+                    case (int) RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VIEW_ID:
+                        values.put(RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VIEW_ID,
+                                in.readString(
+                                        RemoteViewsProto
+                                                .SetViewOutlinePreferredRadiusAction.VIEW_ID));
+                        break;
+                    case (int) RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE_TYPE:
+                        values.put(RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE_TYPE,
+                                in.readInt(
+                                        RemoteViewsProto
+                                                .SetViewOutlinePreferredRadiusAction.VALUE_TYPE));
+                        break;
+                    case (int) RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE:
+                        values.put(RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE,
+                                in.readInt(
+                                        RemoteViewsProto
+                                                .SetViewOutlinePreferredRadiusAction.VALUE));
+                        break;
+                    default:
+                        Log.w(LOG_TAG, "Unhandled field while reading RemoteViews proto!\n"
+                                + ProtoUtils.currentFieldToString(in));
+                }
+            }
+            in.end(token);
+
+            checkContainsKeys(values,
+                    new long[]{RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VIEW_ID,
+                            RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE_TYPE});
+
+            return (context, resources, rootData, depth) -> {
+                int viewId = getAsIdentifier(resources, values,
+                        RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VIEW_ID);
+                return new SetViewOutlinePreferredRadiusAction(viewId,
+                        (int) values.get(RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE,
+                                0), (int) values.get(
+                        RemoteViewsProto.SetViewOutlinePreferredRadiusAction.VALUE_TYPE));
+            };
+        }
     }
 
     /**
@@ -5324,6 +5817,46 @@
         public int getActionTag() {
             return SET_DRAW_INSTRUCTION_TAG;
         }
+
+        @Override
+        public boolean canWriteToProto() {
+            return drawDataParcel();
+        }
+
+        @Override
+        public void writeToProto(ProtoOutputStream out, Context context, Resources appResources) {
+            if (!drawDataParcel()) return;
+            final long token = out.start(RemoteViewsProto.Action.SET_DRAW_INSTRUCTION_ACTION);
+            if (mInstructions != null) {
+                for (byte[] bytes : mInstructions.mInstructions) {
+                    out.write(RemoteViewsProto.SetDrawInstructionAction.INSTRUCTIONS, bytes);
+                }
+            }
+            out.end(token);
+        }
+    }
+
+    @FlaggedApi(FLAG_DRAW_DATA_PARCEL)
+    private PendingResources<Action> createSetDrawInstructionActionFromProto(ProtoInputStream in)
+            throws Exception {
+        List<byte[]> instructions = new ArrayList<byte[]>();
+
+        final long token = in.start(RemoteViewsProto.Action.SET_DRAW_INSTRUCTION_ACTION);
+        while (in.nextField() != NO_MORE_FIELDS) {
+            switch (in.getFieldNumber()) {
+                case (int) RemoteViewsProto.SetDrawInstructionAction.INSTRUCTIONS:
+                    instructions.add(
+                            in.readBytes(RemoteViewsProto.SetDrawInstructionAction.INSTRUCTIONS));
+                    break;
+                default:
+                    Log.w(LOG_TAG, "Unhandled field while reading RemoteViews proto!\n"
+                            + ProtoUtils.currentFieldToString(in));
+            }
+        }
+        in.end(token);
+
+        return (context, resources, rootData, depth) -> new SetDrawInstructionAction(
+                new DrawInstructions.Builder(instructions).build());
     }
 
     /**
@@ -9604,12 +10137,20 @@
             }
             if (ref.mMode == MODE_NORMAL) {
                 rv.setIdealSize(ref.mIdealSize);
+                boolean hasDrawInstructionAction = false;
                 for (PendingResources<Action> pendingAction : ref.mActions) {
                     Action action = pendingAction.create(appContext, appResources, rootData, depth);
                     if (action != null) {
+                        if (action instanceof SetDrawInstructionAction) {
+                            hasDrawInstructionAction = true;
+                        }
                         rv.addAction(action);
                     }
                 }
+                if (rv.mHasDrawInstructions && !hasDrawInstructionAction) {
+                    throw new InvalidProtoException(
+                            "RemoteViews proto is missing DrawInstructions");
+                }
                 return rv;
             } else if (ref.mMode == MODE_HAS_SIZED_REMOTEVIEWS) {
                 List<RemoteViews> sizedViews = new ArrayList<>();
@@ -9685,6 +10226,23 @@
                 return rv.createSetRemoteCollectionItemListAdapterActionFromProto(in);
             case (int) RemoteViewsProto.Action.SET_RIPPLE_DRAWABLE_COLOR_ACTION:
                 return SetRippleDrawableColor.createFromProto(in);
+            case (int) RemoteViewsProto.Action.SET_VIEW_OUTLINE_PREFERRED_RADIUS_ACTION:
+                return SetViewOutlinePreferredRadiusAction.createFromProto(in);
+            case (int) RemoteViewsProto.Action.TEXT_VIEW_DRAWABLE_ACTION:
+                return TextViewDrawableAction.createFromProto(in);
+            case (int) RemoteViewsProto.Action.TEXT_VIEW_SIZE_ACTION:
+                return TextViewSizeAction.createFromProto(in);
+            case (int) RemoteViewsProto.Action.VIEW_GROUP_ADD_ACTION:
+                return rv.createViewGroupActionAddFromProto(in);
+            case (int) RemoteViewsProto.Action.VIEW_GROUP_REMOVE_ACTION:
+                return ViewGroupActionRemove.createFromProto(in);
+            case (int) RemoteViewsProto.Action.VIEW_PADDING_ACTION:
+                return ViewPaddingAction.createFromProto(in);
+            case (int) RemoteViewsProto.Action.SET_DRAW_INSTRUCTION_ACTION:
+                if (!drawDataParcel()) {
+                    return null;
+                }
+                return rv.createSetDrawInstructionActionFromProto(in);
             default:
                 throw new RuntimeException("Unhandled field while reading Action proto!\n"
                         + ProtoUtils.currentFieldToString(in));
diff --git a/core/java/android/window/flags/DesktopModeFlags.java b/core/java/android/window/DesktopModeFlags.java
similarity index 99%
rename from core/java/android/window/flags/DesktopModeFlags.java
rename to core/java/android/window/DesktopModeFlags.java
index 47af50da..6fb82af 100644
--- a/core/java/android/window/flags/DesktopModeFlags.java
+++ b/core/java/android/window/DesktopModeFlags.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package android.window.flags;
+package android.window;
 
 import android.annotation.Nullable;
 import android.app.ActivityThread;
diff --git a/core/java/android/window/KeyguardState.aidl b/core/java/android/window/KeyguardState.aidl
new file mode 100644
index 0000000..9612d95
--- /dev/null
+++ b/core/java/android/window/KeyguardState.aidl
@@ -0,0 +1,22 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.window;
+
+/**
+ * @hide
+ */
+parcelable KeyguardState;
diff --git a/core/java/android/window/KeyguardState.java b/core/java/android/window/KeyguardState.java
new file mode 100644
index 0000000..6584d30
--- /dev/null
+++ b/core/java/android/window/KeyguardState.java
@@ -0,0 +1,160 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.window;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+import java.util.Objects;
+
+/**
+ * Data object of params for Keyguard related {@link WindowContainerTransaction} operation.
+ *
+ * @hide
+ */
+public final class KeyguardState implements Parcelable {
+
+    private final int mDisplayId;
+
+    private final boolean mKeyguardShowing;
+
+    private final boolean mAodShowing;
+
+
+    private KeyguardState(int displayId, boolean keyguardShowing, boolean aodShowing) {
+        mDisplayId = displayId;
+        mKeyguardShowing = keyguardShowing;
+        mAodShowing = aodShowing;
+    }
+
+    private KeyguardState(Parcel in) {
+        mDisplayId = in.readInt();
+        mKeyguardShowing = in.readBoolean();
+        mAodShowing = in.readBoolean();
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        dest.writeInt(mDisplayId);
+        dest.writeBoolean(mKeyguardShowing);
+        dest.writeBoolean(mAodShowing);
+    }
+
+    @NonNull
+    public static final Creator<KeyguardState> CREATOR =
+            new Creator<KeyguardState>() {
+                @Override
+                public KeyguardState createFromParcel(Parcel in) {
+                    return new KeyguardState(in);
+                }
+
+                @Override
+                public KeyguardState[] newArray(int size) {
+                    return new KeyguardState[size];
+                }
+            };
+
+    /**
+     * Gets the display id of this {@link KeyguardState}.
+     */
+    public int getDisplayId() {
+        return mDisplayId;
+    }
+
+    /** Returns the keyguard showing value. */
+    public boolean getKeyguardShowing() {
+        return mKeyguardShowing;
+    }
+
+    /** Returns the aod showing value. */
+    public boolean getAodShowing() {
+        return mAodShowing;
+    }
+
+    @Override
+    public String toString() {
+        return "KeyguardState{ displayId=" + mDisplayId
+                + ", keyguardShowing=" + mKeyguardShowing
+                + ", aodShowing=" + mAodShowing
+                + '}';
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(mDisplayId, mKeyguardShowing, mAodShowing);
+    }
+
+    @Override
+    public boolean equals(@Nullable Object obj) {
+        if (!(obj instanceof KeyguardState other)) {
+            return false;
+        }
+        return mDisplayId == other.mDisplayId
+                && mKeyguardShowing == other.mKeyguardShowing
+                && mAodShowing == other.mAodShowing;
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    /** Builder to construct the {@link KeyguardState}. */
+    public static final class Builder {
+
+        private final int mDisplayId;
+
+        private boolean mKeyguardShowing;
+
+        private boolean mAodShowing;
+
+        /**
+         * @param displayId the display of this {@link KeyguardState}.
+         */
+        public Builder(int displayId) {
+            mDisplayId = displayId;
+        }
+
+        /**
+         * Sets the boolean value for this operation.
+         */
+        @NonNull
+        public Builder setKeyguardShowing(boolean keyguardShowing) {
+            mKeyguardShowing = keyguardShowing;
+            return this;
+        }
+
+        /**
+         * Sets the boolean value for this operation.
+         */
+        @NonNull
+        public Builder setAodShowing(boolean aodShowing) {
+            mAodShowing = aodShowing;
+            return this;
+        }
+
+        /**
+         * Constructs the {@link KeyguardState}.
+         */
+        @NonNull
+        public KeyguardState build() {
+            return new KeyguardState(mDisplayId, mKeyguardShowing, mAodShowing);
+        }
+    }
+}
diff --git a/core/java/android/window/WindowContainerTransaction.java b/core/java/android/window/WindowContainerTransaction.java
index 0dc9263..8e495ec 100644
--- a/core/java/android/window/WindowContainerTransaction.java
+++ b/core/java/android/window/WindowContainerTransaction.java
@@ -869,6 +869,24 @@
     }
 
     /**
+     * Adds a {@link KeyguardState} to apply to the given displays.
+     *
+     * @hide
+     */
+    @NonNull
+    public WindowContainerTransaction addKeyguardState(
+            @NonNull KeyguardState keyguardState) {
+        Objects.requireNonNull(keyguardState);
+        final HierarchyOp hierarchyOp =
+                new HierarchyOp.Builder(
+                        HierarchyOp.HIERARCHY_OP_TYPE_SET_KEYGUARD_STATE)
+                        .setKeyguardState(keyguardState)
+                        .build();
+        mHierarchyOps.add(hierarchyOp);
+        return this;
+    }
+
+    /**
      * Sets/removes the always on top flag for this {@code windowContainer}. See
      * {@link com.android.server.wm.ConfigurationContainer#setAlwaysOnTop(boolean)}.
      * Please note that this method is only intended to be used for a
@@ -1469,6 +1487,7 @@
         public static final int HIERARCHY_OP_TYPE_SET_IS_TRIMMABLE = 19;
         public static final int HIERARCHY_OP_TYPE_RESTORE_BACK_NAVIGATION = 20;
         public static final int HIERARCHY_OP_TYPE_SET_EXCLUDE_INSETS_TYPES = 21;
+        public static final int HIERARCHY_OP_TYPE_SET_KEYGUARD_STATE = 22;
 
         // The following key(s) are for use with mLaunchOptions:
         // When launching a task (eg. from recents), this is the taskId to be launched.
@@ -1516,6 +1535,9 @@
         private TaskFragmentOperation mTaskFragmentOperation;
 
         @Nullable
+        private KeyguardState mKeyguardState;
+
+        @Nullable
         private PendingIntent mPendingIntent;
 
         @Nullable
@@ -1666,6 +1688,7 @@
             mLaunchOptions = copy.mLaunchOptions;
             mActivityIntent = copy.mActivityIntent;
             mTaskFragmentOperation = copy.mTaskFragmentOperation;
+            mKeyguardState = copy.mKeyguardState;
             mPendingIntent = copy.mPendingIntent;
             mShortcutInfo = copy.mShortcutInfo;
             mAlwaysOnTop = copy.mAlwaysOnTop;
@@ -1689,6 +1712,7 @@
             mLaunchOptions = in.readBundle();
             mActivityIntent = in.readTypedObject(Intent.CREATOR);
             mTaskFragmentOperation = in.readTypedObject(TaskFragmentOperation.CREATOR);
+            mKeyguardState = in.readTypedObject(KeyguardState.CREATOR);
             mPendingIntent = in.readTypedObject(PendingIntent.CREATOR);
             mShortcutInfo = in.readTypedObject(ShortcutInfo.CREATOR);
             mAlwaysOnTop = in.readBoolean();
@@ -1770,6 +1794,11 @@
         }
 
         @Nullable
+        public KeyguardState getKeyguardState() {
+            return mKeyguardState;
+        }
+
+        @Nullable
         public PendingIntent getPendingIntent() {
             return mPendingIntent;
         }
@@ -1902,6 +1931,9 @@
                             .append(" mExcludeInsetsTypes= ")
                             .append(WindowInsets.Type.toString(mExcludeInsetsTypes));
                     break;
+                case HIERARCHY_OP_TYPE_SET_KEYGUARD_STATE:
+                    sb.append("KeyguardState= ").append(mKeyguardState);
+                    break;
                 case HIERARCHY_OP_TYPE_SET_IS_TRIMMABLE:
                     sb.append("container= ").append(mContainer)
                             .append(" isTrimmable= ")
@@ -1932,6 +1964,7 @@
             dest.writeBundle(mLaunchOptions);
             dest.writeTypedObject(mActivityIntent, flags);
             dest.writeTypedObject(mTaskFragmentOperation, flags);
+            dest.writeTypedObject(mKeyguardState, flags);
             dest.writeTypedObject(mPendingIntent, flags);
             dest.writeTypedObject(mShortcutInfo, flags);
             dest.writeBoolean(mAlwaysOnTop);
@@ -1993,6 +2026,9 @@
             private TaskFragmentOperation mTaskFragmentOperation;
 
             @Nullable
+            private KeyguardState mKeyguardState;
+
+            @Nullable
             private PendingIntent mPendingIntent;
 
             @Nullable
@@ -2081,6 +2117,12 @@
                 return this;
             }
 
+            Builder setKeyguardState(
+                    @Nullable KeyguardState keyguardState) {
+                mKeyguardState = keyguardState;
+                return this;
+            }
+
             Builder setReparentLeafTaskIfRelaunch(boolean reparentLeafTaskIfRelaunch) {
                 mReparentLeafTaskIfRelaunch = reparentLeafTaskIfRelaunch;
                 return this;
@@ -2130,6 +2172,7 @@
                 hierarchyOp.mPendingIntent = mPendingIntent;
                 hierarchyOp.mAlwaysOnTop = mAlwaysOnTop;
                 hierarchyOp.mTaskFragmentOperation = mTaskFragmentOperation;
+                hierarchyOp.mKeyguardState = mKeyguardState;
                 hierarchyOp.mShortcutInfo = mShortcutInfo;
                 hierarchyOp.mBounds = mBounds;
                 hierarchyOp.mIncludingParents = mIncludingParents;
diff --git a/core/java/android/window/flags/lse_desktop_experience.aconfig b/core/java/android/window/flags/lse_desktop_experience.aconfig
index b22aa22..31bb3a6 100644
--- a/core/java/android/window/flags/lse_desktop_experience.aconfig
+++ b/core/java/android/window/flags/lse_desktop_experience.aconfig
@@ -319,4 +319,11 @@
     namespace: "lse_desktop_experience"
     description: "Enables custom transitions for alt-tab app launches in Desktop Mode."
     bug: "370735595"
-}
\ No newline at end of file
+}
+
+flag {
+    name: "enable_move_to_next_display_shortcut"
+    namespace: "lse_desktop_experience"
+    description: "Add new keyboard shortcut of moving a task into next display"
+    bug: "364154795"
+}
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index c9b93c9..622f8c8 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -173,6 +173,16 @@
 }
 
 flag {
+  name: "filter_irrelevant_input_device_change"
+  namespace: "windowing_frontend"
+  description: "Recompute display configuration only for necessary input device changes"
+  bug: "368461853"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
+
+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/java/com/android/internal/accessibility/dialog/AccessibilityServiceWarning.java b/core/java/com/android/internal/accessibility/dialog/AccessibilityServiceWarning.java
index 0f8ced2..3557633 100644
--- a/core/java/com/android/internal/accessibility/dialog/AccessibilityServiceWarning.java
+++ b/core/java/com/android/internal/accessibility/dialog/AccessibilityServiceWarning.java
@@ -29,6 +29,7 @@
 import android.view.View;
 import android.view.Window;
 import android.view.WindowManager;
+import android.view.accessibility.Flags;
 import android.widget.Button;
 import android.widget.ImageView;
 import android.widget.TextView;
@@ -58,13 +59,15 @@
             @NonNull View.OnClickListener uninstallListener) {
         final AlertDialog ad = new AlertDialog.Builder(context)
                 .setView(createAccessibilityServiceWarningDialogContentView(
-                                context, info, allowListener, denyListener, uninstallListener))
+                        context, info, allowListener, denyListener, uninstallListener))
                 .setCancelable(true)
                 .create();
         Window window = ad.getWindow();
         WindowManager.LayoutParams params = window.getAttributes();
         params.privateFlags |= SYSTEM_FLAG_HIDE_NON_SYSTEM_OVERLAY_WINDOWS;
-        params.type = WindowManager.LayoutParams.TYPE_SYSTEM_DIALOG;
+        if (!Flags.warningUseDefaultDialogType()) {
+            params.type = WindowManager.LayoutParams.TYPE_SYSTEM_DIALOG;
+        }
         window.setAttributes(params);
         return ad;
     }
diff --git a/core/java/com/android/internal/policy/DecorView.java b/core/java/com/android/internal/policy/DecorView.java
index 6faea17..bd746d5 100644
--- a/core/java/com/android/internal/policy/DecorView.java
+++ b/core/java/com/android/internal/policy/DecorView.java
@@ -38,8 +38,8 @@
 import static android.view.WindowManager.LayoutParams.FLAG_TRANSLUCENT_STATUS;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 import static android.view.flags.Flags.customizableWindowHeaders;
-import static android.window.flags.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
-import static android.window.flags.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION_ALWAYS;
+import static android.window.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
+import static android.window.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION_ALWAYS;
 
 import static com.android.internal.policy.PhoneWindow.FEATURE_OPTIONS_PANEL;
 
diff --git a/core/java/com/android/internal/widget/PointerLocationView.java b/core/java/com/android/internal/widget/PointerLocationView.java
index c0a7383..e65b4b6 100644
--- a/core/java/com/android/internal/widget/PointerLocationView.java
+++ b/core/java/com/android/internal/widget/PointerLocationView.java
@@ -16,19 +16,14 @@
 
 package com.android.internal.widget;
 
-import static java.lang.Float.NaN;
-
 import android.compat.annotation.UnsupportedAppUsage;
 import android.content.Context;
 import android.content.res.Configuration;
-import android.graphics.Bitmap;
 import android.graphics.Canvas;
-import android.graphics.Color;
 import android.graphics.Insets;
 import android.graphics.Paint;
 import android.graphics.Paint.FontMetricsInt;
 import android.graphics.Path;
-import android.graphics.PorterDuff;
 import android.graphics.RectF;
 import android.graphics.Region;
 import android.hardware.input.InputManager;
@@ -70,14 +65,11 @@
     private static final PointerState EMPTY_POINTER_STATE = new PointerState();
 
     public static class PointerState {
-        private float mCurrentX = NaN;
-        private float mCurrentY = NaN;
-        private float mPreviousX = NaN;
-        private float mPreviousY = NaN;
-        private float mFirstX = NaN;
-        private float mFirstY = NaN;
-        private boolean mPreviousPointIsHistorical;
-        private boolean mCurrentPointIsHistorical;
+        // Trace of previous points.
+        private float[] mTraceX = new float[32];
+        private float[] mTraceY = new float[32];
+        private boolean[] mTraceCurrent = new boolean[32];
+        private int mTraceCount;
 
         // True if the pointer is down.
         @UnsupportedAppUsage
@@ -104,20 +96,31 @@
         public PointerState() {
         }
 
-        public void addTrace(float x, float y, boolean isHistorical) {
-            if (Float.isNaN(mFirstX)) {
-                mFirstX = x;
-            }
-            if (Float.isNaN(mFirstY)) {
-                mFirstY = y;
+        public void clearTrace() {
+            mTraceCount = 0;
+        }
+
+        public void addTrace(float x, float y, boolean current) {
+            int traceCapacity = mTraceX.length;
+            if (mTraceCount == traceCapacity) {
+                traceCapacity *= 2;
+                float[] newTraceX = new float[traceCapacity];
+                System.arraycopy(mTraceX, 0, newTraceX, 0, mTraceCount);
+                mTraceX = newTraceX;
+
+                float[] newTraceY = new float[traceCapacity];
+                System.arraycopy(mTraceY, 0, newTraceY, 0, mTraceCount);
+                mTraceY = newTraceY;
+
+                boolean[] newTraceCurrent = new boolean[traceCapacity];
+                System.arraycopy(mTraceCurrent, 0, newTraceCurrent, 0, mTraceCount);
+                mTraceCurrent= newTraceCurrent;
             }
 
-            mPreviousX = mCurrentX;
-            mPreviousY = mCurrentY;
-            mCurrentX = x;
-            mCurrentY = y;
-            mPreviousPointIsHistorical = mCurrentPointIsHistorical;
-            mCurrentPointIsHistorical = isHistorical;
+            mTraceX[mTraceCount] = x;
+            mTraceY[mTraceCount] = y;
+            mTraceCurrent[mTraceCount] = current;
+            mTraceCount += 1;
         }
     }
 
@@ -146,12 +149,6 @@
     private final SparseArray<PointerState> mPointers = new SparseArray<PointerState>();
     private final PointerCoords mTempCoords = new PointerCoords();
 
-    // Draw the trace of all pointers in the current gesture in a separate layer
-    // that is not cleared on every frame so that we don't have to re-draw the
-    // entire trace on each frame.
-    private final Bitmap mTraceBitmap;
-    private final Canvas mTraceCanvas;
-
     private final Region mSystemGestureExclusion = new Region();
     private final Region mSystemGestureExclusionRejected = new Region();
     private final Path mSystemGestureExclusionPath = new Path();
@@ -200,10 +197,6 @@
         mPathPaint.setARGB(255, 0, 96, 255);
         mPathPaint.setStyle(Paint.Style.STROKE);
 
-        mTraceBitmap = Bitmap.createBitmap(getResources().getDisplayMetrics().widthPixels,
-                getResources().getDisplayMetrics().heightPixels, Bitmap.Config.ARGB_8888);
-        mTraceCanvas = new Canvas(mTraceBitmap);
-
         configureDensityDependentFactors();
 
         mSystemGestureExclusionPaint = new Paint();
@@ -263,7 +256,7 @@
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
         super.onMeasure(widthMeasureSpec, heightMeasureSpec);
         mTextPaint.getFontMetricsInt(mTextMetrics);
-        mHeaderBottom = mHeaderPaddingTop - mTextMetrics.ascent + mTextMetrics.descent + 2;
+        mHeaderBottom = mHeaderPaddingTop-mTextMetrics.ascent+mTextMetrics.descent+2;
         if (false) {
             Log.i("foo", "Metrics: ascent=" + mTextMetrics.ascent
                     + " descent=" + mTextMetrics.descent
@@ -276,7 +269,6 @@
     // Draw an oval.  When angle is 0 radians, orients the major axis vertically,
     // angles less than or greater than 0 radians rotate the major axis left or right.
     private RectF mReusableOvalRect = new RectF();
-
     private void drawOval(Canvas canvas, float x, float y, float major, float minor,
             float angle, Paint paint) {
         canvas.save(Canvas.MATRIX_SAVE_FLAG);
@@ -293,8 +285,6 @@
     protected void onDraw(Canvas canvas) {
         final int NP = mPointers.size();
 
-        canvas.drawBitmap(mTraceBitmap, 0, 0, null);
-
         if (!mSystemGestureExclusion.isEmpty()) {
             mSystemGestureExclusionPath.reset();
             mSystemGestureExclusion.getBoundaryPath(mSystemGestureExclusionPath);
@@ -313,9 +303,32 @@
         // Pointer trace.
         for (int p = 0; p < NP; p++) {
             final PointerState ps = mPointers.valueAt(p);
-            float lastX = ps.mCurrentX, lastY = ps.mCurrentY;
 
-            if (!Float.isNaN(lastX) && !Float.isNaN(lastY)) {
+            // Draw path.
+            final int N = ps.mTraceCount;
+            float lastX = 0, lastY = 0;
+            boolean haveLast = false;
+            boolean drawn = false;
+            mPaint.setARGB(255, 128, 255, 255);
+            for (int i=0; i < N; i++) {
+                float x = ps.mTraceX[i];
+                float y = ps.mTraceY[i];
+                if (Float.isNaN(x) || Float.isNaN(y)) {
+                    haveLast = false;
+                    continue;
+                }
+                if (haveLast) {
+                    canvas.drawLine(lastX, lastY, x, y, mPathPaint);
+                    final Paint paint = ps.mTraceCurrent[i - 1] ? mCurrentPointPaint : mPaint;
+                    canvas.drawPoint(lastX, lastY, paint);
+                    drawn = true;
+                }
+                lastX = x;
+                lastY = y;
+                haveLast = true;
+            }
+
+            if (drawn) {
                 // Draw velocity vector.
                 mPaint.setARGB(255, 255, 64, 128);
                 float xVel = ps.mXVelocity * (1000 / 60);
@@ -340,7 +353,7 @@
                         Math.max(getHeight(), getWidth()), mTargetPaint);
 
                 // Draw current point.
-                int pressureLevel = (int) (ps.mCoords.pressure * 255);
+                int pressureLevel = (int)(ps.mCoords.pressure * 255);
                 mPaint.setARGB(255, pressureLevel, 255, 255 - pressureLevel);
                 canvas.drawPoint(ps.mCoords.x, ps.mCoords.y, mPaint);
 
@@ -411,7 +424,8 @@
                 .append(" / ").append(mMaxNumPointers)
                 .toString(), 1, base, mTextPaint);
 
-        if ((mCurDown && ps.mCurDown) || Float.isNaN(ps.mCurrentX)) {
+        final int count = ps.mTraceCount;
+        if ((mCurDown && ps.mCurDown) || count == 0) {
             canvas.drawRect(itemW, mHeaderPaddingTop, (itemW * 2) - 1, bottom,
                     mTextBackgroundPaint);
             canvas.drawText(mText.clear()
@@ -423,8 +437,8 @@
                     .append("Y: ").append(ps.mCoords.y, 1)
                     .toString(), 1 + itemW * 2, base, mTextPaint);
         } else {
-            float dx = ps.mCurrentX - ps.mFirstX;
-            float dy = ps.mCurrentY - ps.mFirstY;
+            float dx = ps.mTraceX[count - 1] - ps.mTraceX[0];
+            float dy = ps.mTraceY[count - 1] - ps.mTraceY[0];
             canvas.drawRect(itemW, mHeaderPaddingTop, (itemW * 2) - 1, bottom,
                     Math.abs(dx) < mVC.getScaledTouchSlop()
                             ? mTextBackgroundPaint : mTextLevelPaint);
@@ -551,9 +565,9 @@
                 .append(" TouchMinor=").append(coords.touchMinor, 3)
                 .append(" ToolMajor=").append(coords.toolMajor, 3)
                 .append(" ToolMinor=").append(coords.toolMinor, 3)
-                .append(" Orientation=").append((float) (coords.orientation * 180 / Math.PI), 1)
+                .append(" Orientation=").append((float)(coords.orientation * 180 / Math.PI), 1)
                 .append("deg")
-                .append(" Tilt=").append((float) (
+                .append(" Tilt=").append((float)(
                         coords.getAxisValue(MotionEvent.AXIS_TILT) * 180 / Math.PI), 1)
                 .append("deg")
                 .append(" Distance=").append(coords.getAxisValue(MotionEvent.AXIS_DISTANCE), 1)
@@ -584,7 +598,6 @@
                 mCurNumPointers = 0;
                 mMaxNumPointers = 0;
                 mVelocity.clear();
-                mTraceCanvas.drawColor(Color.TRANSPARENT, PorterDuff.Mode.CLEAR);
                 if (mAltVelocity != null) {
                     mAltVelocity.clear();
                 }
@@ -633,8 +646,7 @@
                     logCoords("Pointer", action, i, coords, id, event);
                 }
                 if (ps != null) {
-                    ps.addTrace(coords.x, coords.y, true);
-                    updateDrawTrace(ps);
+                    ps.addTrace(coords.x, coords.y, false);
                 }
             }
         }
@@ -647,8 +659,7 @@
                 logCoords("Pointer", action, i, coords, id, event);
             }
             if (ps != null) {
-                ps.addTrace(coords.x, coords.y, false);
-                updateDrawTrace(ps);
+                ps.addTrace(coords.x, coords.y, true);
                 ps.mXVelocity = mVelocity.getXVelocity(id);
                 ps.mYVelocity = mVelocity.getYVelocity(id);
                 if (mAltVelocity != null) {
@@ -691,26 +702,13 @@
                 if (mActivePointerId == id) {
                     mActivePointerId = event.getPointerId(index == 0 ? 1 : 0);
                 }
-                ps.addTrace(Float.NaN, Float.NaN, true);
+                ps.addTrace(Float.NaN, Float.NaN, false);
             }
         }
 
         invalidate();
     }
 
-    private void updateDrawTrace(PointerState ps) {
-        mPaint.setARGB(255, 128, 255, 255);
-        float x = ps.mCurrentX;
-        float y = ps.mCurrentY;
-        float lastX = ps.mPreviousX;
-        float lastY = ps.mPreviousY;
-        if (!Float.isNaN(x) && !Float.isNaN(y) && !Float.isNaN(lastX) && !Float.isNaN(lastY)) {
-            mTraceCanvas.drawLine(lastX, lastY, x, y, mPathPaint);
-            Paint paint = ps.mPreviousPointIsHistorical ? mPaint : mCurrentPointPaint;
-            mTraceCanvas.drawPoint(lastX, lastY, paint);
-        }
-    }
-
     @Override
     public boolean onTouchEvent(MotionEvent event) {
         onPointerEvent(event);
@@ -769,7 +767,7 @@
                 return true;
             default:
                 return KeyEvent.isGamepadButton(keyCode)
-                        || KeyEvent.isModifierKey(keyCode);
+                    || KeyEvent.isModifierKey(keyCode);
         }
     }
 
@@ -889,7 +887,7 @@
         public FasterStringBuilder append(int value, int zeroPadWidth) {
             final boolean negative = value < 0;
             if (negative) {
-                value = -value;
+                value = - value;
                 if (value < 0) {
                     append("-2147483648");
                     return this;
@@ -975,27 +973,26 @@
 
     private ISystemGestureExclusionListener mSystemGestureExclusionListener =
             new ISystemGestureExclusionListener.Stub() {
-                @Override
-                public void onSystemGestureExclusionChanged(int displayId,
-                        Region systemGestureExclusion,
-                        Region systemGestureExclusionUnrestricted) {
-                    Region exclusion = Region.obtain(systemGestureExclusion);
-                    Region rejected = Region.obtain();
-                    if (systemGestureExclusionUnrestricted != null) {
-                        rejected.set(systemGestureExclusionUnrestricted);
-                        rejected.op(exclusion, Region.Op.DIFFERENCE);
-                    }
-                    Handler handler = getHandler();
-                    if (handler != null) {
-                        handler.post(() -> {
-                            mSystemGestureExclusion.set(exclusion);
-                            mSystemGestureExclusionRejected.set(rejected);
-                            exclusion.recycle();
-                            invalidate();
-                        });
-                    }
-                }
-            };
+        @Override
+        public void onSystemGestureExclusionChanged(int displayId, Region systemGestureExclusion,
+                Region systemGestureExclusionUnrestricted) {
+            Region exclusion = Region.obtain(systemGestureExclusion);
+            Region rejected = Region.obtain();
+            if (systemGestureExclusionUnrestricted != null) {
+                rejected.set(systemGestureExclusionUnrestricted);
+                rejected.op(exclusion, Region.Op.DIFFERENCE);
+            }
+            Handler handler = getHandler();
+            if (handler != null) {
+                handler.post(() -> {
+                    mSystemGestureExclusion.set(exclusion);
+                    mSystemGestureExclusionRejected.set(rejected);
+                    exclusion.recycle();
+                    invalidate();
+                });
+            }
+        }
+    };
 
     @Override
     protected void onConfigurationChanged(Configuration newConfig) {
diff --git a/core/proto/android/widget/remoteviews.proto b/core/proto/android/widget/remoteviews.proto
index f477d32..6a987a4 100644
--- a/core/proto/android/widget/remoteviews.proto
+++ b/core/proto/android/widget/remoteviews.proto
@@ -32,6 +32,8 @@
  *
  * Do not change the tag number or type of any fields in order to maintain compatibility with
  * previous versions. If a field is deleted, use `reserved` to mark its tag number.
+ *
+ * Next tag: 17
  */
 message RemoteViewsProto {
     option (android.msg_privacy).dest = DEST_AUTOMATIC;
@@ -290,6 +292,7 @@
         }
     }
 
+    // Next tag: 23
     message Action {
         oneof action {
             AttributeReflectionAction attribute_reflection_action = 1;
@@ -307,6 +310,13 @@
             SetRadioGroupCheckedAction set_radio_group_checked_action = 13;
             SetRemoteCollectionItemListAdapterAction set_remote_collection_item_list_adapter_action = 14;
             SetRippleDrawableColorAction set_ripple_drawable_color_action = 15;
+            SetViewOutlinePreferredRadiusAction set_view_outline_preferred_radius_action = 16;
+            TextViewDrawableAction text_view_drawable_action = 17;
+            TextViewSizeAction text_view_size_action = 18;
+            ViewGroupAddAction view_group_add_action = 19;
+            ViewGroupRemoveAction view_group_remove_action = 20;
+            ViewPaddingAction view_padding_action = 21;
+            SetDrawInstructionAction set_draw_instruction_action = 22;
         }
     }
 
@@ -428,6 +438,65 @@
         optional string view_id = 1;
         optional android.content.res.ColorStateListProto color_state_list = 2;
     }
+
+    message SetViewOutlinePreferredRadiusAction {
+        optional string view_id = 1;
+        optional int32 value_type = 2;
+        optional int32 value = 3;
+    }
+
+    message TextViewDrawableAction {
+        optional string view_id = 1;
+        optional bool is_relative = 2;
+        oneof drawables {
+            Resources resources = 3;
+            Icons icons = 4;
+        };
+
+        message Resources {
+            optional string one = 1;
+            optional string two = 2;
+            optional string three = 3;
+            optional string four = 4;
+        }
+
+        message Icons {
+            optional Icon one = 1;
+            optional Icon two = 2;
+            optional Icon three = 3;
+            optional Icon four = 4;
+        }
+    }
+
+    message TextViewSizeAction {
+        optional string view_id = 1;
+        optional int32 units = 2;
+        optional float size = 3;
+    }
+
+    message ViewGroupAddAction {
+        optional string view_id = 1;
+        optional RemoteViewsProto nested_views = 2;
+        optional int32 index = 3;
+        optional int32 stableId = 4;
+    }
+
+    message ViewGroupRemoveAction {
+        optional string view_id = 1;
+        optional string view_id_to_keep = 2;
+    }
+
+    message ViewPaddingAction {
+        optional string view_id = 1;
+        optional int32 left = 2;
+        optional int32 right = 3;
+        optional int32 top = 4;
+        optional int32 bottom = 5;
+    }
+
+    message SetDrawInstructionAction {
+        repeated bytes instructions = 1;
+    }
 }
 
 
diff --git a/core/res/res/values/config.xml b/core/res/res/values/config.xml
index b90ee2b..8a2d767 100644
--- a/core/res/res/values/config.xml
+++ b/core/res/res/values/config.xml
@@ -6918,6 +6918,14 @@
          entering the corresponding modes -->
     <string name="config_rearDisplayPhysicalAddress" translatable="false"></string>
 
+    <!-- The maximum number of virtual displays that can exist at the same time.
+         It must be >= 1. -->
+    <integer name="config_virtualDisplayLimit">100</integer>
+
+    <!-- The maximum number of virtual displays per package that can exist at the same time.
+         It must be >= 1. -->
+    <integer name="config_virtualDisplayLimitPerPackage">50</integer>
+
     <!-- List of certificate to be used for font fs-verity integrity verification -->
     <string-array translatable="false" name="config_fontManagerServiceCerts">
     </string-array>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index b7c8765..c1893ab 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -5240,6 +5240,9 @@
   <java-symbol type="integer" name="config_deviceStateConcurrentRearDisplay" />
   <java-symbol type="string" name="config_rearDisplayPhysicalAddress" />
 
+  <java-symbol type="integer" name="config_virtualDisplayLimit" />
+  <java-symbol type="integer" name="config_virtualDisplayLimitPerPackage" />
+
   <!-- For app language picker -->
   <java-symbol type="string" name="system_locale_title" />
   <java-symbol type="layout" name="app_language_picker_system_default" />
diff --git a/core/tests/coretests/src/android/content/res/ResourcesManagerTest.java b/core/tests/coretests/src/android/content/res/ResourcesManagerTest.java
index b16c237..be8ecbe 100644
--- a/core/tests/coretests/src/android/content/res/ResourcesManagerTest.java
+++ b/core/tests/coretests/src/android/content/res/ResourcesManagerTest.java
@@ -470,6 +470,7 @@
 
     @Test
     @SmallTest
+    @DisabledOnRavenwood(blockedBy = ResourcesManager.class)
     public void testResourceConfigurationAppliedWhenOverrideDoesNotExist() {
         final int width = 240;
         final int height = 360;
diff --git a/core/tests/coretests/src/android/window/flags/DesktopModeFlagsTest.java b/core/tests/coretests/src/android/window/DesktopModeFlagsTest.java
similarity index 88%
rename from core/tests/coretests/src/android/window/flags/DesktopModeFlagsTest.java
rename to core/tests/coretests/src/android/window/DesktopModeFlagsTest.java
index a311d0d..b28e2b0 100644
--- a/core/tests/coretests/src/android/window/flags/DesktopModeFlagsTest.java
+++ b/core/tests/coretests/src/android/window/DesktopModeFlagsTest.java
@@ -14,13 +14,14 @@
  * limitations under the License.
  */
 
-package android.window.flags;
+package android.window;
 
-import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_MODE;
-import static android.window.flags.DesktopModeFlags.ToggleOverride.OVERRIDE_OFF;
-import static android.window.flags.DesktopModeFlags.ToggleOverride.OVERRIDE_ON;
-import static android.window.flags.DesktopModeFlags.ToggleOverride.OVERRIDE_UNSET;
-import static android.window.flags.DesktopModeFlags.ToggleOverride.fromSetting;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_MODE;
+import static android.window.DesktopModeFlags.ToggleOverride.OVERRIDE_OFF;
+import static android.window.DesktopModeFlags.ToggleOverride.OVERRIDE_ON;
+import static android.window.DesktopModeFlags.ToggleOverride.OVERRIDE_UNSET;
+import static android.window.DesktopModeFlags.ToggleOverride.fromSetting;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS;
 
 import static com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE;
 import static com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_TRANSITIONS;
@@ -40,6 +41,7 @@
 import androidx.test.filters.SmallTest;
 import androidx.test.platform.app.InstrumentationRegistry;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -48,7 +50,7 @@
 import java.lang.reflect.Field;
 
 /**
- * Test class for {@link DesktopModeFlags}
+ * Test class for {@link android.window.DesktopModeFlags}
  *
  * Build/Install/Run:
  * atest FrameworksCoreTests:DesktopModeFlagsTest
@@ -68,8 +70,12 @@
     private static final int OVERRIDE_UNSET_SETTING = -1;
 
     @Before
-    public void setUp() throws Exception {
+    public void setUp() {
         mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
+    }
+
+    @After
+    public void tearDown() throws Exception {
         resetCache();
     }
 
@@ -203,7 +209,7 @@
         setOverride(OVERRIDE_UNSET_SETTING);
 
         // For unset overrides, follow flag
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
     }
 
     @Test
@@ -212,7 +218,7 @@
     public void isTrue_dwFlagOn_overrideUnset_featureFlagOff_returnsFalse() {
         setOverride(OVERRIDE_UNSET_SETTING);
         // For unset overrides, follow flag
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
     }
 
     @Test
@@ -225,7 +231,7 @@
         setOverride(OVERRIDE_ON_SETTING);
 
         // When toggle override matches its default state (dw flag), don't override flags
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
     }
 
     @Test
@@ -235,7 +241,7 @@
         setOverride(OVERRIDE_ON_SETTING);
 
         // When toggle override matches its default state (dw flag), don't override flags
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
     }
 
     @Test
@@ -248,7 +254,7 @@
         setOverride(OVERRIDE_OFF_SETTING);
 
         // Follow override if they exist, and is not equal to default toggle state (dw flag)
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
     }
 
     @Test
@@ -258,7 +264,7 @@
         setOverride(OVERRIDE_OFF_SETTING);
 
         // Follow override if they exist, and is not equal to default toggle state (dw flag)
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
     }
 
     @Test
@@ -271,7 +277,7 @@
         setOverride(OVERRIDE_UNSET_SETTING);
 
         // For unset overrides, follow flag
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
     }
 
     @Test
@@ -284,7 +290,7 @@
         setOverride(OVERRIDE_UNSET_SETTING);
 
         // For unset overrides, follow flag
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
     }
 
     @Test
@@ -297,7 +303,7 @@
         setOverride(OVERRIDE_ON_SETTING);
 
         // Follow override if they exist, and is not equal to default toggle state (dw flag)
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
     }
 
     @Test
@@ -310,7 +316,7 @@
         setOverride(OVERRIDE_ON_SETTING);
 
         // Follow override if they exist, and is not equal to default toggle state (dw flag)
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
     }
 
     @Test
@@ -323,7 +329,7 @@
         setOverride(OVERRIDE_OFF_SETTING);
 
         // When toggle override matches its default state (dw flag), don't override flags
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isTrue();
     }
 
     @Test
@@ -336,7 +342,7 @@
         setOverride(OVERRIDE_OFF_SETTING);
 
         // When toggle override matches its default state (dw flag), don't override flags
-        assertThat(DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
+        assertThat(ENABLE_DESKTOP_WINDOWING_TRANSITIONS.isTrue()).isFalse();
     }
 
     @Test
@@ -375,5 +381,6 @@
                 "sCachedToggleOverride");
         cachedToggleOverride.setAccessible(true);
         cachedToggleOverride.set(null, null);
+        setOverride(OVERRIDE_UNSET_SETTING);
     }
 }
diff --git a/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityServiceWarningTest.java b/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityServiceWarningTest.java
index 362eeea..8e906fd 100644
--- a/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityServiceWarningTest.java
+++ b/core/tests/coretests/src/com/android/internal/accessibility/dialog/AccessibilityServiceWarningTest.java
@@ -25,6 +25,8 @@
 import android.app.AlertDialog;
 import android.content.Context;
 import android.os.RemoteException;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
@@ -33,6 +35,7 @@
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.Window;
+import android.view.accessibility.Flags;
 import android.widget.TextView;
 
 import androidx.test.platform.app.InstrumentationRegistry;
@@ -89,7 +92,19 @@
     }
 
     @Test
-    public void createAccessibilityServiceWarningDialog_hasExpectedWindowParams() {
+    @RequiresFlagsDisabled(Flags.FLAG_WARNING_USE_DEFAULT_DIALOG_TYPE)
+    public void createAccessibilityServiceWarningDialog_hasExpectedWindowParams_isSystemDialog() {
+        createAccessibilityServiceWarningDialog_hasExpectedWindowParams(true);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_WARNING_USE_DEFAULT_DIALOG_TYPE)
+    public void createAccessibilityServiceWarningDialog_hasExpectedWindowParams_notSystemDialog() {
+        createAccessibilityServiceWarningDialog_hasExpectedWindowParams(false);
+    }
+
+    private void createAccessibilityServiceWarningDialog_hasExpectedWindowParams(
+            boolean expectSystemDialog) {
         final AlertDialog dialog =
                 AccessibilityServiceWarning.createAccessibilityServiceWarningDialog(
                         mContext,
@@ -101,7 +116,11 @@
         expect.that(dialogWindow.getAttributes().privateFlags
                 & SYSTEM_FLAG_HIDE_NON_SYSTEM_OVERLAY_WINDOWS).isEqualTo(
                 SYSTEM_FLAG_HIDE_NON_SYSTEM_OVERLAY_WINDOWS);
-        expect.that(dialogWindow.getAttributes().type).isEqualTo(TYPE_SYSTEM_DIALOG);
+        if (expectSystemDialog) {
+            expect.that(dialogWindow.getAttributes().type).isEqualTo(TYPE_SYSTEM_DIALOG);
+        } else {
+            expect.that(dialogWindow.getAttributes().type).isNotEqualTo(TYPE_SYSTEM_DIALOG);
+        }
     }
 
     @Test
diff --git a/core/tests/devicestatetests/Android.bp b/core/tests/devicestatetests/Android.bp
index a3303c6..a2aa62d 100644
--- a/core/tests/devicestatetests/Android.bp
+++ b/core/tests/devicestatetests/Android.bp
@@ -26,11 +26,12 @@
     // Include all test java files
     srcs: ["src/**/*.java"],
     static_libs: [
+        "androidx.test.ext.junit",
         "androidx.test.rules",
         "frameworks-base-testutils",
         "mockito-target-minus-junit4",
         "platform-test-annotations",
-        "testng",
+        "truth",
     ],
     libs: ["android.test.runner.stubs.system"],
     platform_apis: true,
diff --git a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateInfoTest.java b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateInfoTest.java
index cf7c549..1b78433 100644
--- a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateInfoTest.java
+++ b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateInfoTest.java
@@ -22,21 +22,15 @@
 import static android.hardware.devicestate.DeviceState.PROPERTY_FOLDABLE_HARDWARE_CONFIGURATION_FOLD_IN_HALF_OPEN;
 import static android.hardware.devicestate.DeviceState.PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS;
 
-import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertNotNull;
-
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertTrue;
+import static com.google.common.truth.Truth.assertThat;
 
 import android.os.Parcel;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
-import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -44,13 +38,13 @@
 
 /**
  * Unit tests for {@link DeviceStateInfo}.
- * <p/>
- * Run with <code>atest DeviceStateInfoTest</code>.
+ *
+ * <p> Build/Install/Run:
+ * atest FrameworksCoreDeviceStateManagerTests:DeviceStateInfoTest
  */
-@RunWith(JUnit4.class)
 @SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class DeviceStateInfoTest {
-
     private static final DeviceState DEVICE_STATE_0 = new DeviceState(
             new DeviceState.Configuration.Builder(0, "STATE_0")
                     .setSystemProperties(
@@ -74,88 +68,113 @@
 
     @Test
     public void create() {
-        final ArrayList<DeviceState> supportedStates = new ArrayList<>(
-                List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
+        final ArrayList<DeviceState> supportedStates =
+                new ArrayList<>(List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
         final DeviceState baseState = DEVICE_STATE_0;
         final DeviceState currentState = DEVICE_STATE_2;
 
-        final DeviceStateInfo info = new DeviceStateInfo(supportedStates, baseState, currentState);
-        assertNotNull(info.supportedStates);
-        assertEquals(supportedStates, info.supportedStates);
-        assertEquals(baseState, info.baseState);
-        assertEquals(currentState, info.currentState);
+        final DeviceStateInfo info =
+                new DeviceStateInfo(supportedStates, baseState, currentState);
+
+        assertThat(info.supportedStates).containsExactlyElementsIn(supportedStates).inOrder();
+        assertThat(info.baseState).isEqualTo(baseState);
+        assertThat(info.currentState).isEqualTo(currentState);
     }
 
     @Test
     public void equals() {
-        final ArrayList<DeviceState> supportedStates = new ArrayList<>(
-                List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
+        final ArrayList<DeviceState> supportedStates =
+                new ArrayList<>(List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
         final DeviceState baseState = DEVICE_STATE_0;
         final DeviceState currentState = DEVICE_STATE_2;
 
         final DeviceStateInfo info = new DeviceStateInfo(supportedStates, baseState, currentState);
-        Assert.assertEquals(info, info);
+        final DeviceStateInfo sameInstance = info;
+        assertThat(info).isEqualTo(sameInstance);
 
-        final DeviceStateInfo sameInfo = new DeviceStateInfo(supportedStates, baseState,
-                currentState);
-        Assert.assertEquals(info, sameInfo);
+        final DeviceStateInfo sameInfo =
+                new DeviceStateInfo(supportedStates, baseState, currentState);
+        assertThat(info).isEqualTo(sameInfo);
+
+        final DeviceStateInfo copiedInfo = new DeviceStateInfo(info);
+        assertThat(info).isEqualTo(copiedInfo);
 
         final DeviceStateInfo differentInfo = new DeviceStateInfo(
                 new ArrayList<>(List.of(DEVICE_STATE_0, DEVICE_STATE_2)), baseState, currentState);
-        assertNotEquals(info, differentInfo);
+        assertThat(differentInfo).isNotEqualTo(info);
+    }
+
+    @Test
+    public void hashCode_sameObject() {
+        final ArrayList<DeviceState> supportedStates =
+                new ArrayList<>(List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
+        final DeviceState baseState = DEVICE_STATE_0;
+        final DeviceState currentState = DEVICE_STATE_2;
+        final DeviceStateInfo info =
+                new DeviceStateInfo(supportedStates, baseState, currentState);
+        final DeviceStateInfo copiedInfo = new DeviceStateInfo(info);
+
+        assertThat(info.hashCode()).isEqualTo(copiedInfo.hashCode());
     }
 
     @Test
     public void diff_sameObject() {
-        final ArrayList<DeviceState> supportedStates = new ArrayList<>(
-                List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
+        final ArrayList<DeviceState> supportedStates =
+                new ArrayList<>(List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
         final DeviceState baseState = DEVICE_STATE_0;
         final DeviceState currentState = DEVICE_STATE_2;
 
         final DeviceStateInfo info = new DeviceStateInfo(supportedStates, baseState, currentState);
-        assertEquals(0, info.diff(info));
+
+        assertThat(info.diff(info)).isEqualTo(0);
     }
 
     @Test
     public void diff_differentSupportedStates() {
-        final DeviceStateInfo info = new DeviceStateInfo(new ArrayList<>(List.of(DEVICE_STATE_1)),
-                DEVICE_STATE_0, DEVICE_STATE_0);
+        final DeviceStateInfo info = new DeviceStateInfo(
+                new ArrayList<>(List.of(DEVICE_STATE_1)), DEVICE_STATE_0, DEVICE_STATE_0);
         final DeviceStateInfo otherInfo = new DeviceStateInfo(
                 new ArrayList<>(List.of(DEVICE_STATE_2)), DEVICE_STATE_0, DEVICE_STATE_0);
+
         final int diff = info.diff(otherInfo);
-        assertTrue((diff & DeviceStateInfo.CHANGED_SUPPORTED_STATES) > 0);
-        assertFalse((diff & DeviceStateInfo.CHANGED_BASE_STATE) > 0);
-        assertFalse((diff & DeviceStateInfo.CHANGED_CURRENT_STATE) > 0);
+
+        assertThat(diff & DeviceStateInfo.CHANGED_SUPPORTED_STATES).isGreaterThan(0);
+        assertThat(diff & DeviceStateInfo.CHANGED_BASE_STATE).isEqualTo(0);
+        assertThat(diff & DeviceStateInfo.CHANGED_CURRENT_STATE).isEqualTo(0);
     }
 
     @Test
     public void diff_differentNonOverrideState() {
-        final DeviceStateInfo info = new DeviceStateInfo(new ArrayList<>(List.of(DEVICE_STATE_1)),
-                DEVICE_STATE_1, DEVICE_STATE_0);
+        final DeviceStateInfo info = new DeviceStateInfo(
+                new ArrayList<>(List.of(DEVICE_STATE_1)), DEVICE_STATE_1, DEVICE_STATE_0);
         final DeviceStateInfo otherInfo = new DeviceStateInfo(
                 new ArrayList<>(List.of(DEVICE_STATE_1)), DEVICE_STATE_2, DEVICE_STATE_0);
+
         final int diff = info.diff(otherInfo);
-        assertFalse((diff & DeviceStateInfo.CHANGED_SUPPORTED_STATES) > 0);
-        assertTrue((diff & DeviceStateInfo.CHANGED_BASE_STATE) > 0);
-        assertFalse((diff & DeviceStateInfo.CHANGED_CURRENT_STATE) > 0);
+
+        assertThat(diff & DeviceStateInfo.CHANGED_SUPPORTED_STATES).isEqualTo(0);
+        assertThat(diff & DeviceStateInfo.CHANGED_BASE_STATE).isGreaterThan(0);
+        assertThat(diff & DeviceStateInfo.CHANGED_CURRENT_STATE).isEqualTo(0);
     }
 
     @Test
     public void diff_differentState() {
-        final DeviceStateInfo info = new DeviceStateInfo(new ArrayList<>(List.of(DEVICE_STATE_1)),
-                DEVICE_STATE_0, DEVICE_STATE_1);
+        final DeviceStateInfo info = new DeviceStateInfo(
+                new ArrayList<>(List.of(DEVICE_STATE_1)), DEVICE_STATE_0, DEVICE_STATE_1);
         final DeviceStateInfo otherInfo = new DeviceStateInfo(
                 new ArrayList<>(List.of(DEVICE_STATE_1)), DEVICE_STATE_0, DEVICE_STATE_2);
+
         final int diff = info.diff(otherInfo);
-        assertFalse((diff & DeviceStateInfo.CHANGED_SUPPORTED_STATES) > 0);
-        assertFalse((diff & DeviceStateInfo.CHANGED_BASE_STATE) > 0);
-        assertTrue((diff & DeviceStateInfo.CHANGED_CURRENT_STATE) > 0);
+
+        assertThat(diff & DeviceStateInfo.CHANGED_SUPPORTED_STATES).isEqualTo(0);
+        assertThat(diff & DeviceStateInfo.CHANGED_BASE_STATE).isEqualTo(0);
+        assertThat(diff & DeviceStateInfo.CHANGED_CURRENT_STATE).isGreaterThan(0);
     }
 
     @Test
     public void writeToParcel() {
-        final ArrayList<DeviceState> supportedStates = new ArrayList<>(
-                List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
+        final ArrayList<DeviceState> supportedStates =
+                new ArrayList<>(List.of(DEVICE_STATE_0, DEVICE_STATE_1, DEVICE_STATE_2));
         final DeviceState nonOverrideState = DEVICE_STATE_0;
         final DeviceState state = DEVICE_STATE_2;
         final DeviceStateInfo originalInfo =
@@ -166,6 +185,6 @@
         parcel.setDataPosition(0);
 
         final DeviceStateInfo info = DeviceStateInfo.CREATOR.createFromParcel(parcel);
-        assertEquals(originalInfo, info);
+        assertThat(info).isEqualTo(originalInfo);
     }
 }
diff --git a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java
index f4d3631..7c01ecc 100644
--- a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java
+++ b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateManagerGlobalTest.java
@@ -16,8 +16,7 @@
 
 package android.hardware.devicestate;
 
-import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertFalse;
+import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
@@ -31,6 +30,9 @@
 import android.os.RemoteException;
 import android.os.test.FakePermissionEnforcer;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.util.ConcurrentUtils;
@@ -38,7 +40,6 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -47,32 +48,35 @@
 
 /**
  * Unit tests for {@link DeviceStateManagerGlobal}.
- * <p/>
- * Run with <code>atest DeviceStateManagerGlobalTest</code>.
+ *
+ * <p> Build/Install/Run:
+ * atest FrameworksCoreDeviceStateManagerTests:DeviceStateManagerGlobalTest
  */
-@RunWith(JUnit4.class)
 @SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class DeviceStateManagerGlobalTest {
     private static final DeviceState DEFAULT_DEVICE_STATE = new DeviceState(
             new DeviceState.Configuration.Builder(0 /* identifier */, "" /* name */).build());
     private static final DeviceState OTHER_DEVICE_STATE = new DeviceState(
             new DeviceState.Configuration.Builder(1 /* identifier */, "" /* name */).build());
 
+    @NonNull
     private TestDeviceStateManagerService mService;
+    @NonNull
     private DeviceStateManagerGlobal mDeviceStateManagerGlobal;
 
     @Before
     public void setUp() {
-        FakePermissionEnforcer permissionEnforcer = new FakePermissionEnforcer();
+        final FakePermissionEnforcer permissionEnforcer = new FakePermissionEnforcer();
         mService = new TestDeviceStateManagerService(permissionEnforcer);
         mDeviceStateManagerGlobal = new DeviceStateManagerGlobal(mService);
-        assertFalse(mService.mCallbacks.isEmpty());
+        assertThat(mService.mCallbacks).isNotEmpty();
     }
 
     @Test
     public void registerCallback() {
-        DeviceStateCallback callback1 = mock(DeviceStateCallback.class);
-        DeviceStateCallback callback2 = mock(DeviceStateCallback.class);
+        final DeviceStateCallback callback1 = mock(DeviceStateCallback.class);
+        final DeviceStateCallback callback2 = mock(DeviceStateCallback.class);
 
         mDeviceStateManagerGlobal.registerDeviceStateCallback(callback1,
                 ConcurrentUtils.DIRECT_EXECUTOR);
@@ -105,8 +109,8 @@
         reset(callback2);
 
         // Change the requested state and verify callback
-        DeviceStateRequest request = DeviceStateRequest.newBuilder(
-                DEFAULT_DEVICE_STATE.getIdentifier()).build();
+        final DeviceStateRequest request =
+                DeviceStateRequest.newBuilder(DEFAULT_DEVICE_STATE.getIdentifier()).build();
         mDeviceStateManagerGlobal.requestState(request, null /* executor */, null /* callback */);
 
         verify(callback1).onDeviceStateChanged(eq(mService.getMergedState()));
@@ -115,10 +119,10 @@
 
     @Test
     public void unregisterCallback() {
-        DeviceStateCallback callback = mock(DeviceStateCallback.class);
+        final DeviceStateCallback callback = mock(DeviceStateCallback.class);
 
-        mDeviceStateManagerGlobal.registerDeviceStateCallback(callback,
-                ConcurrentUtils.DIRECT_EXECUTOR);
+        mDeviceStateManagerGlobal
+                .registerDeviceStateCallback(callback, ConcurrentUtils.DIRECT_EXECUTOR);
 
         // Verify initial callbacks
         verify(callback).onSupportedStatesChanged(eq(mService.getSupportedDeviceStates()));
@@ -134,15 +138,15 @@
 
     @Test
     public void submitRequest() {
-        DeviceStateCallback callback = mock(DeviceStateCallback.class);
-        mDeviceStateManagerGlobal.registerDeviceStateCallback(callback,
-                ConcurrentUtils.DIRECT_EXECUTOR);
+        final DeviceStateCallback callback = mock(DeviceStateCallback.class);
+        mDeviceStateManagerGlobal
+                .registerDeviceStateCallback(callback, ConcurrentUtils.DIRECT_EXECUTOR);
 
         verify(callback).onDeviceStateChanged(eq(mService.getBaseState()));
         reset(callback);
 
-        DeviceStateRequest request = DeviceStateRequest.newBuilder(
-                OTHER_DEVICE_STATE.getIdentifier()).build();
+        final DeviceStateRequest request =
+                DeviceStateRequest.newBuilder(OTHER_DEVICE_STATE.getIdentifier()).build();
         mDeviceStateManagerGlobal.requestState(request, null /* executor */, null /* callback */);
 
         verify(callback).onDeviceStateChanged(eq(OTHER_DEVICE_STATE));
@@ -155,15 +159,15 @@
 
     @Test
     public void submitBaseStateOverrideRequest() {
-        DeviceStateCallback callback = mock(DeviceStateCallback.class);
-        mDeviceStateManagerGlobal.registerDeviceStateCallback(callback,
-                ConcurrentUtils.DIRECT_EXECUTOR);
+        final DeviceStateCallback callback = mock(DeviceStateCallback.class);
+        mDeviceStateManagerGlobal
+                .registerDeviceStateCallback(callback, ConcurrentUtils.DIRECT_EXECUTOR);
 
         verify(callback).onDeviceStateChanged(eq(mService.getBaseState()));
         reset(callback);
 
-        DeviceStateRequest request = DeviceStateRequest.newBuilder(
-                OTHER_DEVICE_STATE.getIdentifier()).build();
+        final DeviceStateRequest request =
+                DeviceStateRequest.newBuilder(OTHER_DEVICE_STATE.getIdentifier()).build();
         mDeviceStateManagerGlobal.requestBaseStateOverride(request, null /* executor */,
                 null /* callback */);
 
@@ -177,28 +181,28 @@
 
     @Test
     public void submitBaseAndEmulatedStateOverride() {
-        DeviceStateCallback callback = mock(DeviceStateCallback.class);
-        mDeviceStateManagerGlobal.registerDeviceStateCallback(callback,
-                ConcurrentUtils.DIRECT_EXECUTOR);
+        final DeviceStateCallback callback = mock(DeviceStateCallback.class);
+        mDeviceStateManagerGlobal
+                .registerDeviceStateCallback(callback, ConcurrentUtils.DIRECT_EXECUTOR);
 
         verify(callback).onDeviceStateChanged(eq(mService.getBaseState()));
         reset(callback);
 
-        DeviceStateRequest request = DeviceStateRequest.newBuilder(
-                OTHER_DEVICE_STATE.getIdentifier()).build();
+        final DeviceStateRequest request =
+                DeviceStateRequest.newBuilder(OTHER_DEVICE_STATE.getIdentifier()).build();
         mDeviceStateManagerGlobal.requestBaseStateOverride(request, null /* executor */,
                 null /* callback */);
 
         verify(callback).onDeviceStateChanged(eq(OTHER_DEVICE_STATE));
-        assertEquals(OTHER_DEVICE_STATE, mService.getBaseState());
+        assertThat(mService.getBaseState()).isEqualTo(OTHER_DEVICE_STATE);
         reset(callback);
 
-        DeviceStateRequest secondRequest = DeviceStateRequest.newBuilder(
-                DEFAULT_DEVICE_STATE.getIdentifier()).build();
+        final DeviceStateRequest secondRequest =
+                DeviceStateRequest.newBuilder(DEFAULT_DEVICE_STATE.getIdentifier()).build();
 
         mDeviceStateManagerGlobal.requestState(secondRequest, null, null);
 
-        assertEquals(OTHER_DEVICE_STATE, mService.getBaseState());
+        assertThat(mService.getBaseState()).isEqualTo(OTHER_DEVICE_STATE);
         verify(callback).onDeviceStateChanged(eq(DEFAULT_DEVICE_STATE));
         reset(callback);
 
@@ -214,10 +218,10 @@
 
     @Test
     public void verifyDeviceStateRequestCallbacksCalled() {
-        DeviceStateRequest.Callback callback = mock(TestDeviceStateRequestCallback.class);
+        final DeviceStateRequest.Callback callback = mock(TestDeviceStateRequestCallback.class);
 
-        DeviceStateRequest request = DeviceStateRequest.newBuilder(
-                OTHER_DEVICE_STATE.getIdentifier()).build();
+        final DeviceStateRequest request =
+                DeviceStateRequest.newBuilder(OTHER_DEVICE_STATE.getIdentifier()).build();
         mDeviceStateManagerGlobal.requestState(request,
                 ConcurrentUtils.DIRECT_EXECUTOR /* executor */,
                 callback /* callback */);
@@ -232,52 +236,55 @@
 
     public static class TestDeviceStateRequestCallback implements DeviceStateRequest.Callback {
         @Override
-        public void onRequestActivated(DeviceStateRequest request) { }
+        public void onRequestActivated(@NonNull DeviceStateRequest request) { }
 
         @Override
-        public void onRequestCanceled(DeviceStateRequest request) { }
+        public void onRequestCanceled(@NonNull DeviceStateRequest request) { }
 
         @Override
-        public void onRequestSuspended(DeviceStateRequest request) { }
+        public void onRequestSuspended(@NonNull DeviceStateRequest request) { }
     }
 
     private static final class TestDeviceStateManagerService extends IDeviceStateManager.Stub {
-        public static final class Request {
-            public final IBinder token;
-            public final int state;
-            public final int flags;
+        static final class Request {
+            @NonNull
+            final IBinder mToken;
+            final int mState;
 
-            private Request(IBinder token, int state, int flags) {
-                this.token = token;
-                this.state = state;
-                this.flags = flags;
+            private Request(@NonNull IBinder token, int state) {
+                this.mToken = token;
+                this.mState = state;
             }
         }
 
-        private List<DeviceState> mSupportedDeviceStates = List.of(DEFAULT_DEVICE_STATE,
-                OTHER_DEVICE_STATE);
-
+        @NonNull
+        private List<DeviceState> mSupportedDeviceStates =
+                List.of(DEFAULT_DEVICE_STATE, OTHER_DEVICE_STATE);
+        @NonNull
         private DeviceState mBaseState = DEFAULT_DEVICE_STATE;
+        @Nullable
         private Request mRequest;
+        @Nullable
         private Request mBaseStateRequest;
 
         private final Set<IDeviceStateManagerCallback> mCallbacks = new HashSet<>();
 
-        TestDeviceStateManagerService(FakePermissionEnforcer enforcer) {
+        TestDeviceStateManagerService(@NonNull FakePermissionEnforcer enforcer) {
             super(enforcer);
         }
 
+        @NonNull
         private DeviceStateInfo getInfo() {
             final int mergedBaseState = mBaseStateRequest == null
-                    ? mBaseState.getIdentifier() : mBaseStateRequest.state;
-            final int mergedState = mRequest == null
-                    ? mergedBaseState : mRequest.state;
+                    ? mBaseState.getIdentifier() : mBaseStateRequest.mState;
+            final int mergedState = mRequest == null ? mergedBaseState : mRequest.mState;
 
+            final ArrayList<DeviceState> supportedStates = new ArrayList<>(mSupportedDeviceStates);
             final DeviceState baseState = new DeviceState(
                     new DeviceState.Configuration.Builder(mergedBaseState, "" /* name */).build());
             final DeviceState state = new DeviceState(
                     new DeviceState.Configuration.Builder(mergedState, "" /* name */).build());
-            return new DeviceStateInfo(new ArrayList<>(mSupportedDeviceStates), baseState, state);
+            return new DeviceStateInfo(supportedStates, baseState, state);
         }
 
         private void notifyDeviceStateInfoChanged() {
@@ -291,6 +298,7 @@
             }
         }
 
+        @NonNull
         @Override
         public DeviceStateInfo getDeviceStateInfo() {
             return getInfo();
@@ -311,18 +319,18 @@
         }
 
         @Override
-        public void requestState(IBinder token, int state, int flags) {
+        public void requestState(@NonNull IBinder token, int state, int unusedFlags) {
             if (mRequest != null) {
                 for (IDeviceStateManagerCallback callback : mCallbacks) {
                     try {
-                        callback.onRequestCanceled(mRequest.token);
+                        callback.onRequestCanceled(mRequest.mToken);
                     } catch (RemoteException e) {
                         e.rethrowFromSystemServer();
                     }
                 }
             }
 
-            final Request request = new Request(token, state, flags);
+            final Request request = new Request(token, state);
             mRequest = request;
             notifyDeviceStateInfoChanged();
 
@@ -337,7 +345,7 @@
 
         @Override
         public void cancelStateRequest() {
-            IBinder token = mRequest.token;
+            final IBinder token = mRequest.mToken;
             mRequest = null;
             for (IDeviceStateManagerCallback callback : mCallbacks) {
                 try {
@@ -350,19 +358,18 @@
         }
 
         @Override
-        public void requestBaseStateOverride(IBinder token, int state, int flags) {
+        public void requestBaseStateOverride(@NonNull IBinder token, int state, int unusedFlags) {
             if (mBaseStateRequest != null) {
                 for (IDeviceStateManagerCallback callback : mCallbacks) {
                     try {
-                        callback.onRequestCanceled(mBaseStateRequest.token);
+                        callback.onRequestCanceled(mBaseStateRequest.mToken);
                     } catch (RemoteException e) {
                         e.rethrowFromSystemServer();
                     }
                 }
             }
 
-            final Request request = new Request(token, state, flags);
-            mBaseStateRequest = request;
+            mBaseStateRequest = new Request(token, state);
             notifyDeviceStateInfoChanged();
 
             for (IDeviceStateManagerCallback callback : mCallbacks) {
@@ -376,7 +383,7 @@
 
         @Override
         public void cancelBaseStateOverride() throws RemoteException {
-            IBinder token = mBaseStateRequest.token;
+            final IBinder token = mBaseStateRequest.mToken;
             mBaseStateRequest = null;
             for (IDeviceStateManagerCallback callback : mCallbacks) {
                 try {
@@ -396,24 +403,27 @@
             onStateRequestOverlayDismissed_enforcePermission();
         }
 
-        public void setSupportedStates(List<DeviceState> states) {
+        public void setSupportedStates(@NonNull List<DeviceState> states) {
             mSupportedDeviceStates = states;
             notifyDeviceStateInfoChanged();
         }
 
+        @NonNull
         public List<DeviceState> getSupportedDeviceStates() {
             return mSupportedDeviceStates;
         }
 
-        public void setBaseState(DeviceState state) {
+        public void setBaseState(@NonNull DeviceState state) {
             mBaseState = state;
             notifyDeviceStateInfoChanged();
         }
 
+        @NonNull
         public DeviceState getBaseState() {
             return getInfo().baseState;
         }
 
+        @NonNull
         public DeviceState getMergedState() {
             return getInfo().currentState;
         }
diff --git a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateTest.java b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateTest.java
index 78d4324..83b5ff3 100644
--- a/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateTest.java
+++ b/core/tests/devicestatetests/src/android/hardware/devicestate/DeviceStateTest.java
@@ -21,17 +21,16 @@
 import static android.hardware.devicestate.DeviceState.PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS;
 import static android.hardware.devicestate.DeviceStateManager.MINIMUM_DEVICE_STATE_IDENTIFIER;
 
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertTrue;
+import static com.google.common.truth.Truth.assertThat;
 
 import android.os.Parcel;
 import android.platform.test.annotations.Presubmit;
 
-import junit.framework.Assert;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
 
 import java.util.HashSet;
 import java.util.List;
@@ -39,28 +38,32 @@
 
 /**
  * Unit tests for {@link android.hardware.devicestate.DeviceState}.
- * <p/>
- * Run with <code>atest DeviceStateTest</code>.
+ *
+ * <p> Build/Install/Run:
+ * atest FrameworksCoreDeviceStateManagerTests:DeviceStateTest
  */
 @Presubmit
-@RunWith(JUnit4.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public final class DeviceStateTest {
     @Test
     public void testConstruct() {
-        DeviceState.Configuration config = new DeviceState.Configuration.Builder(
+        final DeviceState.Configuration config = new DeviceState.Configuration.Builder(
                 MINIMUM_DEVICE_STATE_IDENTIFIER, "TEST_CLOSED")
                 .setSystemProperties(
                         new HashSet<>(List.of(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS)))
                 .build();
+
         final DeviceState state = new DeviceState(config);
-        assertEquals(state.getIdentifier(), MINIMUM_DEVICE_STATE_IDENTIFIER);
-        assertEquals(state.getName(), "TEST_CLOSED");
-        assertTrue(state.hasProperty(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS));
+
+        assertThat(state.getIdentifier()).isEqualTo(MINIMUM_DEVICE_STATE_IDENTIFIER);
+        assertThat(state.getName()).isEqualTo("TEST_CLOSED");
+        assertThat(state.hasProperty(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS)).isTrue();
     }
 
     @Test
     public void testHasProperties() {
-        DeviceState.Configuration config = new DeviceState.Configuration.Builder(
+        final DeviceState.Configuration config = new DeviceState.Configuration.Builder(
                 MINIMUM_DEVICE_STATE_IDENTIFIER, "TEST")
                 .setSystemProperties(new HashSet<>(List.of(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS,
                         PROPERTY_POLICY_AVAILABLE_FOR_APP_REQUEST)))
@@ -68,10 +71,10 @@
 
         final DeviceState state = new DeviceState(config);
 
-        assertTrue(state.hasProperty(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS));
-        assertTrue(state.hasProperty(PROPERTY_POLICY_AVAILABLE_FOR_APP_REQUEST));
-        assertTrue(state.hasProperties(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS,
-                PROPERTY_POLICY_AVAILABLE_FOR_APP_REQUEST));
+        assertThat(state.hasProperty(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS)).isTrue();
+        assertThat(state.hasProperty(PROPERTY_POLICY_AVAILABLE_FOR_APP_REQUEST)).isTrue();
+        assertThat(state.hasProperties(PROPERTY_POLICY_CANCEL_OVERRIDE_REQUESTS,
+                PROPERTY_POLICY_AVAILABLE_FOR_APP_REQUEST)).isTrue();
     }
 
     @Test
@@ -91,7 +94,7 @@
         final DeviceState.Configuration stateConfiguration =
                 DeviceState.Configuration.CREATOR.createFromParcel(parcel);
 
-        Assert.assertEquals(originalState, new DeviceState(stateConfiguration));
+        assertThat(originalState).isEqualTo(new DeviceState(stateConfiguration));
     }
 
     @Test
@@ -109,6 +112,6 @@
         final DeviceState.Configuration stateConfiguration =
                 DeviceState.Configuration.CREATOR.createFromParcel(parcel);
 
-        Assert.assertEquals(originalState, new DeviceState(stateConfiguration));
+        assertThat(originalState).isEqualTo(new DeviceState(stateConfiguration));
     }
 }
diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index e493ed1..4642fe5 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -234,10 +234,6 @@
         // *.kt sources are inside a filegroup.
         "kotlin-annotations",
     ],
-    required: [
-        "wmshell.protolog.json.gz",
-        "wmshell.protolog.pb",
-    ],
     flags_packages: [
         "com_android_wm_shell_flags",
     ],
@@ -246,3 +242,11 @@
     plugins: ["dagger2-compiler"],
     use_resource_processor: true,
 }
+
+java_defaults {
+    name: "wmshell_defaults",
+    required: [
+        "wmshell.protolog.json.gz",
+        "wmshell.protolog.pb",
+    ],
+}
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/FocusTransitionListener.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/FocusTransitionListener.java
index 26aae2d..02a7991 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/FocusTransitionListener.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/FocusTransitionListener.java
@@ -26,5 +26,11 @@
     /**
      * Called when a transition changes the top, focused display.
      */
-    void onFocusedDisplayChanged(int displayId);
+    default void onFocusedDisplayChanged(int displayId) {}
+
+    /**
+     * Called when the per-app or system-wide focus state has changed for a task.
+     */
+    default void onFocusedTaskChanged(int taskId, boolean isFocusedOnDisplay,
+            boolean isFocusedGlobally) {}
 }
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/FlyoutDrawableLoader.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/FlyoutDrawableLoader.kt
new file mode 100644
index 0000000..5a17330
--- /dev/null
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/FlyoutDrawableLoader.kt
@@ -0,0 +1,47 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.shared.bubbles
+
+import android.content.Context
+import android.content.Intent
+import android.graphics.drawable.Drawable
+import android.graphics.drawable.Icon
+import android.util.Log
+
+object FlyoutDrawableLoader {
+
+    private const val TAG = "FlyoutDrawableLoader"
+
+    /** Loads the flyout icon as a [Drawable]. */
+    @JvmStatic
+    fun Icon?.loadFlyoutDrawable(context: Context): Drawable? {
+        if (this == null) return null
+        try {
+            if (this.type == Icon.TYPE_URI || this.type == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
+                context.grantUriPermission(
+                    context.packageName,
+                    this.uri,
+                    Intent.FLAG_GRANT_READ_URI_PERMISSION
+                )
+            }
+            return loadDrawable(context)
+        } catch (e: Exception) {
+            Log.w(TAG, "loadFlyoutDrawable failed: ${e.message}")
+            return null
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeStatus.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeStatus.java
index 647a555a..0150bcd 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeStatus.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeStatus.java
@@ -19,7 +19,7 @@
 import android.annotation.NonNull;
 import android.content.Context;
 import android.os.SystemProperties;
-import android.window.flags.DesktopModeFlags;
+import android.window.DesktopModeFlags;
 
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleLogger.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleLogger.java
index c88a58b..1abe119 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleLogger.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleLogger.java
@@ -36,6 +36,8 @@
     @VisibleForTesting
     public enum Event implements UiEventLogger.UiEventEnum {
 
+        // region bubble events
+
         @UiEvent(doc = "User dismissed the bubble via gesture, add bubble to overflow.")
         BUBBLE_OVERFLOW_ADD_USER_GESTURE(483),
 
@@ -64,7 +66,89 @@
         BUBBLE_OVERFLOW_SELECTED(600),
 
         @UiEvent(doc = "Restore bubble to overflow after phone reboot.")
-        BUBBLE_OVERFLOW_RECOVER(691);
+        BUBBLE_OVERFLOW_RECOVER(691),
+
+        // endregion
+
+        // region bubble bar events
+
+        @UiEvent(doc = "new bubble posted")
+        BUBBLE_BAR_BUBBLE_POSTED(1927),
+
+        @UiEvent(doc = "existing bubble updated")
+        BUBBLE_BAR_BUBBLE_UPDATED(1928),
+
+        @UiEvent(doc = "expanded a bubble from bubble bar")
+        BUBBLE_BAR_EXPANDED(1929),
+
+        @UiEvent(doc = "bubble bar collapsed")
+        BUBBLE_BAR_COLLAPSED(1930),
+
+        @UiEvent(doc = "dismissed single bubble from bubble bar by dragging it to dismiss target")
+        BUBBLE_BAR_BUBBLE_DISMISSED_DRAG_BUBBLE(1931),
+
+        @UiEvent(doc = "dismissed single bubble from bubble bar by dragging the expanded view to "
+                + "dismiss target")
+        BUBBLE_BAR_BUBBLE_DISMISSED_DRAG_EXP_VIEW(1932),
+
+        @UiEvent(doc = "dismiss bubble from app handle menu")
+        BUBBLE_BAR_BUBBLE_DISMISSED_APP_MENU(1933),
+
+        @UiEvent(doc = "bubble is dismissed due to app finishing the bubble activity")
+        BUBBLE_BAR_BUBBLE_ACTIVITY_FINISH(1934),
+
+        @UiEvent(doc = "dismissed the bubble bar by dragging it to dismiss target")
+        BUBBLE_BAR_DISMISSED_DRAG_BAR(1935),
+
+        @UiEvent(doc = "bubble bar moved to the left edge of the screen by dragging from the "
+                + "expanded view")
+        BUBBLE_BAR_MOVED_LEFT_DRAG_EXP_VIEW(1936),
+
+        @UiEvent(doc = "bubble bar moved to the left edge of the screen by dragging from a single"
+                + " bubble")
+        BUBBLE_BAR_MOVED_LEFT_DRAG_BUBBLE(1937),
+
+        @UiEvent(doc = "bubble bar moved to the left edge of the screen by dragging the bubble bar")
+        BUBBLE_BAR_MOVED_LEFT_DRAG_BAR(1938),
+
+        @UiEvent(doc = "bubble bar moved to the right edge of the screen by dragging from the "
+                + "expanded view")
+        BUBBLE_BAR_MOVED_RIGHT_DRAG_EXP_VIEW(1939),
+
+        @UiEvent(doc = "bubble bar moved to the right edge of the screen by dragging from a "
+                + "single bubble")
+        BUBBLE_BAR_MOVED_RIGHT_DRAG_BUBBLE(1940),
+
+        @UiEvent(doc = "bubble bar moved to the right edge of the screen by dragging the bubble "
+                + "bar")
+        BUBBLE_BAR_MOVED_RIGHT_DRAG_BAR(1941),
+
+        @UiEvent(doc = "stop bubbling conversation from app handle menu")
+        BUBBLE_BAR_APP_MENU_OPT_OUT(1942),
+
+        @UiEvent(doc = "open app settings from app handle menu")
+        BUBBLE_BAR_APP_MENU_GO_TO_SETTINGS(1943),
+
+        @UiEvent(doc = "flyout shown for a bubble")
+        BUBBLE_BAR_FLYOUT(1944),
+
+        @UiEvent(doc = "notification for the bubble was canceled")
+        BUBBLE_BAR_BUBBLE_REMOVED_CANCELED(1945),
+
+        @UiEvent(doc = "user turned off bubbles from settings")
+        BUBBLE_BAR_BUBBLE_REMOVED_BLOCKED(1946),
+
+        @UiEvent(doc = "bubble bar overflow opened")
+        BUBBLE_BAR_OVERFLOW_SELECTED(1947),
+
+        @UiEvent(doc = "max number of bubbles was reached in bubble bar, move bubble to overflow")
+        BUBBLE_BAR_OVERFLOW_ADD_AGED(1948),
+
+        @UiEvent(doc = "bubble promoted from overflow back to bubble bar")
+        BUBBLE_BAR_OVERFLOW_REMOVE_BACK_TO_BAR(1949),
+
+        // endregion
+        ;
 
         private final int mId;
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java
index c5e3afd..39fb2f49 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTask.java
@@ -21,11 +21,10 @@
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.TAG_BUBBLES;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_BUBBLES;
+import static com.android.wm.shell.shared.bubbles.FlyoutDrawableLoader.loadFlyoutDrawable;
 
-import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
-import android.content.Intent;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
@@ -34,7 +33,6 @@
 import android.graphics.Matrix;
 import android.graphics.Path;
 import android.graphics.drawable.Drawable;
-import android.graphics.drawable.Icon;
 import android.util.Log;
 import android.util.PathParser;
 import android.view.LayoutInflater;
@@ -51,7 +49,6 @@
 import com.android.wm.shell.shared.handles.RegionSamplingHelper;
 
 import java.lang.ref.WeakReference;
-import java.util.Objects;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -340,7 +337,7 @@
             info.flyoutMessage = b.getFlyoutMessage();
             if (info.flyoutMessage != null) {
                 info.flyoutMessage.senderAvatar =
-                        loadSenderAvatar(c, info.flyoutMessage.senderIcon);
+                        loadFlyoutDrawable(info.flyoutMessage.senderIcon, c);
             }
             return info;
         }
@@ -422,21 +419,4 @@
                 Color.WHITE, WHITE_SCRIM_ALPHA);
         return true;
     }
-
-    @Nullable
-    static Drawable loadSenderAvatar(@NonNull final Context context, @Nullable final Icon icon) {
-        Objects.requireNonNull(context);
-        if (icon == null) return null;
-        try {
-            if (icon.getType() == Icon.TYPE_URI
-                    || icon.getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
-                context.grantUriPermission(context.getPackageName(),
-                        icon.getUri(), Intent.FLAG_GRANT_READ_URI_PERMISSION);
-            }
-            return icon.loadDrawable(context);
-        } catch (Exception e) {
-            Log.w(TAG, "loadSenderAvatar failed: " + e.getMessage());
-            return null;
-        }
-    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTaskLegacy.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTaskLegacy.java
index c12822a..e9a5933 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTaskLegacy.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleViewInfoTaskLegacy.java
@@ -20,11 +20,10 @@
 import static com.android.wm.shell.bubbles.BadgedImageView.WHITE_SCRIM_ALPHA;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.TAG_BUBBLES;
 import static com.android.wm.shell.bubbles.BubbleDebugConfig.TAG_WITH_CLASS_NAME;
+import static com.android.wm.shell.shared.bubbles.FlyoutDrawableLoader.loadFlyoutDrawable;
 
-import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.content.Context;
-import android.content.Intent;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
@@ -33,7 +32,6 @@
 import android.graphics.Matrix;
 import android.graphics.Path;
 import android.graphics.drawable.Drawable;
-import android.graphics.drawable.Icon;
 import android.os.AsyncTask;
 import android.util.Log;
 import android.util.PathParser;
@@ -50,7 +48,6 @@
 import com.android.wm.shell.shared.handles.RegionSamplingHelper;
 
 import java.lang.ref.WeakReference;
-import java.util.Objects;
 import java.util.concurrent.Executor;
 
 /**
@@ -264,7 +261,7 @@
             info.flyoutMessage = b.getFlyoutMessage();
             if (info.flyoutMessage != null) {
                 info.flyoutMessage.senderAvatar =
-                        loadSenderAvatar(c, info.flyoutMessage.senderIcon);
+                        loadFlyoutDrawable(info.flyoutMessage.senderIcon, c);
             }
             return info;
         }
@@ -346,21 +343,4 @@
                 Color.WHITE, WHITE_SCRIM_ALPHA);
         return true;
     }
-
-    @Nullable
-    static Drawable loadSenderAvatar(@NonNull final Context context, @Nullable final Icon icon) {
-        Objects.requireNonNull(context);
-        if (icon == null) return null;
-        try {
-            if (icon.getType() == Icon.TYPE_URI
-                    || icon.getType() == Icon.TYPE_URI_ADAPTIVE_BITMAP) {
-                context.grantUriPermission(context.getPackageName(),
-                        icon.getUri(), Intent.FLAG_GRANT_READ_URI_PERMISSION);
-            }
-            return icon.loadDrawable(context);
-        } catch (Exception e) {
-            Log.w(TAG, "loadSenderAvatar failed: " + e.getMessage());
-            return null;
-        }
-    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
index 4d15605c..2128cbc 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
@@ -27,7 +27,7 @@
 import android.util.Pair;
 import android.view.LayoutInflater;
 import android.view.View;
-import android.window.flags.DesktopModeFlags;
+import android.window.DesktopModeFlags;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.window.flags.Flags;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 79c31e0..75adef4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -16,8 +16,8 @@
 
 package com.android.wm.shell.dagger;
 
-import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_EXIT_TRANSITIONS;
-import static android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TASK_LIMIT;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_EXIT_TRANSITIONS;
+import static android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TASK_LIMIT;
 
 import android.annotation.Nullable;
 import android.app.KeyguardManager;
@@ -71,10 +71,10 @@
 import com.android.wm.shell.desktopmode.DesktopModeEventLogger;
 import com.android.wm.shell.desktopmode.DesktopModeLoggerTransitionObserver;
 import com.android.wm.shell.desktopmode.DesktopRepository;
+import com.android.wm.shell.desktopmode.DesktopTaskChangeListener;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.desktopmode.DesktopTasksLimiter;
 import com.android.wm.shell.desktopmode.DesktopTasksTransitionObserver;
-import com.android.wm.shell.desktopmode.DesktopTaskChangeListener;
 import com.android.wm.shell.desktopmode.DragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.EnterDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.ExitDesktopTaskTransitionHandler;
@@ -92,9 +92,9 @@
 import com.android.wm.shell.freeform.FreeformTaskListener;
 import com.android.wm.shell.freeform.FreeformTaskTransitionHandler;
 import com.android.wm.shell.freeform.FreeformTaskTransitionObserver;
-import com.android.wm.shell.freeform.TaskChangeListener;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarterInitializer;
+import com.android.wm.shell.freeform.TaskChangeListener;
 import com.android.wm.shell.keyguard.KeyguardTransitionHandler;
 import com.android.wm.shell.onehanded.OneHandedController;
 import com.android.wm.shell.pip.PipTransitionController;
@@ -111,6 +111,7 @@
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.taskview.TaskViewTransitions;
 import com.android.wm.shell.transition.DefaultMixedHandler;
+import com.android.wm.shell.transition.FocusTransitionObserver;
 import com.android.wm.shell.transition.HomeTransitionObserver;
 import com.android.wm.shell.transition.MixedTransitionHandler;
 import com.android.wm.shell.transition.Transitions;
@@ -391,10 +392,11 @@
             Transitions transitions,
             Optional<DesktopFullImmersiveTransitionHandler> desktopImmersiveTransitionHandler,
             WindowDecorViewModel windowDecorViewModel,
-            Optional<TaskChangeListener> taskChangeListener) {
+            Optional<TaskChangeListener> taskChangeListener,
+            FocusTransitionObserver focusTransitionObserver) {
         return new FreeformTaskTransitionObserver(
                 context, shellInit, transitions, desktopImmersiveTransitionHandler,
-                windowDecorViewModel, taskChangeListener);
+                windowDecorViewModel, taskChangeListener, focusTransitionObserver);
     }
 
     @WMSingleton
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
index 2b38cda..b8507e3 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
@@ -68,9 +68,7 @@
     private val idSequence: InstanceIdSequence by lazy { InstanceIdSequence(Int.MAX_VALUE) }
 
     init {
-        if (
-            Transitions.ENABLE_SHELL_TRANSITIONS && DesktopModeStatus.canEnterDesktopMode(context)
-        ) {
+        if (DesktopModeStatus.canEnterDesktopMode(context)) {
             shellInit.addInitCallback(this::onInit, this)
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
index 7b2a5d3..c175133 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
@@ -80,6 +80,14 @@
             freeformTasksInZOrder = ArrayList(freeformTasksInZOrder),
             fullImmersiveTaskId = fullImmersiveTaskId
         )
+        fun clear() {
+            activeTasks.clear()
+            visibleTasks.clear()
+            minimizedTasks.clear()
+            closingTasks.clear()
+            freeformTasksInZOrder.clear()
+            fullImmersiveTaskId = null
+        }
     }
 
     /* Current wallpaper activity token to remove wallpaper activity when last task is removed. */
@@ -414,6 +422,19 @@
     }
 
     /**
+     * Removes the desktop for the given [displayId] and returns the active tasks on that desktop.
+     */
+    fun removeDesktop(displayId: Int): ArraySet<Int> {
+        if (!desktopTaskDataByDisplayId.contains(displayId)) {
+            logW("Could not find desktop to remove: displayId=%d", displayId)
+            return ArraySet()
+        }
+        val activeTasks = ArraySet(desktopTaskDataByDisplayId[displayId].activeTasks)
+        desktopTaskDataByDisplayId[displayId].clear()
+        return activeTasks
+    }
+
+    /**
      * Updates active desktop gesture exclusion regions.
      *
      * If [desktopExclusionRegions] is accepted by [desktopGestureExclusionListener], updates it in
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 5f41326..75c795b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -16,8 +16,8 @@
 
 package com.android.wm.shell.desktopmode
 
-import android.app.ActivityManager.RunningTaskInfo
 import android.app.ActivityManager
+import android.app.ActivityManager.RunningTaskInfo
 import android.app.ActivityOptions
 import android.app.KeyguardManager
 import android.app.PendingIntent
@@ -44,9 +44,14 @@
 import android.view.DragEvent
 import android.view.SurfaceControl
 import android.view.WindowManager.TRANSIT_CHANGE
+import android.view.WindowManager.TRANSIT_CLOSE
 import android.view.WindowManager.TRANSIT_NONE
 import android.view.WindowManager.TRANSIT_OPEN
 import android.view.WindowManager.TRANSIT_TO_FRONT
+import android.window.DesktopModeFlags
+import android.window.DesktopModeFlags.DISABLE_NON_RESIZABLE_APP_SNAP_RESIZE
+import android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
+import android.window.DesktopModeFlags.ENABLE_WINDOWING_DYNAMIC_INITIAL_BOUNDS
 import android.window.RemoteTransition
 import android.window.TransitionInfo
 import android.window.TransitionRequestInfo
@@ -86,10 +91,6 @@
 import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.annotations.ExternalThread
 import com.android.wm.shell.shared.annotations.ShellMainThread
-import android.window.flags.DesktopModeFlags
-import android.window.flags.DesktopModeFlags.DISABLE_NON_RESIZABLE_APP_SNAP_RESIZE
-import android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
-import android.window.flags.DesktopModeFlags.ENABLE_WINDOWING_DYNAMIC_INITIAL_BOUNDS
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus.DESKTOP_DENSITY_OVERRIDE
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus.useDesktopOverrideDensity
@@ -261,17 +262,12 @@
         val wct = WindowContainerTransaction()
         bringDesktopAppsToFront(displayId, wct)
 
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            val transitionType = transitionType(remoteTransition)
-            val handler =
-                remoteTransition?.let {
-                    OneShotRemoteHandler(transitions.mainExecutor, remoteTransition)
-                }
-            transitions.startTransition(transitionType, wct, handler).also { t ->
-                handler?.setTransition(t)
-            }
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
+        val transitionType = transitionType(remoteTransition)
+        val handler = remoteTransition?.let {
+            OneShotRemoteHandler(transitions.mainExecutor, remoteTransition)
+        }
+        transitions.startTransition(transitionType, wct, handler).also { t ->
+            handler?.setTransition(t)
         }
     }
 
@@ -388,12 +384,8 @@
             bringDesktopAppsToFrontBeforeShowingNewTask(task.displayId, wct, task.taskId)
         addMoveToDesktopChanges(wct, task)
 
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            val transition = enterDesktopTaskTransitionHandler.moveToDesktop(wct, transitionSource)
-            addPendingMinimizeTransition(transition, taskToMinimize)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-        }
+        val transition = enterDesktopTaskTransitionHandler.moveToDesktop(wct, transitionSource)
+        addPendingMinimizeTransition(transition, taskToMinimize)
     }
 
     /**
@@ -499,11 +491,9 @@
         // Rather than set windowing mode to multi-window at task level, set it to
         // undefined and inherit from split stage.
         wct.setWindowingMode(task.token, WINDOWING_MODE_UNDEFINED)
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            transitions.startTransition(TRANSIT_CHANGE, wct, null /* handler */)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-        }
+
+        transitions.startTransition(TRANSIT_CHANGE, wct, null /* handler */)
+
     }
 
     private fun exitSplitIfApplicable(wct: WindowContainerTransaction, taskInfo: RunningTaskInfo) {
@@ -537,17 +527,12 @@
         val wct = WindowContainerTransaction()
         addMoveToFullscreenChanges(wct, task)
 
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            exitDesktopTaskTransitionHandler.startTransition(
+        exitDesktopTaskTransitionHandler.startTransition(
                 transitionSource,
                 wct,
                 position,
                 mOnAnimationFinishedCallback
             )
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-            releaseVisualIndicator()
-        }
     }
 
     /** Move a task to the front */
@@ -584,12 +569,9 @@
         wct.reorder(taskInfo.token, true /* onTop */, true /* includingParents */)
         val taskToMinimize =
             addAndGetMinimizeChangesIfNeeded(taskInfo.displayId, wct, taskInfo.taskId)
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            val transition = transitions.startTransition(TRANSIT_TO_FRONT, wct, null /* handler */)
-            addPendingMinimizeTransition(transition, taskToMinimize)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-        }
+
+        val transition = transitions.startTransition(TRANSIT_TO_FRONT, wct, null /* handler */)
+        addPendingMinimizeTransition(transition, taskToMinimize)
     }
 
     /**
@@ -645,11 +627,9 @@
 
         val wct = WindowContainerTransaction()
         wct.reparent(task.token, displayAreaInfo.token, true /* onTop */)
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            transitions.startTransition(TRANSIT_CHANGE, wct, null /* handler */)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-        }
+
+        transitions.startTransition(TRANSIT_CHANGE, wct, null /* handler */)
+
     }
 
     /** Moves a task in/out of full immersive state within the desktop. */
@@ -733,11 +713,9 @@
 
         taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding)
         val wct = WindowContainerTransaction().setBounds(taskInfo.token, destinationBounds)
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            toggleResizeDesktopTaskTransitionHandler.startTransition(wct)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-        }
+
+        toggleResizeDesktopTaskTransitionHandler.startTransition(wct)
+
     }
 
     private fun getMaximizeBounds(taskInfo: RunningTaskInfo, stableBounds: Rect): Rect {
@@ -847,11 +825,9 @@
 
         taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(true)
         val wct = WindowContainerTransaction().setBounds(taskInfo.token, destinationBounds)
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            toggleResizeDesktopTaskTransitionHandler.startTransition(wct, currentDragBounds)
-        } else {
-            shellTaskOrganizer.applyTransaction(wct)
-        }
+
+        toggleResizeDesktopTaskTransitionHandler.startTransition(wct, currentDragBounds)
+
     }
 
     @VisibleForTesting
@@ -1498,6 +1474,22 @@
         }
     }
 
+    fun removeDesktop(displayId: Int) {
+        if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION.isTrue()) return
+
+        val tasksToRemove = taskRepository.removeDesktop(displayId)
+        val wct = WindowContainerTransaction()
+        tasksToRemove.forEach {
+            val task = shellTaskOrganizer.getRunningTaskInfo(it)
+            if (task != null) {
+                wct.removeTask(task.token)
+            } else {
+                recentTasksController?.removeBackgroundTask(it)
+            }
+        }
+        if (!wct.isEmpty) transitions.startTransition(TRANSIT_CLOSE, wct, null)
+    }
+
     /** Enter split by using the focused desktop task in given `displayId`. */
     fun enterSplit(displayId: Int, leftOrTop: Boolean) {
         getFocusedFreeformTask(displayId)?.let { requestSplit(it, leftOrTop) }
@@ -2025,6 +2017,12 @@
                 c.moveTaskToDesktop(taskId, transitionSource = transitionSource)
             }
         }
+
+        override fun removeDesktop(displayId: Int) {
+            executeRemoteCallWithTaskPermission(controller, "removeDesktop") { c ->
+                c.removeDesktop(displayId)
+            }
+        }
     }
 
     private fun logV(msg: String, vararg arguments: Any?) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
index 37bec21..d6b7212 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksLimiter.kt
@@ -24,7 +24,7 @@
 import android.view.WindowManager.TRANSIT_TO_BACK
 import android.window.TransitionInfo
 import android.window.WindowContainerTransaction
-import android.window.flags.DesktopModeFlags
+import android.window.DesktopModeFlags
 import androidx.annotation.VisibleForTesting
 import com.android.internal.jank.Cuj.CUJ_DESKTOP_MODE_MINIMIZE_WINDOW
 import com.android.internal.jank.InteractionJankMonitor
@@ -39,7 +39,7 @@
  * Limits the number of tasks shown in Desktop Mode.
  *
  * This class should only be used if
- * [android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TASK_LIMIT]
+ * [android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_TASK_LIMIT]
  * is enabled and [maxTasksLimit] is strictly greater than 0.
  */
 class DesktopTasksLimiter (
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
index e086e40..a4bc2fe 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
@@ -24,8 +24,8 @@
 import android.view.WindowManager.TRANSIT_TO_BACK
 import android.window.TransitionInfo
 import android.window.WindowContainerTransaction
-import android.window.flags.DesktopModeFlags
-import android.window.flags.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
+import android.window.DesktopModeFlags
+import android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
 import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
@@ -48,9 +48,7 @@
 ) : Transitions.TransitionObserver {
 
     init {
-        if (
-            Transitions.ENABLE_SHELL_TRANSITIONS && DesktopModeStatus.canEnterDesktopMode(context)
-        ) {
+        if (DesktopModeStatus.canEnterDesktopMode(context)) {
             shellInit.addInitCallback(::onInit, this)
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopMode.aidl b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopMode.aidl
index b036e40e..1090a46 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopMode.aidl
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopMode.aidl
@@ -49,4 +49,7 @@
 
     /** Move a task with given `taskId` to desktop */
     void moveToDesktop(int taskId, in DesktopModeTransitionSource transitionSource);
+
+    /** Remove desktop on the given display */
+    void removeDesktop(int displayId);
 }
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index fbd3c10..ae65892 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -24,7 +24,7 @@
 import android.content.Context;
 import android.util.SparseArray;
 import android.view.SurfaceControl;
-import android.window.flags.DesktopModeFlags;
+import android.window.DesktopModeFlags;
 
 import com.android.internal.protolog.ProtoLog;
 import com.android.wm.shell.ShellTaskOrganizer;
@@ -99,11 +99,6 @@
         state.mTaskInfo = taskInfo;
         state.mLeash = leash;
         mTasks.put(taskInfo.taskId, state);
-        if (!Transitions.ENABLE_SHELL_TRANSITIONS) {
-            SurfaceControl.Transaction t = new SurfaceControl.Transaction();
-            mWindowDecorationViewModel.onTaskOpening(taskInfo, leash, t, t);
-            t.apply();
-        }
 
         if (DesktopModeStatus.canEnterDesktopMode(mContext)) {
             mDesktopRepository.ifPresent(repository -> {
@@ -139,9 +134,6 @@
             });
         }
         mWindowDecorationViewModel.onTaskVanished(taskInfo);
-        if (!Transitions.ENABLE_SHELL_TRANSITIONS) {
-            mWindowDecorationViewModel.destroyWindowDecoration(taskInfo);
-        }
         updateLaunchAdjacentController();
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserver.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserver.java
index 056f6b9..4106a10 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserver.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserver.java
@@ -30,6 +30,7 @@
 import com.android.window.flags.Flags;
 import com.android.wm.shell.desktopmode.DesktopFullImmersiveTransitionHandler;
 import com.android.wm.shell.sysui.ShellInit;
+import com.android.wm.shell.transition.FocusTransitionObserver;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.WindowDecorViewModel;
 
@@ -50,6 +51,7 @@
     private final Optional<DesktopFullImmersiveTransitionHandler> mImmersiveTransitionHandler;
     private final WindowDecorViewModel mWindowDecorViewModel;
     private final Optional<TaskChangeListener> mTaskChangeListener;
+    private final FocusTransitionObserver mFocusTransitionObserver;
 
     private final Map<IBinder, List<ActivityManager.RunningTaskInfo>> mTransitionToTaskInfo =
             new HashMap<>();
@@ -60,12 +62,14 @@
             Transitions transitions,
             Optional<DesktopFullImmersiveTransitionHandler> immersiveTransitionHandler,
             WindowDecorViewModel windowDecorViewModel,
-            Optional<TaskChangeListener> taskChangeListener) {
+            Optional<TaskChangeListener> taskChangeListener,
+            FocusTransitionObserver focusTransitionObserver) {
         mTransitions = transitions;
         mImmersiveTransitionHandler = immersiveTransitionHandler;
         mWindowDecorViewModel = windowDecorViewModel;
         mTaskChangeListener = taskChangeListener;
-        if (Transitions.ENABLE_SHELL_TRANSITIONS && FreeformComponents.isFreeformEnabled(context)) {
+        mFocusTransitionObserver = focusTransitionObserver;
+        if (FreeformComponents.isFreeformEnabled(context)) {
             shellInit.addInitCallback(this::onInit, this);
         }
     }
@@ -87,6 +91,9 @@
             //  Otherwise window decoration relayout won't run with the immersive state up to date.
             mImmersiveTransitionHandler.ifPresent(h -> h.onTransitionReady(transition));
         }
+        // Update focus state first to ensure the correct state can be queried from listeners.
+        // TODO(371503964): Remove this once the unified task repository is ready.
+        mFocusTransitionObserver.updateFocusState(info);
 
         final ArrayList<ActivityManager.RunningTaskInfo> taskInfoList = new ArrayList<>();
         final ArrayList<WindowContainerToken> taskParents = new ArrayList<>();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitionHandler.java
index abec3b9..f8d2011 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitionHandler.java
@@ -28,6 +28,8 @@
 import static android.view.WindowManager.TRANSIT_FLAG_KEYGUARD_OCCLUDING;
 import static android.view.WindowManager.TRANSIT_FLAG_KEYGUARD_UNOCCLUDING;
 import static android.view.WindowManager.TRANSIT_SLEEP;
+import static android.view.WindowManager.TRANSIT_TO_BACK;
+import static android.view.WindowManager.TRANSIT_TO_FRONT;
 
 import static com.android.wm.shell.shared.TransitionUtil.isOpeningType;
 
@@ -44,6 +46,7 @@
 import android.view.WindowManager;
 import android.window.IRemoteTransition;
 import android.window.IRemoteTransitionFinishedCallback;
+import android.window.KeyguardState;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
 import android.window.WindowContainerToken;
@@ -388,5 +391,18 @@
             mMainExecutor.execute(() ->
                     mIsLaunchingActivityOverLockscreen = isLaunchingActivityOverLockscreen);
         }
+
+        @Override
+        public void startKeyguardTransition(boolean keyguardShowing, boolean aodShowing) {
+            final WindowContainerTransaction wct = new WindowContainerTransaction();
+            final KeyguardState keyguardState =
+                    new KeyguardState.Builder(android.view.Display.DEFAULT_DISPLAY)
+                            .setKeyguardShowing(keyguardShowing).setAodShowing(aodShowing).build();
+            wct.addKeyguardState(keyguardState);
+            mMainExecutor.execute(() -> {
+                mTransitions.startTransition(keyguardShowing ? TRANSIT_TO_FRONT : TRANSIT_TO_BACK,
+                        wct, KeyguardTransitionHandler.this);
+            });
+        }
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitions.java
index b7245b9..1d349e6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/keyguard/KeyguardTransitions.java
@@ -44,4 +44,11 @@
      * Notify whether keyguard has created a remote animation runner for next app launch.
      */
     default void setLaunchingActivityOverLockscreen(boolean isLaunchingActivityOverLockscreen) {}
+
+    /**
+     * Notifies Shell to start a keyguard transition directly.
+     * @param keyguardShowing whether keyguard is showing or not.
+     * @param aodShowing whether aod is showing or not.
+     */
+    default void startKeyguardTransition(boolean keyguardShowing, boolean aodShowing) {}
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterAnimator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterAnimator.java
new file mode 100644
index 0000000..f40a87c
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterAnimator.java
@@ -0,0 +1,159 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.pip2.animation;
+
+import android.animation.Animator;
+import android.animation.RectEvaluator;
+import android.animation.ValueAnimator;
+import android.content.Context;
+import android.graphics.PointF;
+import android.graphics.Rect;
+import android.view.Surface;
+import android.view.SurfaceControl;
+import android.window.TransitionInfo;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.wm.shell.R;
+import com.android.wm.shell.common.pip.PipUtils;
+import com.android.wm.shell.pip2.PipSurfaceTransactionHelper;
+import com.android.wm.shell.shared.animation.Interpolators;
+
+/**
+ * Animator that handles bounds animations for entering PIP.
+ */
+public class PipEnterAnimator extends ValueAnimator
+        implements ValueAnimator.AnimatorUpdateListener, ValueAnimator.AnimatorListener {
+    @NonNull private final SurfaceControl mLeash;
+    private final SurfaceControl.Transaction mStartTransaction;
+    private final SurfaceControl.Transaction mFinishTransaction;
+
+    // Bounds updated by the evaluator as animator is running.
+    private final Rect mAnimatedRect = new Rect();
+
+    private final RectEvaluator mRectEvaluator;
+    private final Rect mEndBounds = new Rect();
+    @Nullable private final Rect mSourceRectHint;
+    private final @Surface.Rotation int mRotation;
+    @Nullable private Runnable mAnimationStartCallback;
+    @Nullable private Runnable mAnimationEndCallback;
+
+    private final PipSurfaceTransactionHelper.SurfaceControlTransactionFactory
+            mSurfaceControlTransactionFactory;
+
+    // Internal state representing initial transform - cached to avoid recalculation.
+    private final PointF mInitScale = new PointF();
+    private final PointF mInitPos = new PointF();
+    private final Rect mInitCrop = new Rect();
+
+    public PipEnterAnimator(Context context,
+            @NonNull SurfaceControl leash,
+            SurfaceControl.Transaction startTransaction,
+            SurfaceControl.Transaction finishTransaction,
+            @NonNull Rect endBounds,
+            @Nullable Rect sourceRectHint,
+            @Surface.Rotation int rotation) {
+        mLeash = leash;
+        mStartTransaction = startTransaction;
+        mFinishTransaction = finishTransaction;
+        mRectEvaluator = new RectEvaluator(mAnimatedRect);
+        mEndBounds.set(endBounds);
+        mSourceRectHint = sourceRectHint != null ? new Rect(sourceRectHint) : null;
+        mRotation = rotation;
+        mSurfaceControlTransactionFactory =
+                new PipSurfaceTransactionHelper.VsyncSurfaceControlTransactionFactory();
+
+        final int enterAnimationDuration = context.getResources()
+                .getInteger(R.integer.config_pipEnterAnimationDuration);
+        setDuration(enterAnimationDuration);
+        setFloatValues(0f, 1f);
+        setInterpolator(Interpolators.FAST_OUT_SLOW_IN);
+        addListener(this);
+        addUpdateListener(this);
+    }
+
+    public void setAnimationStartCallback(@NonNull Runnable runnable) {
+        mAnimationStartCallback = runnable;
+    }
+
+    public void setAnimationEndCallback(@NonNull Runnable runnable) {
+        mAnimationEndCallback = runnable;
+    }
+
+    @Override
+    public void onAnimationStart(@NonNull Animator animation) {
+        if (mAnimationStartCallback != null) {
+            mAnimationStartCallback.run();
+        }
+        if (mStartTransaction != null) {
+            onEnterAnimationUpdate(mInitScale, mInitPos, mInitCrop,
+                    0f /* fraction */, mStartTransaction);
+            mStartTransaction.apply();
+        }
+    }
+
+    @Override
+    public void onAnimationEnd(@NonNull Animator animation) {
+        if (mAnimationEndCallback != null) {
+            mAnimationEndCallback.run();
+        }
+    }
+
+    @Override
+    public void onAnimationUpdate(@NonNull ValueAnimator animation) {
+        final SurfaceControl.Transaction tx = mSurfaceControlTransactionFactory.getTransaction();
+        final float fraction = getAnimatedFraction();
+        onEnterAnimationUpdate(mInitScale, mInitPos, mInitCrop, fraction, tx);
+        tx.apply();
+    }
+
+    private void onEnterAnimationUpdate(PointF initScale, PointF initPos, Rect initCrop,
+            float fraction, SurfaceControl.Transaction tx) {
+        float scaleX = 1 + (initScale.x - 1) * (1 - fraction);
+        float scaleY = 1 + (initScale.y - 1) * (1 - fraction);
+        tx.setScale(mLeash, scaleX, scaleY);
+
+        float posX = initPos.x + (mEndBounds.left - initPos.x) * fraction;
+        float posY = initPos.y + (mEndBounds.top - initPos.y) * fraction;
+        tx.setPosition(mLeash, posX, posY);
+
+        Rect endCrop = new Rect(mEndBounds);
+        endCrop.offsetTo(0, 0);
+        mRectEvaluator.evaluate(fraction, initCrop, endCrop);
+        tx.setCrop(mLeash, mAnimatedRect);
+    }
+
+    // no-ops
+
+    @Override
+    public void onAnimationCancel(@NonNull Animator animation) {}
+
+    @Override
+    public void onAnimationRepeat(@NonNull Animator animation) {}
+
+    /**
+     * Caches the initial transform relevant values for the bounds enter animation.
+     *
+     * Since enter PiP makes use of a config-at-end transition, initial transform needs to be
+     * calculated differently from generic transitions.
+     * @param pipChange PiP change received as a transition target.
+     */
+    public void setEnterStartState(@NonNull TransitionInfo.Change pipChange) {
+        PipUtils.calcStartTransform(pipChange, mInitScale, mInitPos, mInitCrop);
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterExitAnimator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipExpandAnimator.java
similarity index 77%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterExitAnimator.java
rename to libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipExpandAnimator.java
index 8ebdc96..8fa5aa9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterExitAnimator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipExpandAnimator.java
@@ -19,7 +19,6 @@
 import android.animation.Animator;
 import android.animation.RectEvaluator;
 import android.animation.ValueAnimator;
-import android.annotation.IntDef;
 import android.content.Context;
 import android.graphics.Rect;
 import android.view.Surface;
@@ -30,35 +29,22 @@
 
 import com.android.wm.shell.R;
 import com.android.wm.shell.pip2.PipSurfaceTransactionHelper;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
+import com.android.wm.shell.shared.animation.Interpolators;
 
 /**
- * Animator that handles bounds animations for entering / exiting PIP.
+ * Animator that handles bounds animations for exit-via-expanding PIP.
  */
-public class PipEnterExitAnimator extends ValueAnimator
+public class PipExpandAnimator extends ValueAnimator
         implements ValueAnimator.AnimatorUpdateListener, ValueAnimator.AnimatorListener {
-    @IntDef(prefix = {"BOUNDS_"}, value = {
-            BOUNDS_ENTER,
-            BOUNDS_EXIT
-    })
-
-    @Retention(RetentionPolicy.SOURCE)
-    public @interface BOUNDS {}
-
-    public static final int BOUNDS_ENTER = 0;
-    public static final int BOUNDS_EXIT = 1;
-
-    @NonNull private final SurfaceControl mLeash;
+    @NonNull
+    private final SurfaceControl mLeash;
     private final SurfaceControl.Transaction mStartTransaction;
     private final SurfaceControl.Transaction mFinishTransaction;
-    private final int mEnterExitAnimationDuration;
-    private final @BOUNDS int mDirection;
     private final @Surface.Rotation int mRotation;
 
     // optional callbacks for tracking animation start and end
-    @Nullable private Runnable mAnimationStartCallback;
+    @Nullable
+    private Runnable mAnimationStartCallback;
     @Nullable private Runnable mAnimationEndCallback;
 
     private final Rect mBaseBounds = new Rect();
@@ -78,7 +64,7 @@
     private final RectEvaluator mInsetEvaluator;
     private final PipSurfaceTransactionHelper mPipSurfaceTransactionHelper;
 
-    public PipEnterExitAnimator(Context context,
+    public PipExpandAnimator(Context context,
             @NonNull SurfaceControl leash,
             SurfaceControl.Transaction startTransaction,
             SurfaceControl.Transaction finishTransaction,
@@ -86,7 +72,6 @@
             @NonNull Rect startBounds,
             @NonNull Rect endBounds,
             @Nullable Rect sourceRectHint,
-            @BOUNDS int direction,
             @Surface.Rotation int rotation) {
         mLeash = leash;
         mStartTransaction = startTransaction;
@@ -98,7 +83,6 @@
         mRectEvaluator = new RectEvaluator(mAnimatedRect);
         mInsetEvaluator = new RectEvaluator(new Rect());
         mPipSurfaceTransactionHelper = new PipSurfaceTransactionHelper(context);
-        mDirection = direction;
         mRotation = rotation;
 
         mSourceRectHint = sourceRectHint != null ? new Rect(sourceRectHint) : null;
@@ -113,12 +97,14 @@
 
         mSurfaceControlTransactionFactory =
                 new PipSurfaceTransactionHelper.VsyncSurfaceControlTransactionFactory();
-        mEnterExitAnimationDuration = context.getResources()
+
+        final int enterAnimationDuration = context.getResources()
                 .getInteger(R.integer.config_pipEnterAnimationDuration);
+        setDuration(enterAnimationDuration);
 
         setObjectValues(startBounds, endBounds);
-        setDuration(mEnterExitAnimationDuration);
         setEvaluator(mRectEvaluator);
+        setInterpolator(Interpolators.FAST_OUT_SLOW_IN);
         addListener(this);
         addUpdateListener(this);
     }
@@ -147,9 +133,10 @@
             // finishTransaction might override some state (eg. corner radii) so we want to
             // manually set the state to the end of the animation
             mPipSurfaceTransactionHelper.scaleAndCrop(mFinishTransaction, mLeash, mSourceRectHint,
-                            mBaseBounds, mAnimatedRect, getInsets(1f), isInPipDirection(), 1f)
-                    .round(mFinishTransaction, mLeash, isInPipDirection())
-                    .shadow(mFinishTransaction, mLeash, isInPipDirection());
+                            mBaseBounds, mAnimatedRect, getInsets(1f),
+                            false /* isInPipDirection */, 1f)
+                    .round(mFinishTransaction, mLeash, false /* applyCornerRadius */)
+                    .shadow(mFinishTransaction, mLeash, false /* applyCornerRadius */);
         }
         if (mAnimationEndCallback != null) {
             mAnimationEndCallback.run();
@@ -160,32 +147,22 @@
     public void onAnimationUpdate(@NonNull ValueAnimator animation) {
         final SurfaceControl.Transaction tx = mSurfaceControlTransactionFactory.getTransaction();
         final float fraction = getAnimatedFraction();
-        Rect insets = getInsets(fraction);
 
         // TODO (b/350801661): implement fixed rotation
 
+        Rect insets = getInsets(fraction);
         mPipSurfaceTransactionHelper.scaleAndCrop(tx, mLeash, mSourceRectHint,
-                mBaseBounds, mAnimatedRect, insets, isInPipDirection(), fraction)
-                .round(tx, mLeash, isInPipDirection())
-                .shadow(tx, mLeash, isInPipDirection());
+                        mBaseBounds, mAnimatedRect, insets, false /* isInPipDirection */, fraction)
+                .round(tx, mLeash, false /* applyCornerRadius */)
+                .shadow(tx, mLeash, false /* applyCornerRadius */);
         tx.apply();
     }
-
     private Rect getInsets(float fraction) {
-        Rect startInsets = isInPipDirection() ? mZeroInsets : mSourceRectHintInsets;
-        Rect endInsets = isInPipDirection() ? mSourceRectHintInsets : mZeroInsets;
-
+        final Rect startInsets = mSourceRectHintInsets;
+        final Rect endInsets = mZeroInsets;
         return mInsetEvaluator.evaluate(fraction, startInsets, endInsets);
     }
 
-    private boolean isInPipDirection() {
-        return mDirection == BOUNDS_ENTER;
-    }
-
-    private boolean isOutPipDirection() {
-        return mDirection == BOUNDS_EXIT;
-    }
-
     // no-ops
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
index 62a60fa..b57f51a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
@@ -56,7 +56,8 @@
 import com.android.wm.shell.common.pip.PipUtils;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip2.animation.PipAlphaAnimator;
-import com.android.wm.shell.pip2.animation.PipEnterExitAnimator;
+import com.android.wm.shell.pip2.animation.PipEnterAnimator;
+import com.android.wm.shell.pip2.animation.PipExpandAnimator;
 import com.android.wm.shell.shared.TransitionUtil;
 import com.android.wm.shell.shared.pip.PipContentOverlay;
 import com.android.wm.shell.sysui.ShellInit;
@@ -218,6 +219,7 @@
             @NonNull SurfaceControl.Transaction startTransaction,
             @NonNull SurfaceControl.Transaction finishTransaction,
             @NonNull Transitions.TransitionFinishCallback finishCallback) {
+        mFinishCallback = finishCallback;
         if (transition == mEnterTransition || info.getType() == TRANSIT_PIP) {
             mEnterTransition = null;
             // If we are in swipe PiP to Home transition we are ENTERING_PIP as a jumpcut transition
@@ -258,6 +260,7 @@
         if (isRemovePipTransition(info)) {
             return removePipImmediately(info, startTransaction, finishTransaction, finishCallback);
         }
+        mFinishCallback = null;
         return false;
     }
 
@@ -297,7 +300,6 @@
             mBoundsChangeDuration = BOUNDS_CHANGE_JUMPCUT_DURATION;
         }
 
-        mFinishCallback = finishCallback;
         mPipTransitionState.setState(PipTransitionState.CHANGING_PIP_BOUNDS, extra);
         return true;
     }
@@ -349,7 +351,6 @@
             startTransaction.setMatrix(pipLeash, transformTensor, matrixTmp);
         }
         startTransaction.apply();
-        finishCallback.onTransitionFinished(null /* finishWct */);
         finishInner();
         return true;
     }
@@ -386,14 +387,6 @@
             return false;
         }
 
-        WindowContainerToken pipTaskToken = pipChange.getContainer();
-        if (pipTaskToken == null) {
-            return false;
-        }
-
-        WindowContainerTransaction finishWct = new WindowContainerTransaction();
-        SurfaceControl.Transaction tx = new SurfaceControl.Transaction();
-
         Rect startBounds = pipChange.getStartAbsBounds();
         Rect endBounds = pipChange.getEndAbsBounds();
         SurfaceControl pipLeash = mPipTransitionState.mPinnedTaskLeash;
@@ -405,29 +398,22 @@
             sourceRectHint = pipChange.getTaskInfo().pictureInPictureParams.getSourceRectHint();
         }
 
-        // For opening type transitions, if there is a non-pip change of mode TO_FRONT/OPEN,
+        // For opening type transitions, if there is a change of mode TO_FRONT/OPEN,
         // make sure that change has alpha of 1f, since it's init state might be set to alpha=0f
         // by the Transitions framework to simplify Task opening transitions.
         if (TransitionUtil.isOpeningType(info.getType())) {
             for (TransitionInfo.Change change : info.getChanges()) {
-                if (change.getLeash() == null || change == pipChange) continue;
+                if (change.getLeash() == null) continue;
                 if (change.getMode() == TRANSIT_OPEN || change.getMode() == TRANSIT_TO_FRONT) {
                     startTransaction.setAlpha(change.getLeash(), 1f);
                 }
             }
         }
 
-        PipEnterExitAnimator animator = new PipEnterExitAnimator(mContext, pipLeash,
-                startTransaction, finishTransaction, startBounds, startBounds, endBounds,
-                sourceRectHint, PipEnterExitAnimator.BOUNDS_ENTER, Surface.ROTATION_0);
-
-        tx.addTransactionCommittedListener(mPipScheduler.getMainExecutor(),
-                this::finishInner);
-        finishWct.setBoundsChangeTransaction(pipTaskToken, tx);
-
-        animator.setAnimationEndCallback(() ->
-                finishCallback.onTransitionFinished(finishWct));
-
+        PipEnterAnimator animator = new PipEnterAnimator(mContext, pipLeash,
+                startTransaction, finishTransaction, endBounds, sourceRectHint, Surface.ROTATION_0);
+        animator.setAnimationStartCallback(() -> animator.setEnterStartState(pipChange));
+        animator.setAnimationEndCallback(this::finishInner);
         animator.start();
         return true;
     }
@@ -452,11 +438,8 @@
 
         PipAlphaAnimator animator = new PipAlphaAnimator(mContext, pipLeash, startTransaction,
                 PipAlphaAnimator.FADE_IN);
-        animator.setAnimationEndCallback(() -> {
-            finishCallback.onTransitionFinished(null);
-            // This should update the pip transition state accordingly after we stop playing.
-            finishInner();
-        });
+        // This should update the pip transition state accordingly after we stop playing.
+        animator.setAnimationEndCallback(this::finishInner);
 
         animator.start();
         return true;
@@ -510,9 +493,9 @@
             sourceRectHint = mPipTaskListener.getPictureInPictureParams().getSourceRectHint();
         }
 
-        PipEnterExitAnimator animator = new PipEnterExitAnimator(mContext, pipLeash,
+        PipExpandAnimator animator = new PipExpandAnimator(mContext, pipLeash,
                 startTransaction, finishTransaction, endBounds, startBounds, endBounds,
-                sourceRectHint, PipEnterExitAnimator.BOUNDS_EXIT, Surface.ROTATION_0);
+                sourceRectHint, Surface.ROTATION_0);
 
         animator.setAnimationEndCallback(() -> {
             mPipTransitionState.setState(PipTransitionState.EXITED_PIP);
@@ -631,6 +614,7 @@
     //
 
     private void finishInner() {
+        finishTransition(null /* tx */);
         if (mPipTransitionState.getSwipePipToHomeOverlay() != null) {
             startOverlayFadeoutAnimation();
         } else if (mPipTransitionState.getState() == PipTransitionState.ENTERING_PIP) {
@@ -652,6 +636,7 @@
         }
         if (mFinishCallback != null) {
             mFinishCallback.onTransitionFinished(wct);
+            mFinishCallback = null;
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
index 95cb3df..6086801 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
@@ -38,8 +38,8 @@
 import android.util.Slog;
 import android.util.SparseArray;
 import android.util.SparseIntArray;
+import android.window.DesktopModeFlags;
 import android.window.WindowContainerToken;
-import android.window.flags.DesktopModeFlags;
 
 import androidx.annotation.BinderThread;
 import androidx.annotation.NonNull;
@@ -540,6 +540,14 @@
         return null;
     }
 
+    /**
+     * Remove the background task that match the given taskId. This will remove the task regardless
+     * of whether it's active or recent.
+     */
+    public boolean removeBackgroundTask(int taskId) {
+        return mActivityTaskManager.removeTask(taskId);
+    }
+
     public void dump(@NonNull PrintWriter pw, String prefix) {
         final String innerPrefix = prefix + "  ";
         pw.println(prefix + TAG);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt
index f2c08dc..1af99f9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt
@@ -21,7 +21,7 @@
 import android.util.ArrayMap
 import android.view.SurfaceControl
 import android.window.TransitionInfo
-import android.window.flags.DesktopModeFlags
+import android.window.DesktopModeFlags
 import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
@@ -42,9 +42,7 @@
         ArrayMap<TaskStackTransitionObserverListener, Executor>()
 
     init {
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            shellInit.addInitCallback(::onInit, this)
-        }
+        shellInit.addInitCallback(::onInit, this)
     }
 
     fun onInit() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/FocusTransitionObserver.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/FocusTransitionObserver.java
index 399e39a..6d01e24 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/FocusTransitionObserver.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/FocusTransitionObserver.java
@@ -16,7 +16,8 @@
 
 package com.android.wm.shell.transition;
 
-import static android.view.Display.INVALID_DISPLAY;
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowManager.TRANSIT_OPEN;
 import static android.window.TransitionInfo.FLAG_IS_DISPLAY;
 import static android.window.TransitionInfo.FLAG_MOVED_TO_TOP;
 
@@ -24,10 +25,11 @@
 import static com.android.wm.shell.transition.Transitions.TransitionObserver;
 
 import android.annotation.NonNull;
-import android.os.IBinder;
+import android.app.ActivityManager.RunningTaskInfo;
 import android.os.RemoteException;
+import android.util.ArraySet;
 import android.util.Slog;
-import android.view.SurfaceControl;
+import android.util.SparseArray;
 import android.window.TransitionInfo;
 
 import com.android.wm.shell.shared.FocusTransitionListener;
@@ -43,44 +45,64 @@
  * It reports transitions to callers outside of the process via {@link IFocusTransitionListener},
  * and callers within the process via {@link FocusTransitionListener}.
  */
-public class FocusTransitionObserver implements TransitionObserver {
+public class FocusTransitionObserver {
     private static final String TAG = FocusTransitionObserver.class.getSimpleName();
 
     private IFocusTransitionListener mRemoteListener;
     private final Map<FocusTransitionListener, Executor> mLocalListeners =
             new HashMap<>();
 
-    private int mFocusedDisplayId = INVALID_DISPLAY;
+    private int mFocusedDisplayId = DEFAULT_DISPLAY;
+    private final SparseArray<RunningTaskInfo> mFocusedTaskOnDisplay = new SparseArray<>();
+
+    private final ArraySet<RunningTaskInfo> mTmpTasksToBeNotified = new ArraySet<>();
 
     public FocusTransitionObserver() {}
 
-    @Override
-    public void onTransitionReady(@NonNull IBinder transition,
-            @NonNull TransitionInfo info,
-            @NonNull SurfaceControl.Transaction startTransaction,
-            @NonNull SurfaceControl.Transaction finishTransaction) {
+    /**
+     * Update display/window focus state from the given transition info and notifies changes if any.
+     */
+    public void updateFocusState(@NonNull TransitionInfo info) {
+        if (!enableDisplayFocusInShellTransitions()) {
+            return;
+        }
         final List<TransitionInfo.Change> changes = info.getChanges();
         for (int i = changes.size() - 1; i >= 0; i--) {
             final TransitionInfo.Change change = changes.get(i);
+
+            final RunningTaskInfo task = change.getTaskInfo();
+            if (task != null
+                    && (change.hasFlags(FLAG_MOVED_TO_TOP) || change.getMode() == TRANSIT_OPEN)) {
+                final RunningTaskInfo lastFocusedTaskOnDisplay =
+                        mFocusedTaskOnDisplay.get(task.displayId);
+                if (lastFocusedTaskOnDisplay != null) {
+                    mTmpTasksToBeNotified.add(lastFocusedTaskOnDisplay);
+                }
+                mTmpTasksToBeNotified.add(task);
+                mFocusedTaskOnDisplay.put(task.displayId, task);
+            }
+
             if (change.hasFlags(FLAG_IS_DISPLAY) && change.hasFlags(FLAG_MOVED_TO_TOP)) {
                 if (mFocusedDisplayId != change.getEndDisplayId()) {
+                    final RunningTaskInfo lastGloballyFocusedTask =
+                            mFocusedTaskOnDisplay.get(mFocusedDisplayId);
+                    if (lastGloballyFocusedTask != null) {
+                        mTmpTasksToBeNotified.add(lastGloballyFocusedTask);
+                    }
                     mFocusedDisplayId = change.getEndDisplayId();
                     notifyFocusedDisplayChanged();
+                    final RunningTaskInfo currentGloballyFocusedTask =
+                            mFocusedTaskOnDisplay.get(mFocusedDisplayId);
+                    if (currentGloballyFocusedTask != null) {
+                        mTmpTasksToBeNotified.add(currentGloballyFocusedTask);
+                    }
                 }
-                return;
             }
         }
+        mTmpTasksToBeNotified.forEach(this::notifyTaskFocusChanged);
+        mTmpTasksToBeNotified.clear();
     }
 
-    @Override
-    public void onTransitionStarting(@NonNull IBinder transition) {}
-
-    @Override
-    public void onTransitionMerged(@NonNull IBinder merged, @NonNull IBinder playing) {}
-
-    @Override
-    public void onTransitionFinished(@NonNull IBinder transition, boolean aborted) {}
-
     /**
      * Sets the focus transition listener that receives any transitions resulting in focus switch.
      * This is for calls from outside the Shell, within the host process.
@@ -92,7 +114,10 @@
             return;
         }
         mLocalListeners.put(listener, executor);
-        executor.execute(() -> listener.onFocusedDisplayChanged(mFocusedDisplayId));
+        executor.execute(() -> {
+            listener.onFocusedDisplayChanged(mFocusedDisplayId);
+            mTmpTasksToBeNotified.forEach(this::notifyTaskFocusChanged);
+        });
     }
 
     /**
@@ -120,13 +145,20 @@
         notifyFocusedDisplayChangedToRemote();
     }
 
-    /**
-     * Notifies the listener that display focus has changed.
-     */
-    public void notifyFocusedDisplayChanged() {
+    private void notifyTaskFocusChanged(RunningTaskInfo task) {
+        final boolean isFocusedOnDisplay = isFocusedOnDisplay(task);
+        final boolean isFocusedGlobally = hasGlobalFocus(task);
+        mLocalListeners.forEach((listener, executor) ->
+                executor.execute(() -> listener.onFocusedTaskChanged(task.taskId,
+                        isFocusedOnDisplay, isFocusedGlobally)));
+    }
+
+    private void notifyFocusedDisplayChanged() {
         notifyFocusedDisplayChangedToRemote();
         mLocalListeners.forEach((listener, executor) ->
-                executor.execute(() -> listener.onFocusedDisplayChanged(mFocusedDisplayId)));
+                executor.execute(() -> {
+                    listener.onFocusedDisplayChanged(mFocusedDisplayId);
+                }));
     }
 
     private void notifyFocusedDisplayChangedToRemote() {
@@ -138,4 +170,23 @@
             }
         }
     }
+
+    private boolean isFocusedOnDisplay(@NonNull RunningTaskInfo task) {
+        if (!enableDisplayFocusInShellTransitions()) {
+            return task.isFocused;
+        }
+        final RunningTaskInfo focusedTaskOnDisplay = mFocusedTaskOnDisplay.get(task.displayId);
+        return focusedTaskOnDisplay != null && focusedTaskOnDisplay.taskId == task.taskId;
+    }
+
+    /**
+     * Checks whether the given task has focused globally on the system.
+     * (Note {@link RunningTaskInfo#isFocused} represents per-display focus.)
+     */
+    public boolean hasGlobalFocus(@NonNull RunningTaskInfo task) {
+        if (!enableDisplayFocusInShellTransitions()) {
+            return task.isFocused;
+        }
+        return task.displayId == mFocusedDisplayId && isFocusedOnDisplay(task);
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
index d5e92e6..346f21b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
@@ -392,8 +392,6 @@
 
         mShellCommandHandler.addCommandCallback("transitions", this, this);
         mShellCommandHandler.addDumpCallback(this::dump, this);
-
-        registerObserver(mFocusTransitionObserver);
     }
 
     public boolean isRegistered() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
index 839973f..576c911 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/CaptionWindowDecoration.java
@@ -16,7 +16,7 @@
 
 package com.android.wm.shell.windowdecor;
 
-import static android.window.flags.DesktopModeFlags.ENABLE_WINDOWING_SCALED_RESIZING;
+import static android.window.DesktopModeFlags.ENABLE_WINDOWING_SCALED_RESIZING;
 
 import static com.android.wm.shell.windowdecor.DragResizeWindowGeometry.getFineResizeCornerSize;
 import static com.android.wm.shell.windowdecor.DragResizeWindowGeometry.getLargeResizeCornerSize;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index bcf48d9..e55bc67 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -77,10 +77,10 @@
 import android.view.View;
 import android.view.ViewConfiguration;
 import android.widget.Toast;
+import android.window.DesktopModeFlags;
 import android.window.TaskSnapshot;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
-import android.window.flags.DesktopModeFlags;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.OptIn;
@@ -103,8 +103,8 @@
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler;
-import com.android.wm.shell.desktopmode.DesktopRepository;
 import com.android.wm.shell.desktopmode.DesktopModeVisualIndicator;
+import com.android.wm.shell.desktopmode.DesktopRepository;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition;
 import com.android.wm.shell.desktopmode.DesktopTasksLimiter;
@@ -133,14 +133,14 @@
 import kotlin.Pair;
 import kotlin.Unit;
 
-import kotlinx.coroutines.ExperimentalCoroutinesApi;
-
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Supplier;
 
+import kotlinx.coroutines.ExperimentalCoroutinesApi;
+
 /**
  * View model for the window decoration with a caption and shadows. Works with
  * {@link DesktopModeWindowDecoration}.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index 25d37fc..a78fb9b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -24,8 +24,8 @@
 import static android.view.MotionEvent.ACTION_CANCEL;
 import static android.view.MotionEvent.ACTION_DOWN;
 import static android.view.MotionEvent.ACTION_UP;
-import static android.window.flags.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
-import static android.window.flags.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION_ALWAYS;
+import static android.window.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
+import static android.window.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION_ALWAYS;
 
 import static com.android.launcher3.icons.BaseIconFactory.MODE_DEFAULT;
 import static com.android.wm.shell.shared.desktopmode.DesktopModeStatus.canEnterDesktopMode;
@@ -72,9 +72,9 @@
 import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.widget.ImageButton;
+import android.window.DesktopModeFlags;
 import android.window.TaskSnapshot;
 import android.window.WindowContainerTransaction;
-import android.window.flags.DesktopModeFlags;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.policy.ScreenDecorationsUtils;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragPositioningCallbackUtility.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragPositioningCallbackUtility.java
index 38f9cfa..60c9222 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragPositioningCallbackUtility.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragPositioningCallbackUtility.java
@@ -27,7 +27,7 @@
 import android.graphics.Rect;
 import android.util.DisplayMetrics;
 import android.view.SurfaceControl;
-import android.window.flags.DesktopModeFlags;
+import android.window.DesktopModeFlags;
 
 import androidx.annotation.NonNull;
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragResizeWindowGeometry.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragResizeWindowGeometry.java
index d726f50..33d1c26 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragResizeWindowGeometry.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DragResizeWindowGeometry.java
@@ -18,7 +18,7 @@
 
 import static android.view.InputDevice.SOURCE_MOUSE;
 import static android.view.InputDevice.SOURCE_TOUCHSCREEN;
-import static android.window.flags.DesktopModeFlags.ENABLE_WINDOWING_EDGE_DRAG_RESIZE;
+import static android.window.DesktopModeFlags.ENABLE_WINDOWING_EDGE_DRAG_RESIZE;
 
 import static com.android.wm.shell.windowdecor.DragPositioningCallback.CTRL_TYPE_BOTTOM;
 import static com.android.wm.shell.windowdecor.DragPositioningCallback.CTRL_TYPE_LEFT;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt
index 68a58ee0..376cd2a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeButtonView.kt
@@ -33,7 +33,7 @@
 import androidx.core.animation.doOnStart
 import androidx.core.content.ContextCompat
 import com.android.wm.shell.R
-import android.window.flags.DesktopModeFlags
+import android.window.DesktopModeFlags
 
 private const val OPEN_MAXIMIZE_MENU_DELAY_ON_HOVER_MS = 350
 private const val MAX_DRAWABLE_ALPHA = 255
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
index 52bf400..c2af1d4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/viewholder/AppHeaderViewHolder.kt
@@ -49,7 +49,7 @@
 import com.android.window.flags.Flags
 import com.android.window.flags.Flags.enableMinimizeButton
 import com.android.wm.shell.R
-import android.window.flags.DesktopModeFlags
+import android.window.DesktopModeFlags
 import com.android.wm.shell.windowdecor.MaximizeButtonView
 import com.android.wm.shell.windowdecor.common.DecorThemeUtil
 import com.android.wm.shell.windowdecor.common.OPACITY_100
diff --git a/libs/WindowManager/Shell/tests/flicker/Android.bp b/libs/WindowManager/Shell/tests/flicker/Android.bp
index 58559ac..7b6cfe3 100644
--- a/libs/WindowManager/Shell/tests/flicker/Android.bp
+++ b/libs/WindowManager/Shell/tests/flicker/Android.bp
@@ -30,7 +30,6 @@
 
 java_library {
     name: "wm-shell-flicker-utils",
-    platform_apis: true,
     optimize: {
         enabled: false,
     },
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt
index db65ae1..daf7e7d 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt
@@ -110,16 +110,11 @@
     shellInit = spy(ShellInit(testExecutor))
     desktopModeEventLogger = mock<DesktopModeEventLogger>()
 
-    transitionObserver =
-        DesktopModeLoggerTransitionObserver(
-            context, mockShellInit, transitions, desktopModeEventLogger)
-    if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-      val initRunnableCaptor = ArgumentCaptor.forClass(Runnable::class.java)
-      verify(mockShellInit).addInitCallback(initRunnableCaptor.capture(), same(transitionObserver))
-      initRunnableCaptor.value.run()
-    } else {
-      transitionObserver.onInit()
-    }
+    transitionObserver = DesktopModeLoggerTransitionObserver(
+        context, mockShellInit, transitions, desktopModeEventLogger)
+    val initRunnableCaptor = ArgumentCaptor.forClass(Runnable::class.java)
+    verify(mockShellInit).addInitCallback(initRunnableCaptor.capture(), same(transitionObserver))
+    initRunnableCaptor.value.run()
   }
 
   @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
index 55b9724..1308114 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
@@ -940,6 +940,23 @@
         assertThat(repo.isTaskInFullImmersiveState(taskId = 2)).isTrue()
     }
 
+    @Test
+    fun removeDesktop_multipleTasks_removesAll() {
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 1)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 2)
+        repo.addActiveTask(displayId = DEFAULT_DISPLAY, taskId = 3)
+        // The front-most task will be the one added last through `addOrMoveFreeformTaskToTop`
+        repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 3)
+        repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 2)
+        repo.addOrMoveFreeformTaskToTop(displayId = DEFAULT_DISPLAY, taskId = 1)
+        repo.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = 2)
+
+        val tasksBeforeRemoval = repo.removeDesktop(displayId = DEFAULT_DISPLAY)
+
+        assertThat(tasksBeforeRemoval).containsExactly(1, 2, 3).inOrder()
+        assertThat(repo.getActiveTasks(displayId = DEFAULT_DISPLAY)).isEmpty()
+    }
+
     class TestListener : DesktopRepository.ActiveTasksListener {
         var activeChangesOnDefaultDisplay = 0
         var activeChangesOnSecondaryDisplay = 0
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index ae4772e..27deb0b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -2466,6 +2466,41 @@
   }
 
   @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
+  fun removeDesktop_multipleTasks_removesAll() {
+    val task1 = setUpFreeformTask()
+    val task2 = setUpFreeformTask()
+    val task3 = setUpFreeformTask()
+    taskRepository.minimizeTask(DEFAULT_DISPLAY, task2.taskId)
+
+    controller.removeDesktop(displayId = DEFAULT_DISPLAY)
+
+    val wct = getLatestWct(TRANSIT_CLOSE)
+    assertThat(wct.hierarchyOps).hasSize(3)
+    wct.assertRemoveAt(index = 0, task1.token)
+    wct.assertRemoveAt(index = 1, task2.token)
+    wct.assertRemoveAt(index = 2, task3.token)
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
+  fun removeDesktop_multipleTasksWithBackgroundTask_removesAll() {
+    val task1 = setUpFreeformTask()
+    val task2 = setUpFreeformTask()
+    val task3 = setUpFreeformTask()
+    taskRepository.minimizeTask(DEFAULT_DISPLAY, task2.taskId)
+    whenever(shellTaskOrganizer.getRunningTaskInfo(task3.taskId)).thenReturn(null)
+
+    controller.removeDesktop(displayId = DEFAULT_DISPLAY)
+
+    val wct = getLatestWct(TRANSIT_CLOSE)
+    assertThat(wct.hierarchyOps).hasSize(2)
+    wct.assertRemoveAt(index = 0, task1.token)
+    wct.assertRemoveAt(index = 1, task2.token)
+    verify(recentTasksController).removeBackgroundTask(task3.taskId)
+  }
+
+  @Test
   @EnableFlags(Flags.FLAG_ENABLE_WINDOWING_DYNAMIC_INITIAL_BOUNDS)
   fun dragToDesktop_landscapeDevice_resizable_undefinedOrientation_defaultLandscapeBounds() {
     val spyController = spy(controller)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserverTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserverTest.java
index da95315..145819f3 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserverTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/freeform/FreeformTaskTransitionObserverTest.java
@@ -22,7 +22,6 @@
 import static android.view.WindowManager.TRANSIT_OPEN;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
-import static android.view.WindowManager.TRANSIT_CHANGE;
 
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -44,18 +43,14 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.window.flags.Flags;
-
-import com.android.wm.shell.desktopmode.DesktopTaskChangeListener;
 import com.android.wm.shell.desktopmode.DesktopFullImmersiveTransitionHandler;
 import com.android.wm.shell.sysui.ShellInit;
+import com.android.wm.shell.transition.FocusTransitionObserver;
 import com.android.wm.shell.transition.TransitionInfoBuilder;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.WindowDecorViewModel;
 
-import java.util.Optional;
-
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
@@ -80,6 +75,9 @@
     private WindowDecorViewModel mWindowDecorViewModel;
     @Mock
     private TaskChangeListener mTaskChangeListener;
+    @Mock
+    private FocusTransitionObserver mFocusTransitionObserver;
+
     private FreeformTaskTransitionObserver mTransitionObserver;
 
     @Before
@@ -95,16 +93,13 @@
         mTransitionObserver = new FreeformTaskTransitionObserver(
                 context, mShellInit, mTransitions,
                 Optional.of(mDesktopFullImmersiveTransitionHandler),
-                mWindowDecorViewModel, Optional.of(mTaskChangeListener));
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            final ArgumentCaptor<Runnable> initRunnableCaptor = ArgumentCaptor.forClass(
-                    Runnable.class);
-            verify(mShellInit).addInitCallback(initRunnableCaptor.capture(),
-                    same(mTransitionObserver));
-            initRunnableCaptor.getValue().run();
-        } else {
-            mTransitionObserver.onInit();
-        }
+                mWindowDecorViewModel, Optional.of(mTaskChangeListener), mFocusTransitionObserver);
+
+        final ArgumentCaptor<Runnable> initRunnableCaptor = ArgumentCaptor.forClass(
+                Runnable.class);
+        verify(mShellInit).addInitCallback(initRunnableCaptor.capture(),
+                same(mTransitionObserver));
+        initRunnableCaptor.getValue().run();
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/FocusTransitionObserverTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/FocusTransitionObserverTest.java
index d63158c..015ea20 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/FocusTransitionObserverTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/FocusTransitionObserverTest.java
@@ -23,6 +23,7 @@
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -30,9 +31,6 @@
 import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager.RunningTaskInfo;
-import android.os.Handler;
-import android.os.IBinder;
-import android.os.Looper;
 import android.os.RemoteException;
 import android.platform.test.annotations.RequiresFlagsEnabled;
 import android.platform.test.flag.junit.CheckFlagsRule;
@@ -43,17 +41,11 @@
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
-import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.window.flags.Flags;
-import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
-import com.android.wm.shell.common.DisplayController;
-import com.android.wm.shell.shared.IFocusTransitionListener;
-import com.android.wm.shell.shared.TransactionPool;
-import com.android.wm.shell.sysui.ShellController;
-import com.android.wm.shell.sysui.ShellInit;
+import com.android.wm.shell.shared.FocusTransitionListener;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -75,57 +67,64 @@
 
     @Rule
     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
-    private IFocusTransitionListener mListener;
-    private Transitions mTransition;
+    private FocusTransitionListener mListener;
+    private final TestShellExecutor mShellExecutor = new TestShellExecutor();
     private FocusTransitionObserver mFocusTransitionObserver;
 
     @Before
     public void setUp() {
-        mListener = mock(IFocusTransitionListener.class);
-        when(mListener.asBinder()).thenReturn(mock(IBinder.class));
-
+        mListener = mock(FocusTransitionListener.class);
         mFocusTransitionObserver = new FocusTransitionObserver();
-        mTransition =
-                new Transitions(InstrumentationRegistry.getInstrumentation().getTargetContext(),
-                        mock(ShellInit.class), mock(ShellController.class),
-                        mock(ShellTaskOrganizer.class), mock(TransactionPool.class),
-                        mock(DisplayController.class), new TestShellExecutor(),
-                        new Handler(Looper.getMainLooper()), new TestShellExecutor(),
-                        mock(HomeTransitionObserver.class),
-                mFocusTransitionObserver);
-        mFocusTransitionObserver.setRemoteFocusTransitionListener(mTransition, mListener);
+        mFocusTransitionObserver.setLocalFocusTransitionListener(mListener, mShellExecutor);
+        mShellExecutor.flushAll();
+        clearInvocations(mListener);
     }
 
     @Test
-    public void testOnlyDisplayChangeAffectsDisplayFocus() throws RemoteException {
-        final IBinder binder = mock(IBinder.class);
+    public void testBasicTaskAndDisplayFocusSwitch() throws RemoteException {
         final SurfaceControl.Transaction tx = mock(SurfaceControl.Transaction.class);
 
-        // Open a task on the secondary display, but it doesn't change display focus because it only
-        // has a task change.
+        // First, open a task on the default display.
         TransitionInfo info = mock(TransitionInfo.class);
         final List<TransitionInfo.Change> changes = new ArrayList<>();
-        setupTaskChange(changes, 123 /* taskId */, TRANSIT_OPEN, SECONDARY_DISPLAY_ID,
-                true /* focused */);
+        setupTaskChange(changes, 1 /* taskId */, TRANSIT_OPEN,
+                DEFAULT_DISPLAY, true /* focused */);
         when(info.getChanges()).thenReturn(changes);
-        mFocusTransitionObserver.onTransitionReady(binder, info, tx, tx);
-        verify(mListener, never()).onFocusedDisplayChanged(SECONDARY_DISPLAY_ID);
+        mFocusTransitionObserver.updateFocusState(info);
+        mShellExecutor.flushAll();
+        verify(mListener, never()).onFocusedDisplayChanged(anyInt());
+        verify(mListener, times(1)).onFocusedTaskChanged(1 /* taskId */,
+                true /* isFocusedOnDisplay */, true /* isFocusedGlobally */);
         clearInvocations(mListener);
 
-        // Moving the secondary display to front must change display focus to it.
-        changes.clear();
+        // Open a task on the secondary display.
+        setupTaskChange(changes, 2 /* taskId */, TRANSIT_OPEN,
+                SECONDARY_DISPLAY_ID, true /* focused */);
         setupDisplayToTopChange(changes, SECONDARY_DISPLAY_ID);
         when(info.getChanges()).thenReturn(changes);
-        mFocusTransitionObserver.onTransitionReady(binder, info, tx, tx);
+        mFocusTransitionObserver.updateFocusState(info);
+        mShellExecutor.flushAll();
         verify(mListener, times(1))
                 .onFocusedDisplayChanged(SECONDARY_DISPLAY_ID);
+        verify(mListener, times(1)).onFocusedTaskChanged(1 /* taskId */,
+                true /* isFocusedOnDisplay */, false /* isFocusedGlobally */);
+        verify(mListener, times(1)).onFocusedTaskChanged(2 /* taskId */,
+                true /* isFocusedOnDisplay */, true /* isFocusedGlobally */);
+        clearInvocations(mListener);
 
-        // Moving the secondary display to front must change display focus back to it.
+        // Moving only the default display back to front, and verify that affected tasks are also
+        // notified.
         changes.clear();
         setupDisplayToTopChange(changes, DEFAULT_DISPLAY);
         when(info.getChanges()).thenReturn(changes);
-        mFocusTransitionObserver.onTransitionReady(binder, info, tx, tx);
-        verify(mListener, times(1)).onFocusedDisplayChanged(DEFAULT_DISPLAY);
+        mFocusTransitionObserver.updateFocusState(info);
+        mShellExecutor.flushAll();
+        verify(mListener, times(1))
+                .onFocusedDisplayChanged(DEFAULT_DISPLAY);
+        verify(mListener, times(1)).onFocusedTaskChanged(1 /* taskId */,
+                true /* isFocusedOnDisplay */, true /* isFocusedGlobally */);
+        verify(mListener, times(1)).onFocusedTaskChanged(2 /* taskId */,
+                true /* isFocusedOnDisplay */, false /* isFocusedGlobally */);
     }
 
     private void setupTaskChange(List<TransitionInfo.Change> changes, int taskId,
diff --git a/libs/androidfw/Idmap.cpp b/libs/androidfw/Idmap.cpp
index f066e46..3ecd82b 100644
--- a/libs/androidfw/Idmap.cpp
+++ b/libs/androidfw/Idmap.cpp
@@ -65,13 +65,7 @@
   uint32_t string_pool_index_offset;
 };
 
-struct Idmap_target_entry {
-  uint32_t target_id;
-  uint32_t overlay_id;
-};
-
 struct Idmap_target_entry_inline {
-  uint32_t target_id;
   uint32_t start_value_index;
   uint32_t value_count;
 };
@@ -81,10 +75,9 @@
   Res_value value;
 };
 
-struct Idmap_overlay_entry {
-  uint32_t overlay_id;
-  uint32_t target_id;
-};
+static constexpr uint32_t convert_dev_target_id(uint32_t dev_target_id) {
+  return (0x00FFFFFFU & dtohl(dev_target_id));
+}
 
 OverlayStringPool::OverlayStringPool(const LoadedIdmap* loaded_idmap)
     : data_header_(loaded_idmap->data_header_),
@@ -117,27 +110,29 @@
 }
 
 OverlayDynamicRefTable::OverlayDynamicRefTable(const Idmap_data_header* data_header,
-                                               const Idmap_overlay_entry* entries,
+                                               Idmap_overlay_entries entries,
                                                uint8_t target_assigned_package_id)
     : data_header_(data_header),
       entries_(entries),
-      target_assigned_package_id_(target_assigned_package_id) {}
+      target_assigned_package_id_(target_assigned_package_id) {
+}
 
 status_t OverlayDynamicRefTable::lookupResourceId(uint32_t* resId) const {
-  const Idmap_overlay_entry* first_entry = entries_;
-  const Idmap_overlay_entry* end_entry = entries_ + dtohl(data_header_->overlay_entry_count);
-  auto entry = std::lower_bound(first_entry, end_entry, *resId,
-                                [](const Idmap_overlay_entry& e1, const uint32_t overlay_id) {
-    return dtohl(e1.overlay_id) < overlay_id;
-  });
+  const auto count = dtohl(data_header_->overlay_entry_count);
+  const auto overlay_it_end = entries_.overlay_id + count;
+  const auto entry_it = std::lower_bound(entries_.overlay_id, overlay_it_end, *resId,
+                                         [](uint32_t dev_overlay_id, uint32_t overlay_id) {
+                                           return dtohl(dev_overlay_id) < overlay_id;
+                                         });
 
-  if (entry == end_entry || dtohl(entry->overlay_id) != *resId) {
+  if (entry_it == overlay_it_end || dtohl(*entry_it) != *resId) {
     // A mapping for the target resource id could not be found.
     return DynamicRefTable::lookupResourceId(resId);
   }
 
-  *resId = (0x00FFFFFFU & dtohl(entry->target_id))
-      | (((uint32_t) target_assigned_package_id_) << 24U);
+  const auto index = entry_it - entries_.overlay_id;
+  *resId = convert_dev_target_id(entries_.target_id[index]) |
+           (((uint32_t)target_assigned_package_id_) << 24U);
   return NO_ERROR;
 }
 
@@ -145,12 +140,10 @@
   return DynamicRefTable::lookupResourceId(resId);
 }
 
-IdmapResMap::IdmapResMap(const Idmap_data_header* data_header,
-                         const Idmap_target_entry* entries,
-                         const Idmap_target_entry_inline* inline_entries,
+IdmapResMap::IdmapResMap(const Idmap_data_header* data_header, Idmap_target_entries entries,
+                         Idmap_target_inline_entries inline_entries,
                          const Idmap_target_entry_inline_value* inline_entry_values,
-                         const ConfigDescription* configs,
-                         uint8_t target_assigned_package_id,
+                         const ConfigDescription* configs, uint8_t target_assigned_package_id,
                          const OverlayDynamicRefTable* overlay_ref_table)
     : data_header_(data_header),
       entries_(entries),
@@ -158,7 +151,8 @@
       inline_entry_values_(inline_entry_values),
       configurations_(configs),
       target_assigned_package_id_(target_assigned_package_id),
-      overlay_ref_table_(overlay_ref_table) { }
+      overlay_ref_table_(overlay_ref_table) {
+}
 
 IdmapResMap::Result IdmapResMap::Lookup(uint32_t target_res_id) const {
   if ((target_res_id >> 24U) != target_assigned_package_id_) {
@@ -171,15 +165,15 @@
   target_res_id &= 0x00FFFFFFU;
 
   // Check if the target resource is mapped to an overlay resource.
-  auto first_entry = entries_;
-  auto end_entry = entries_ + dtohl(data_header_->target_entry_count);
-  auto entry = std::lower_bound(first_entry, end_entry, target_res_id,
-                                [](const Idmap_target_entry& e, const uint32_t target_id) {
-    return (0x00FFFFFFU & dtohl(e.target_id)) < target_id;
-  });
+  const auto target_end = entries_.target_id + dtohl(data_header_->target_entry_count);
+  auto target_it = std::lower_bound(entries_.target_id, target_end, target_res_id,
+                                    [](uint32_t dev_target_id, uint32_t target_id) {
+                                      return convert_dev_target_id(dev_target_id) < target_id;
+                                    });
 
-  if (entry != end_entry && (0x00FFFFFFU & dtohl(entry->target_id)) == target_res_id) {
-    uint32_t overlay_resource_id = dtohl(entry->overlay_id);
+  if (target_it != target_end && convert_dev_target_id(*target_it) == target_res_id) {
+    const auto index = target_it - entries_.target_id;
+    uint32_t overlay_resource_id = dtohl(entries_.overlay_id[index]);
     // Lookup the resource without rewriting the overlay resource id back to the target resource id
     // being looked up.
     overlay_ref_table_->lookupResourceIdNoRewrite(&overlay_resource_id);
@@ -187,20 +181,22 @@
   }
 
   // Check if the target resources is mapped to an inline table entry.
-  auto first_inline_entry = inline_entries_;
-  auto end_inline_entry = inline_entries_ + dtohl(data_header_->target_inline_entry_count);
-  auto inline_entry = std::lower_bound(first_inline_entry, end_inline_entry, target_res_id,
-                                       [](const Idmap_target_entry_inline& e,
-                                          const uint32_t target_id) {
-    return (0x00FFFFFFU & dtohl(e.target_id)) < target_id;
-  });
+  const auto inline_entry_target_end =
+      inline_entries_.target_id + dtohl(data_header_->target_inline_entry_count);
+  const auto inline_entry_target_it =
+      std::lower_bound(inline_entries_.target_id, inline_entry_target_end, target_res_id,
+                       [](uint32_t dev_target_id, uint32_t target_id) {
+                         return convert_dev_target_id(dev_target_id) < target_id;
+                       });
 
-  if (inline_entry != end_inline_entry &&
-      (0x00FFFFFFU & dtohl(inline_entry->target_id)) == target_res_id) {
+  if (inline_entry_target_it != inline_entry_target_end &&
+      convert_dev_target_id(*inline_entry_target_it) == target_res_id) {
+    const auto index = inline_entry_target_it - inline_entries_.target_id;
     std::map<ConfigDescription, Res_value> values_map;
-    for (int i = 0; i < inline_entry->value_count; i++) {
-      const auto& value = inline_entry_values_[inline_entry->start_value_index + i];
-      const auto& config = configurations_[value.config_index];
+    const auto& inline_entry = inline_entries_.entry[index];
+    for (int i = 0; i < dtohl(inline_entry.value_count); i++) {
+      const auto& value = inline_entry_values_[dtohl(inline_entry.start_value_index) + i];
+      const auto& config = configurations_[dtohl(value.config_index)];
       values_map[config] = value.value;
     }
     return Result(std::move(values_map));
@@ -210,15 +206,15 @@
 
 namespace {
 template <typename T>
-const T* ReadType(const uint8_t** in_out_data_ptr, size_t* in_out_size, const std::string& label,
+const T* ReadType(const uint8_t** in_out_data_ptr, size_t* in_out_size, const char* label,
                   size_t count = 1) {
   if (!util::IsFourByteAligned(*in_out_data_ptr)) {
-    LOG(ERROR) << "Idmap " << label << " is not word aligned.";
+    LOG(ERROR) << "Idmap " << label << " in " << __func__ << " is not word aligned.";
     return {};
   }
   if ((*in_out_size / sizeof(T)) < count) {
-    LOG(ERROR) << "Idmap too small for the number of " << label << " entries ("
-               << count << ").";
+    LOG(ERROR) << "Idmap too small for the number of " << label << " in " << __func__
+               << " entries (" << count << ").";
     return nullptr;
   }
   auto data_ptr = *in_out_data_ptr;
@@ -229,8 +225,8 @@
 }
 
 std::optional<std::string_view> ReadString(const uint8_t** in_out_data_ptr, size_t* in_out_size,
-                                           const std::string& label) {
-  const auto* len = ReadType<uint32_t>(in_out_data_ptr, in_out_size, label + " length");
+                                           const char* label) {
+  const auto* len = ReadType<uint32_t>(in_out_data_ptr, in_out_size, label);
   if (len == nullptr) {
     return {};
   }
@@ -242,7 +238,7 @@
   const uint32_t padding_size = (4U - ((size_t)*in_out_data_ptr & 0x3U)) % 4U;
   for (uint32_t i = 0; i < padding_size; i++) {
     if (**in_out_data_ptr != 0) {
-      LOG(ERROR) << " Idmap padding of " << label << " is non-zero.";
+      LOG(ERROR) << " Idmap padding of " << label << " in " << __func__ << " is non-zero.";
       return {};
     }
     *in_out_data_ptr += sizeof(uint8_t);
@@ -258,12 +254,10 @@
 #endif
 
 LoadedIdmap::LoadedIdmap(const std::string& idmap_path, const Idmap_header* header,
-                         const Idmap_data_header* data_header,
-                         const Idmap_target_entry* target_entries,
-                         const Idmap_target_entry_inline* target_inline_entries,
+                         const Idmap_data_header* data_header, Idmap_target_entries target_entries,
+                         Idmap_target_inline_entries target_inline_entries,
                          const Idmap_target_entry_inline_value* inline_entry_values,
-                         const ConfigDescription* configs,
-                         const Idmap_overlay_entry* overlay_entries,
+                         const ConfigDescription* configs, Idmap_overlay_entries overlay_entries,
                          std::unique_ptr<ResStringPool>&& string_pool,
                          std::string_view overlay_apk_path, std::string_view target_apk_path)
     : header_(header),
@@ -274,10 +268,12 @@
       configurations_(configs),
       overlay_entries_(overlay_entries),
       string_pool_(std::move(string_pool)),
-      idmap_fd_(android::base::utf8::open(idmap_path.c_str(), O_RDONLY|O_CLOEXEC|O_BINARY|O_PATH)),
+      idmap_fd_(
+          android::base::utf8::open(idmap_path.c_str(), O_RDONLY | O_CLOEXEC | O_BINARY | O_PATH)),
       overlay_apk_path_(overlay_apk_path),
       target_apk_path_(target_apk_path),
-      idmap_last_mod_time_(getFileModDate(idmap_fd_.get())) {}
+      idmap_last_mod_time_(getFileModDate(idmap_fd_.get())) {
+}
 
 std::unique_ptr<LoadedIdmap> LoadedIdmap::Load(StringPiece idmap_path, StringPiece idmap_data) {
   ATRACE_CALL();
@@ -319,14 +315,21 @@
   if (data_header == nullptr) {
     return {};
   }
-  auto target_entries = ReadType<Idmap_target_entry>(&data_ptr, &data_size, "target",
-                                                     dtohl(data_header->target_entry_count));
-  if (target_entries == nullptr) {
+  Idmap_target_entries target_entries{
+      .target_id = ReadType<uint32_t>(&data_ptr, &data_size, "entries.target_id",
+                                      dtohl(data_header->target_entry_count)),
+      .overlay_id = ReadType<uint32_t>(&data_ptr, &data_size, "entries.overlay_id",
+                                       dtohl(data_header->target_entry_count)),
+  };
+  if (!target_entries.target_id || !target_entries.overlay_id) {
     return {};
   }
-  auto target_inline_entries = ReadType<Idmap_target_entry_inline>(
-      &data_ptr, &data_size, "target inline", dtohl(data_header->target_inline_entry_count));
-  if (target_inline_entries == nullptr) {
+  Idmap_target_inline_entries target_inline_entries{
+      .target_id = ReadType<uint32_t>(&data_ptr, &data_size, "target inline.target_id",
+                                      dtohl(data_header->target_inline_entry_count)),
+      .entry = ReadType<Idmap_target_entry_inline>(&data_ptr, &data_size, "target inline.entry",
+                                                   dtohl(data_header->target_inline_entry_count))};
+  if (!target_inline_entries.target_id || !target_inline_entries.entry) {
     return {};
   }
 
@@ -344,9 +347,13 @@
     return {};
   }
 
-  auto overlay_entries = ReadType<Idmap_overlay_entry>(&data_ptr, &data_size, "target inline",
-                                                       dtohl(data_header->overlay_entry_count));
-  if (overlay_entries == nullptr) {
+  Idmap_overlay_entries overlay_entries{
+      .overlay_id = ReadType<uint32_t>(&data_ptr, &data_size, "overlay entries.overlay_id",
+                                       dtohl(data_header->overlay_entry_count)),
+      .target_id = ReadType<uint32_t>(&data_ptr, &data_size, "overlay entries.target_id",
+                                      dtohl(data_header->overlay_entry_count)),
+  };
+  if (!overlay_entries.overlay_id || !overlay_entries.target_id) {
     return {};
   }
   std::optional<std::string_view> string_pool = ReadString(&data_ptr, &data_size, "string pool");
diff --git a/libs/androidfw/include/androidfw/Idmap.h b/libs/androidfw/include/androidfw/Idmap.h
index 64b1f0c..e213fbd 100644
--- a/libs/androidfw/include/androidfw/Idmap.h
+++ b/libs/androidfw/include/androidfw/Idmap.h
@@ -40,6 +40,19 @@
 struct Idmap_target_entry_inline_value;
 struct Idmap_overlay_entry;
 
+struct Idmap_target_entries {
+  const uint32_t* target_id = nullptr;
+  const uint32_t* overlay_id = nullptr;
+};
+struct Idmap_target_inline_entries {
+  const uint32_t* target_id = nullptr;
+  const Idmap_target_entry_inline* entry = nullptr;
+};
+struct Idmap_overlay_entries {
+  const uint32_t* overlay_id = nullptr;
+  const uint32_t* target_id = nullptr;
+};
+
 // A string pool for overlay apk assets. The string pool holds the strings of the overlay resources
 // table and additionally allows for loading strings from the idmap string pool. The idmap string
 // pool strings are offset after the end of the overlay resource table string pool entries so
@@ -67,7 +80,7 @@
 
  private:
   explicit OverlayDynamicRefTable(const Idmap_data_header* data_header,
-                                  const Idmap_overlay_entry* entries,
+                                  Idmap_overlay_entries entries,
                                   uint8_t target_assigned_package_id);
 
   // Rewrites a compile-time overlay resource id to the runtime resource id of corresponding target
@@ -75,8 +88,8 @@
   status_t lookupResourceIdNoRewrite(uint32_t* resId) const;
 
   const Idmap_data_header* data_header_;
-  const Idmap_overlay_entry* entries_;
-  const int8_t target_assigned_package_id_;
+  Idmap_overlay_entries entries_;
+  uint8_t target_assigned_package_id_;
 
   friend LoadedIdmap;
   friend IdmapResMap;
@@ -131,17 +144,15 @@
   }
 
  private:
-  explicit IdmapResMap(const Idmap_data_header* data_header,
-                       const Idmap_target_entry* entries,
-                       const Idmap_target_entry_inline* inline_entries,
+  explicit IdmapResMap(const Idmap_data_header* data_header, Idmap_target_entries entries,
+                       Idmap_target_inline_entries inline_entries,
                        const Idmap_target_entry_inline_value* inline_entry_values,
-                       const ConfigDescription* configs,
-                       uint8_t target_assigned_package_id,
+                       const ConfigDescription* configs, uint8_t target_assigned_package_id,
                        const OverlayDynamicRefTable* overlay_ref_table);
 
   const Idmap_data_header* data_header_;
-  const Idmap_target_entry* entries_;
-  const Idmap_target_entry_inline* inline_entries_;
+  Idmap_target_entries entries_;
+  Idmap_target_inline_entries inline_entries_;
   const Idmap_target_entry_inline_value* inline_entry_values_;
   const ConfigDescription* configurations_;
   const uint8_t target_assigned_package_id_;
@@ -192,11 +203,11 @@
 
   const Idmap_header* header_;
   const Idmap_data_header* data_header_;
-  const Idmap_target_entry* target_entries_;
-  const Idmap_target_entry_inline* target_inline_entries_;
+  Idmap_target_entries target_entries_;
+  Idmap_target_inline_entries target_inline_entries_;
   const Idmap_target_entry_inline_value* inline_entry_values_;
   const ConfigDescription* configurations_;
-  const Idmap_overlay_entry* overlay_entries_;
+  const Idmap_overlay_entries overlay_entries_;
   const std::unique_ptr<ResStringPool> string_pool_;
 
   android::base::unique_fd idmap_fd_;
@@ -207,17 +218,13 @@
  private:
   DISALLOW_COPY_AND_ASSIGN(LoadedIdmap);
 
-  explicit LoadedIdmap(const std::string& idmap_path,
-                       const Idmap_header* header,
-                       const Idmap_data_header* data_header,
-                       const Idmap_target_entry* target_entries,
-                       const Idmap_target_entry_inline* target_inline_entries,
+  explicit LoadedIdmap(const std::string& idmap_path, const Idmap_header* header,
+                       const Idmap_data_header* data_header, Idmap_target_entries target_entries,
+                       Idmap_target_inline_entries target_inline_entries,
                        const Idmap_target_entry_inline_value* inline_entry_values_,
-                       const ConfigDescription* configs,
-                       const Idmap_overlay_entry* overlay_entries,
+                       const ConfigDescription* configs, Idmap_overlay_entries overlay_entries,
                        std::unique_ptr<ResStringPool>&& string_pool,
-                       std::string_view overlay_apk_path,
-                       std::string_view target_apk_path);
+                       std::string_view overlay_apk_path, std::string_view target_apk_path);
 
   friend OverlayStringPool;
 };
diff --git a/libs/androidfw/include/androidfw/ResourceTypes.h b/libs/androidfw/include/androidfw/ResourceTypes.h
index c264890..e330410 100644
--- a/libs/androidfw/include/androidfw/ResourceTypes.h
+++ b/libs/androidfw/include/androidfw/ResourceTypes.h
@@ -48,7 +48,7 @@
 namespace android {
 
 constexpr const uint32_t kIdmapMagic = 0x504D4449u;
-constexpr const uint32_t kIdmapCurrentVersion = 0x00000009u;
+constexpr const uint32_t kIdmapCurrentVersion = 0x0000000Au;
 
 // This must never change.
 constexpr const uint32_t kFabricatedOverlayMagic = 0x4f525246; // FRRO (big endian)
diff --git a/libs/androidfw/tests/data/overlay/overlay.idmap b/libs/androidfw/tests/data/overlay/overlay.idmap
index 8e847e8..7e4b261 100644
--- a/libs/androidfw/tests/data/overlay/overlay.idmap
+++ b/libs/androidfw/tests/data/overlay/overlay.idmap
Binary files differ
diff --git a/libs/hwui/pipeline/skia/BackdropFilterDrawable.cpp b/libs/hwui/pipeline/skia/BackdropFilterDrawable.cpp
index e81cbfb..c0ef4b14 100644
--- a/libs/hwui/pipeline/skia/BackdropFilterDrawable.cpp
+++ b/libs/hwui/pipeline/skia/BackdropFilterDrawable.cpp
@@ -29,37 +29,6 @@
 namespace uirenderer {
 namespace skiapipeline {
 
-BackdropFilterDrawable::~BackdropFilterDrawable() {}
-
-bool BackdropFilterDrawable::prepareToDraw(SkCanvas* canvas, const RenderProperties& properties,
-                                           int backdropImageWidth, int backdropImageHeight) {
-    // the drawing bounds for blurred content.
-    mDstBounds.setWH(properties.getWidth(), properties.getHeight());
-
-    float alphaMultiplier = 1.0f;
-    RenderNodeDrawable::setViewProperties(properties, canvas, &alphaMultiplier, true);
-
-    // get proper subset for previous content.
-    canvas->getTotalMatrix().mapRect(&mImageSubset, mDstBounds);
-    SkRect imageSubset(mImageSubset);
-    // ensure the subset is inside bounds of previous content.
-    if (!mImageSubset.intersect(SkRect::MakeWH(backdropImageWidth, backdropImageHeight))) {
-        return false;
-    }
-
-    // correct the drawing bounds if subset was changed.
-    if (mImageSubset != imageSubset) {
-        SkMatrix inverse;
-        if (canvas->getTotalMatrix().invert(&inverse)) {
-            inverse.mapRect(&mDstBounds, mImageSubset);
-        }
-    }
-
-    // follow the alpha from the target RenderNode.
-    mPaint.setAlpha(properties.layerProperties().alpha() * alphaMultiplier);
-    return true;
-}
-
 void BackdropFilterDrawable::onDraw(SkCanvas* canvas) {
     const RenderProperties& properties = mTargetRenderNode->properties();
     auto* backdropFilter = properties.layerProperties().getBackdropImageFilter();
@@ -68,27 +37,43 @@
         return;
     }
 
-    auto backdropImage = surface->makeImageSnapshot();
-    // sync necessary properties from target RenderNode.
-    if (!prepareToDraw(canvas, properties, backdropImage->width(), backdropImage->height())) {
+    SkRect srcBounds = SkRect::MakeWH(properties.getWidth(), properties.getHeight());
+
+    float alphaMultiplier = 1.0f;
+    RenderNodeDrawable::setViewProperties(properties, canvas, &alphaMultiplier, true);
+    SkPaint paint;
+    paint.setAlpha(properties.layerProperties().alpha() * alphaMultiplier);
+
+    SkRect surfaceSubset;
+    canvas->getTotalMatrix().mapRect(&surfaceSubset, srcBounds);
+    if (!surfaceSubset.intersect(SkRect::MakeWH(surface->width(), surface->height()))) {
         return;
     }
 
-    auto imageSubset = mImageSubset.roundOut();
+    auto backdropImage = surface->makeImageSnapshot(surfaceSubset.roundOut());
+
+    SkIRect imageBounds = SkIRect::MakeWH(backdropImage->width(), backdropImage->height());
+    SkIPoint offset;
+    SkIRect imageSubset;
+
 #ifdef __ANDROID__
     if (canvas->recordingContext()) {
         backdropImage =
                 SkImages::MakeWithFilter(canvas->recordingContext(), backdropImage, backdropFilter,
-                                         imageSubset, imageSubset, &mOutSubset, &mOutOffset);
+                                         imageBounds, imageBounds, &imageSubset, &offset);
     } else
 #endif
     {
-        backdropImage = SkImages::MakeWithFilter(backdropImage, backdropFilter, imageSubset,
-                                                 imageSubset, &mOutSubset, &mOutOffset);
+        backdropImage = SkImages::MakeWithFilter(backdropImage, backdropFilter, imageBounds,
+                                                 imageBounds, &imageSubset, &offset);
     }
-    canvas->drawImageRect(backdropImage, SkRect::Make(mOutSubset), mDstBounds,
-                          SkSamplingOptions(SkFilterMode::kLinear), &mPaint,
-                          SkCanvas::kStrict_SrcRectConstraint);
+
+    canvas->save();
+    canvas->resetMatrix();
+    canvas->drawImageRect(backdropImage, SkRect::Make(imageSubset), surfaceSubset,
+                          SkSamplingOptions(SkFilterMode::kLinear), &paint,
+                          SkCanvas::kFast_SrcRectConstraint);
+    canvas->restore();
 }
 
 }  // namespace skiapipeline
diff --git a/libs/hwui/pipeline/skia/BackdropFilterDrawable.h b/libs/hwui/pipeline/skia/BackdropFilterDrawable.h
index 9e35837..5e216a1 100644
--- a/libs/hwui/pipeline/skia/BackdropFilterDrawable.h
+++ b/libs/hwui/pipeline/skia/BackdropFilterDrawable.h
@@ -37,23 +37,10 @@
     BackdropFilterDrawable(RenderNode* renderNode, SkCanvas* canvas)
             : mTargetRenderNode(renderNode), mBounds(canvas->getLocalClipBounds()) {}
 
-    ~BackdropFilterDrawable();
+    ~BackdropFilterDrawable() = default;
 
 private:
     RenderNode* mTargetRenderNode;
-    SkPaint mPaint;
-
-    SkRect mDstBounds;
-    SkRect mImageSubset;
-    SkIRect mOutSubset;
-    SkIPoint mOutOffset;
-
-    /**
-     * Check all necessary properties before actual drawing.
-     * Return true if ready to draw.
-     */
-    bool prepareToDraw(SkCanvas* canvas, const RenderProperties& properties, int backdropImageWidth,
-                       int backdropImageHeight);
 
 protected:
     void onDraw(SkCanvas* canvas) override;
diff --git a/libs/hwui/tests/common/scenes/BackdropBlur.cpp b/libs/hwui/tests/common/scenes/BackdropBlur.cpp
new file mode 100644
index 0000000..a1133ff
--- /dev/null
+++ b/libs/hwui/tests/common/scenes/BackdropBlur.cpp
@@ -0,0 +1,67 @@
+/*
+ * 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.
+ */
+
+#include <SkBlendMode.h>
+
+#include "SkImageFilter.h"
+#include "SkImageFilters.h"
+#include "TestSceneBase.h"
+#include "utils/Blur.h"
+
+class BackdropBlurAnimation : public TestScene {
+private:
+    std::unique_ptr<TestScene> listView;
+
+public:
+    explicit BackdropBlurAnimation(const TestScene::Options& opts) {
+        listView.reset(TestScene::testMap()["listview"].createScene(opts));
+    }
+
+    void createContent(int width, int height, Canvas& canvas) override {
+        sp<RenderNode> list = TestUtils::createNode(
+                0, 0, width, height,
+                [this, width, height](RenderProperties& props, Canvas& canvas) {
+                    props.setClipToBounds(false);
+                    listView->createContent(width, height, canvas);
+                });
+
+        canvas.drawRenderNode(list.get());
+
+        int x = width / 8;
+        int y = height / 4;
+        sp<RenderNode> blurNode = TestUtils::createNode(
+                x, y, width - x, height - y, [](RenderProperties& props, Canvas& canvas) {
+                    props.mutableOutline().setRoundRect(0, 0, props.getWidth(), props.getHeight(),
+                                                        dp(16), 1);
+                    props.mutableOutline().setShouldClip(true);
+                    sk_sp<SkImageFilter> blurFilter = SkImageFilters::Blur(
+                            Blur::convertRadiusToSigma(dp(8)), Blur::convertRadiusToSigma(dp(8)),
+                            SkTileMode::kClamp, nullptr, nullptr);
+                    props.mutateLayerProperties().setBackdropImageFilter(blurFilter.get());
+                    canvas.drawColor(0x33000000, SkBlendMode::kSrcOver);
+                });
+
+        canvas.drawRenderNode(blurNode.get());
+    }
+
+    void doFrame(int frameNr) override { listView->doFrame(frameNr); }
+};
+
+static TestScene::Registrar _BackdropBlur(TestScene::Info{
+        "backdropblur", "A rounded rect that does a blur-behind of a sky animation.",
+        [](const TestScene::Options& opts) -> test::TestScene* {
+            return new BackdropBlurAnimation(opts);
+        }});
diff --git a/libs/hwui/tests/unit/RenderNodeDrawableTests.cpp b/libs/hwui/tests/unit/RenderNodeDrawableTests.cpp
index ca54087..4b29100 100644
--- a/libs/hwui/tests/unit/RenderNodeDrawableTests.cpp
+++ b/libs/hwui/tests/unit/RenderNodeDrawableTests.cpp
@@ -1280,7 +1280,7 @@
     canvas->drawDrawable(&backdropDrawable);
     // the drawable is still visible, ok to draw.
     EXPECT_EQ(2, canvas->mDrawCounter);
-    EXPECT_EQ(SkRect::MakeLTRB(0, 0, CANVAS_WIDTH - 30, CANVAS_HEIGHT - 30), canvas->mDstBounds);
+    EXPECT_EQ(SkRect::MakeLTRB(30, 30, CANVAS_WIDTH, CANVAS_HEIGHT), canvas->mDstBounds);
 
     canvas->translate(CANVAS_WIDTH, CANVAS_HEIGHT);
     canvas->drawDrawable(&drawable);
diff --git a/location/lib/Android.bp b/location/lib/Android.bp
index b10019a..67d5774 100644
--- a/location/lib/Android.bp
+++ b/location/lib/Android.bp
@@ -29,6 +29,10 @@
     libs: [
         "androidx.annotation_annotation",
     ],
+    stub_only_libs: [
+        // Needed for javadoc references.
+        "framework-location.stubs.system",
+    ],
     api_packages: [
         "android.location",
         "com.android.location.provider",
diff --git a/media/java/android/media/MediaCas.java b/media/java/android/media/MediaCas.java
index 2d7db5e..8b04644 100644
--- a/media/java/android/media/MediaCas.java
+++ b/media/java/android/media/MediaCas.java
@@ -16,6 +16,9 @@
 
 package android.media;
 
+import static com.android.media.flags.Flags.FLAG_UPDATE_CLIENT_PROFILE_PRIORITY;
+
+import android.annotation.FlaggedApi;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -970,6 +973,27 @@
         registerClient(context, tvInputServiceSessionId, priorityHint);
     }
 
+    /**
+     * Updates client priority with an arbitrary value along with a nice value.
+     *
+     * <p>Tuner resource manager (TRM) uses the client priority value to decide whether it is able
+     * to reclaim insufficient resources from another client.
+     *
+     * <p>The nice value represents how much the client intends to give up the resource when an
+     * insufficient resource situation happens.
+     *
+     * @see <a
+     *     href="https://source.android.com/docs/devices/tv/tuner-framework#priority-nice-value">
+     *     Priority value and nice value</a>
+     * @param priority the new priority. Any negative value would cause no-op on priority setting
+     *     and the API would only process nice value setting in that case.
+     * @param niceValue the nice value.
+     */
+    @FlaggedApi(FLAG_UPDATE_CLIENT_PROFILE_PRIORITY)
+    public boolean updateResourcePriority(int priority, int niceValue) {
+        return mTunerResourceManager.updateClientPriority(mClientId, priority, niceValue);
+    }
+
     IHwBinder getBinder() {
         if (mICas != null) {
             return null; // Return IHwBinder only for HIDL
diff --git a/media/java/android/media/flags/media_better_together.aconfig b/media/java/android/media/flags/media_better_together.aconfig
index 7252135..1ef98f2 100644
--- a/media/java/android/media/flags/media_better_together.aconfig
+++ b/media/java/android/media/flags/media_better_together.aconfig
@@ -70,6 +70,13 @@
 }
 
 flag {
+    name: "update_client_profile_priority"
+    namespace: "media"
+    description : "Feature flag to add updateResourcePriority api to MediaCas"
+    bug: "300565729"
+}
+
+flag {
      name: "enable_built_in_speaker_route_suitability_statuses"
      is_exported: true
      namespace: "media_solutions"
diff --git a/packages/CompanionDeviceManager/res/layout/activity_confirmation.xml b/packages/CompanionDeviceManager/res/layout/activity_confirmation.xml
index 08155dd..5805332 100644
--- a/packages/CompanionDeviceManager/res/layout/activity_confirmation.xml
+++ b/packages/CompanionDeviceManager/res/layout/activity_confirmation.xml
@@ -31,6 +31,12 @@
 
             <!-- A header for selfManaged devices only. -->
             <include layout="@layout/vendor_header" />
+            <!-- A device icon for selfManaged devices only. -->
+            <ImageView
+                android:id="@+id/device_icon"
+                android:visibility="gone"
+                android:contentDescription="@null"
+                style="@style/DeviceIcon" />
 
             <!-- Do NOT change the ID of the root LinearLayout above:
             it's referenced in CTS tests. -->
diff --git a/packages/CompanionDeviceManager/res/values/styles.xml b/packages/CompanionDeviceManager/res/values/styles.xml
index fe7cfc6..a161a50 100644
--- a/packages/CompanionDeviceManager/res/values/styles.xml
+++ b/packages/CompanionDeviceManager/res/values/styles.xml
@@ -137,4 +137,10 @@
         <item name="android:textAppearance">@android:style/TextAppearance.DeviceDefault.Medium</item>
     </style>
 
+    <style name="DeviceIcon">
+        <item name="android:layout_width">24dp</item>
+        <item name="android:layout_height">24dp</item>
+        <item name="android:layout_gravity">center</item>
+    </style>
+
 </resources>
\ No newline at end of file
diff --git a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionAssociationActivity.java b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionAssociationActivity.java
index 7974a37..39bbc25 100644
--- a/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionAssociationActivity.java
+++ b/packages/CompanionDeviceManager/src/com/android/companiondevicemanager/CompanionAssociationActivity.java
@@ -61,6 +61,7 @@
 import android.graphics.BlendModeColorFilter;
 import android.graphics.Color;
 import android.graphics.drawable.Drawable;
+import android.graphics.drawable.Icon;
 import android.net.MacAddress;
 import android.os.Bundle;
 import android.os.Handler;
@@ -130,6 +131,8 @@
 
     // Present for single device and multiple device only.
     private ImageView mProfileIcon;
+    // Present for self managed association only;
+    private ImageView mDeviceIcon;
 
     // Only present for selfManaged devices.
     private ImageView mVendorHeaderImage;
@@ -306,6 +309,8 @@
         mVendorHeaderName = findViewById(R.id.vendor_header_name);
         mVendorHeaderButton = findViewById(R.id.vendor_header_button);
 
+        mDeviceIcon = findViewById(R.id.device_icon);
+
         mDeviceListRecyclerView = findViewById(R.id.device_list);
 
         mMultipleDeviceSpinner = findViewById(R.id.spinner_multiple_device);
@@ -430,6 +435,7 @@
         final Drawable vendorIcon;
         final CharSequence vendorName;
         final Spanned title;
+        final Icon deviceIcon = mRequest.getDeviceIcon();
 
         if (!SUPPORTED_SELF_MANAGED_PROFILES.contains(deviceProfile)) {
             throw new RuntimeException("Unsupported profile " + deviceProfile);
@@ -452,6 +458,11 @@
         title = getHtmlFromResources(this, PROFILE_TITLES.get(deviceProfile), mAppLabel,
                 getString(R.string.device_type), deviceName);
 
+        if (deviceIcon != null) {
+            mDeviceIcon.setImageIcon(deviceIcon);
+            mDeviceIcon.setVisibility(View.VISIBLE);
+        }
+
         if (PROFILE_SUMMARIES.containsKey(deviceProfile)) {
             final int summaryResourceId = PROFILE_SUMMARIES.get(deviceProfile);
             final Spanned summary = getHtmlFromResources(this, summaryResourceId,
diff --git a/packages/SettingsLib/Preference/src/com/android/settingslib/preference/PreferenceFragment.kt b/packages/SettingsLib/Preference/src/com/android/settingslib/preference/PreferenceFragment.kt
index a270681..debaf3e 100644
--- a/packages/SettingsLib/Preference/src/com/android/settingslib/preference/PreferenceFragment.kt
+++ b/packages/SettingsLib/Preference/src/com/android/settingslib/preference/PreferenceFragment.kt
@@ -18,6 +18,7 @@
 
 import android.content.Context
 import android.os.Bundle
+import android.util.Log
 import androidx.annotation.XmlRes
 import androidx.preference.PreferenceFragmentCompat
 import androidx.preference.PreferenceScreen
@@ -42,7 +43,9 @@
     override fun createPreferenceScreen(factory: PreferenceScreenFactory): PreferenceScreen? {
         val context = factory.context
         fun createPreferenceScreenFromResource() =
-            factory.inflate(getPreferenceScreenResId(context))
+            factory.inflate(getPreferenceScreenResId(context))?.also {
+                Log.i(TAG, "Load screen " + it.key + " from resource")
+            }
 
         val screenCreator =
             getPreferenceScreenCreator(context) ?: return createPreferenceScreenFromResource()
@@ -50,10 +53,12 @@
         val preferenceHierarchy = screenCreator.getPreferenceHierarchy(context)
         val preferenceScreen =
             if (screenCreator.hasCompleteHierarchy()) {
+                Log.i(TAG, "Load screen " + screenCreator.key + " from hierarchy")
                 factory.getOrCreatePreferenceScreen().apply {
                     inflatePreferenceHierarchy(preferenceBindingFactory, preferenceHierarchy)
                 }
             } else {
+                Log.i(TAG, "Screen " + screenCreator.key + " is hybrid")
                 createPreferenceScreenFromResource()?.also {
                     bindRecursively(it, preferenceBindingFactory, preferenceHierarchy)
                 } ?: return null
@@ -83,4 +88,8 @@
         preferenceScreenBindingHelper?.close()
         super.onDestroy()
     }
+
+    companion object {
+        private const val TAG = "PreferenceFragment"
+    }
 }
diff --git a/packages/SettingsLib/src/com/android/settingslib/bluetooth/CachedBluetoothDevice.java b/packages/SettingsLib/src/com/android/settingslib/bluetooth/CachedBluetoothDevice.java
index 92da2be..8845d2e 100644
--- a/packages/SettingsLib/src/com/android/settingslib/bluetooth/CachedBluetoothDevice.java
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/CachedBluetoothDevice.java
@@ -38,7 +38,6 @@
 import android.os.Message;
 import android.os.ParcelUuid;
 import android.os.SystemClock;
-import android.provider.Settings;
 import android.text.SpannableStringBuilder;
 import android.text.TextUtils;
 import android.text.style.ForegroundColorSpan;
@@ -1291,12 +1290,10 @@
         if (BluetoothUtils.hasConnectedBroadcastSource(this, mBluetoothManager)) {
             // Gets summary for the buds which are in the audio sharing.
             int groupId = BluetoothUtils.getGroupId(this);
-            if (groupId != BluetoothCsipSetCoordinator.GROUP_ID_INVALID
-                    && groupId
-                            == Settings.Secure.getInt(
-                                    mContext.getContentResolver(),
-                                    "bluetooth_le_broadcast_fallback_active_group_id",
-                                    BluetoothCsipSetCoordinator.GROUP_ID_INVALID)) {
+            int primaryGroupId = BluetoothUtils.getPrimaryGroupIdForBroadcast(
+                    mContext.getContentResolver());
+            if ((primaryGroupId != BluetoothCsipSetCoordinator.GROUP_ID_INVALID)
+                    ? (groupId == primaryGroupId) : isActiveDevice(BluetoothProfile.LE_AUDIO)) {
                 // The buds are primary buds
                 return getSummaryWithBatteryInfo(
                         R.string.bluetooth_active_battery_level_untethered,
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/InputRouteManager.java b/packages/SettingsLib/src/com/android/settingslib/media/InputRouteManager.java
index 9164b64..a72ba8d 100644
--- a/packages/SettingsLib/src/com/android/settingslib/media/InputRouteManager.java
+++ b/packages/SettingsLib/src/com/android/settingslib/media/InputRouteManager.java
@@ -212,13 +212,14 @@
 
     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;
+        // Using 100 for now since it matches the maximum input gain index in classic ChromeOS.
+        return 100;
     }
 
     public int getCurrentInputGain() {
         // TODO (b/357123335): use real input gain implementation.
-        return 8;
+        // Show a fixed full gain in UI before it really works per UX requirement.
+        return 100;
     }
 
     public boolean isInputGainFixed() {
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/CachedBluetoothDeviceTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/CachedBluetoothDeviceTest.java
index 0d81494..70cb2ef 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/CachedBluetoothDeviceTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/CachedBluetoothDeviceTest.java
@@ -1977,12 +1977,12 @@
     }
 
     @Test
-    public void getConnectionSummary_isBroadcastPrimary_returnActive() {
+    public void getConnectionSummary_isBroadcastPrimary_fallbackDevice_returnActive() {
         when(mBroadcast.isEnabled(any())).thenReturn(true);
         when(mCachedDevice.getDevice()).thenReturn(mDevice);
         Settings.Secure.putInt(
                 mContext.getContentResolver(),
-                "bluetooth_le_broadcast_fallback_active_group_id",
+                BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
                 1);
 
         List<Long> bisSyncState = new ArrayList<>();
@@ -1992,12 +1992,30 @@
         sourceList.add(mLeBroadcastReceiveState);
         when(mAssistant.getAllSources(any())).thenReturn(sourceList);
 
-        when(mCachedDevice.getGroupId())
-                .thenReturn(
-                        Settings.Secure.getInt(
-                                mContext.getContentResolver(),
-                                "bluetooth_le_broadcast_fallback_active_group_id",
-                                BluetoothCsipSetCoordinator.GROUP_ID_INVALID));
+        when(mCachedDevice.getGroupId()).thenReturn(1);
+
+        assertThat(mCachedDevice.getConnectionSummary(false))
+                .isEqualTo(mContext.getString(R.string.bluetooth_active_no_battery_level));
+    }
+
+    @Test
+    public void getConnectionSummary_isBroadcastPrimary_activeDevice_returnActive() {
+        when(mBroadcast.isEnabled(any())).thenReturn(true);
+        when(mCachedDevice.getDevice()).thenReturn(mDevice);
+        Settings.Secure.putInt(
+                mContext.getContentResolver(),
+                BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
+                BluetoothCsipSetCoordinator.GROUP_ID_INVALID);
+
+        List<Long> bisSyncState = new ArrayList<>();
+        bisSyncState.add(1L);
+        when(mLeBroadcastReceiveState.getBisSyncState()).thenReturn(bisSyncState);
+        List<BluetoothLeBroadcastReceiveState> sourceList = new ArrayList<>();
+        sourceList.add(mLeBroadcastReceiveState);
+        when(mAssistant.getAllSources(any())).thenReturn(sourceList);
+
+        when(mCachedDevice.getGroupId()).thenReturn(1);
+        when(mCachedDevice.isActiveDevice(BluetoothProfile.LE_AUDIO)).thenReturn(true);
 
         assertThat(mCachedDevice.getConnectionSummary(false))
                 .isEqualTo(mContext.getString(R.string.bluetooth_active_no_battery_level));
@@ -2009,7 +2027,7 @@
         when(mCachedDevice.getDevice()).thenReturn(mDevice);
         Settings.Secure.putInt(
                 mContext.getContentResolver(),
-                "bluetooth_le_broadcast_fallback_active_group_id",
+                BluetoothUtils.getPrimaryGroupIdUriForBroadcast(),
                 1);
 
         List<Long> bisSyncState = new ArrayList<>();
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
index 16c0c1c..29cc403 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputRouteManagerTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputRouteManagerTest.java
@@ -255,12 +255,12 @@
 
     @Test
     public void getMaxInputGain_returnMaxInputGain() {
-        assertThat(mInputRouteManager.getMaxInputGain()).isEqualTo(15);
+        assertThat(mInputRouteManager.getMaxInputGain()).isEqualTo(100);
     }
 
     @Test
     public void getCurrentInputGain_returnCurrentInputGain() {
-        assertThat(mInputRouteManager.getCurrentInputGain()).isEqualTo(8);
+        assertThat(mInputRouteManager.getCurrentInputGain()).isEqualTo(100);
     }
 
     @Test
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index bd7067b..1f10ead 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -1022,6 +1022,7 @@
     defaults: [
         "platform_app_defaults",
         "SystemUI_optimized_defaults",
+        "wmshell_defaults",
     ],
     static_libs: [
         "SystemUI-core",
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index c6238e8..a21a805 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -813,16 +813,6 @@
 }
 
 flag {
-    name: "register_zen_mode_content_observer_background"
-    namespace: "systemui"
-    description: "Decide whether to register zen mode content observers in the background thread."
-    bug: "324515627"
-    metadata {
-        purpose: PURPOSE_BUGFIX
-    }
-}
-
-flag {
     name: "clipboard_noninteractive_on_lockscreen"
     namespace: "systemui"
     description: "Prevents the interactive clipboard UI from appearing when device is locked"
@@ -1483,3 +1473,11 @@
        purpose: PURPOSE_BUGFIX
    }
 }
+
+flag {
+   name: "magic_portrait_wallpapers"
+   namespace: "systemui"
+   description: "Magic Portrait related changes in systemui"
+   bug: "370863642"
+}
+
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/activity/EdgeToEdgeActivitContent.kt b/packages/SystemUI/compose/core/src/com/android/compose/activity/EdgeToEdgeActivitContent.kt
deleted file mode 100644
index 97c8076..0000000
--- a/packages/SystemUI/compose/core/src/com/android/compose/activity/EdgeToEdgeActivitContent.kt
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.compose.activity
-
-import androidx.compose.foundation.background
-import androidx.compose.foundation.isSystemInDarkTheme
-import androidx.compose.foundation.layout.Box
-import androidx.compose.foundation.layout.fillMaxSize
-import androidx.compose.material3.LocalContentColor
-import androidx.compose.material3.MaterialTheme
-import androidx.compose.material3.contentColorFor
-import androidx.compose.runtime.Composable
-import androidx.compose.runtime.CompositionLocalProvider
-import androidx.compose.runtime.DisposableEffect
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.graphics.Color
-import com.android.compose.rememberSystemUiController
-import com.android.compose.theme.PlatformTheme
-
-/** Scaffolding for an edge-to-edge activity content. */
-@Composable
-fun EdgeToEdgeActivityContent(
-    modifier: Modifier = Modifier,
-    content: @Composable () -> Unit,
-) {
-    // Make the status and navigation bars transparent, ensuring that the status bar icons are dark
-    // when the theme is light and vice-versa.
-    val systemUiController = rememberSystemUiController()
-    val isDarkTheme = isSystemInDarkTheme()
-    val useDarkIcons = !isDarkTheme
-    DisposableEffect(systemUiController, useDarkIcons) {
-        systemUiController.setSystemBarsColor(
-            color = Color.Transparent,
-            darkIcons = useDarkIcons,
-        )
-        onDispose {}
-    }
-
-    PlatformTheme(isDarkTheme) {
-        val backgroundColor = MaterialTheme.colorScheme.background
-        Box(modifier.fillMaxSize().background(backgroundColor)) {
-            CompositionLocalProvider(LocalContentColor provides contentColorFor(backgroundColor)) {
-                content()
-            }
-        }
-    }
-}
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt b/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt
index b28655b..37c37b0 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/theme/AndroidColorScheme.kt
@@ -84,40 +84,8 @@
     val secondary = getColor(context, R.attr.materialColorSecondary)
     val tertiary = getColor(context, R.attr.materialColorTertiary)
 
-    @Deprecated("Use the new android tokens: go/sysui-colors")
-    val deprecated = DeprecatedValues(context)
-
-    class DeprecatedValues(context: Context) {
-        val colorPrimary = getColor(context, R.attr.colorPrimary)
-        val colorPrimaryDark = getColor(context, R.attr.colorPrimaryDark)
-        val colorAccent = getColor(context, R.attr.colorAccent)
-        val colorAccentPrimary = getColor(context, R.attr.colorAccentPrimary)
-        val colorAccentSecondary = getColor(context, R.attr.colorAccentSecondary)
-        val colorAccentTertiary = getColor(context, R.attr.colorAccentTertiary)
-        val colorAccentPrimaryVariant = getColor(context, R.attr.colorAccentPrimaryVariant)
-        val colorAccentSecondaryVariant = getColor(context, R.attr.colorAccentSecondaryVariant)
-        val colorAccentTertiaryVariant = getColor(context, R.attr.colorAccentTertiaryVariant)
-        val colorSurface = getColor(context, R.attr.colorSurface)
-        val colorSurfaceHighlight = getColor(context, R.attr.colorSurfaceHighlight)
-        val colorSurfaceVariant = getColor(context, R.attr.colorSurfaceVariant)
-        val colorSurfaceHeader = getColor(context, R.attr.colorSurfaceHeader)
-        val colorError = getColor(context, R.attr.colorError)
-        val colorBackground = getColor(context, R.attr.colorBackground)
-        val colorBackgroundFloating = getColor(context, R.attr.colorBackgroundFloating)
-        val panelColorBackground = getColor(context, R.attr.panelColorBackground)
-        val textColorPrimary = getColor(context, R.attr.textColorPrimary)
-        val textColorSecondary = getColor(context, R.attr.textColorSecondary)
-        val textColorTertiary = getColor(context, R.attr.textColorTertiary)
-        val textColorPrimaryInverse = getColor(context, R.attr.textColorPrimaryInverse)
-        val textColorSecondaryInverse = getColor(context, R.attr.textColorSecondaryInverse)
-        val textColorTertiaryInverse = getColor(context, R.attr.textColorTertiaryInverse)
-        val textColorOnAccent = getColor(context, R.attr.textColorOnAccent)
-        val colorForeground = getColor(context, R.attr.colorForeground)
-        val colorForegroundInverse = getColor(context, R.attr.colorForegroundInverse)
-    }
-
     companion object {
-        fun getColor(context: Context, attr: Int): Color {
+        internal fun getColor(context: Context, attr: Int): Color {
             val ta = context.obtainStyledAttributes(intArrayOf(attr))
             @ColorInt val color = ta.getColor(0, 0)
             ta.recycle()
diff --git a/packages/SystemUI/compose/core/tests/src/com/android/compose/theme/SystemUIThemeTest.kt b/packages/SystemUI/compose/core/tests/src/com/android/compose/theme/PlatformThemeTest.kt
similarity index 90%
rename from packages/SystemUI/compose/core/tests/src/com/android/compose/theme/SystemUIThemeTest.kt
rename to packages/SystemUI/compose/core/tests/src/com/android/compose/theme/PlatformThemeTest.kt
index fe34017..23538e3 100644
--- a/packages/SystemUI/compose/core/tests/src/com/android/compose/theme/SystemUIThemeTest.kt
+++ b/packages/SystemUI/compose/core/tests/src/com/android/compose/theme/PlatformThemeTest.kt
@@ -27,7 +27,7 @@
 import org.junit.runner.RunWith
 
 @RunWith(AndroidJUnit4::class)
-class SystemUIThemeTest {
+class PlatformThemeTest {
     @get:Rule val composeRule = createComposeRule()
 
     @Test
@@ -40,9 +40,7 @@
     @Test
     fun testAndroidColorsAreAvailableInsideTheme() {
         composeRule.setContent {
-            PlatformTheme {
-                Text("foo", color = LocalAndroidColorScheme.current.deprecated.colorAccent)
-            }
+            PlatformTheme { Text("foo", color = LocalAndroidColorScheme.current.primaryFixed) }
         }
 
         composeRule.onNodeWithText("foo").assertIsDisplayed()
@@ -52,7 +50,7 @@
     fun testAccessingAndroidColorsWithoutThemeThrows() {
         assertThrows(IllegalStateException::class.java) {
             composeRule.setContent {
-                Text("foo", color = LocalAndroidColorScheme.current.deprecated.colorAccent)
+                Text("foo", color = LocalAndroidColorScheme.current.primaryFixed)
             }
         }
     }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
index bf8f6ce..7dc2901 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
@@ -522,13 +522,7 @@
     val currentSceneKey =
         if (isSplitAroundTheFold) SceneKeys.SplitSceneKey else SceneKeys.ContiguousSceneKey
 
-    val state = remember {
-        MutableSceneTransitionLayoutState(
-            currentSceneKey,
-            SceneTransitions,
-            enableInterruptions = false,
-        )
-    }
+    val state = remember { MutableSceneTransitionLayoutState(currentSceneKey, SceneTransitions) }
 
     // Update state whenever currentSceneKey has changed.
     LaunchedEffect(state, currentSceneKey) {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreen.kt b/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreen.kt
index b808044..4f1acdc 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreen.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreen.kt
@@ -46,7 +46,7 @@
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
-import com.android.compose.theme.LocalAndroidColorScheme
+import com.android.compose.theme.colorAttr
 import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.people.ui.viewmodel.PeopleTileViewModel
 import com.android.systemui.people.ui.viewmodel.PeopleViewModel
@@ -60,10 +60,7 @@
  *   the Activity/Fragment/View hosting this Composable once a result is available.
  */
 @Composable
-fun PeopleScreen(
-    viewModel: PeopleViewModel,
-    onResult: (PeopleViewModel.Result) -> Unit,
-) {
+fun PeopleScreen(viewModel: PeopleViewModel, onResult: (PeopleViewModel.Result) -> Unit) {
     val priorityTiles by viewModel.priorityTiles.collectAsStateWithLifecycle()
     val recentTiles by viewModel.recentTiles.collectAsStateWithLifecycle()
 
@@ -79,10 +76,9 @@
 
     // Make sure to use the Android colors and not the default Material3 colors to have the exact
     // same colors as the View implementation.
-    val androidColors = LocalAndroidColorScheme.current.deprecated
     Surface(
-        color = androidColors.colorBackground,
-        contentColor = androidColors.textColorPrimary,
+        color = colorAttr(com.android.internal.R.attr.colorBackground),
+        contentColor = colorAttr(com.android.internal.R.attr.textColorPrimary),
         modifier = Modifier.fillMaxSize(),
     ) {
         if (priorityTiles.isNotEmpty() || recentTiles.isNotEmpty()) {
@@ -99,9 +95,7 @@
     recentTiles: List<PeopleTileViewModel>,
     onTileClicked: (PeopleTileViewModel) -> Unit,
 ) {
-    Column(
-        Modifier.sysuiResTag("top_level_with_conversations"),
-    ) {
+    Column(Modifier.sysuiResTag("top_level_with_conversations")) {
         Column(
             Modifier.fillMaxWidth().padding(PeopleSpacePadding),
             horizontalAlignment = Alignment.CenterHorizontally,
@@ -126,12 +120,7 @@
             Modifier.fillMaxWidth()
                 .sysuiResTag("scroll_view")
                 .verticalScroll(rememberScrollState())
-                .padding(
-                    top = 16.dp,
-                    bottom = PeopleSpacePadding,
-                    start = 8.dp,
-                    end = 8.dp,
-                ),
+                .padding(top = 16.dp, bottom = PeopleSpacePadding, start = 8.dp, end = 8.dp)
         ) {
             val hasPriorityConversations = priorityTiles.isNotEmpty()
             if (hasPriorityConversations) {
@@ -153,13 +142,13 @@
 private fun ConversationList(
     @StringRes headerTextResource: Int,
     tiles: List<PeopleTileViewModel>,
-    onTileClicked: (PeopleTileViewModel) -> Unit
+    onTileClicked: (PeopleTileViewModel) -> Unit,
 ) {
     Text(
         stringResource(headerTextResource),
         Modifier.padding(start = 16.dp),
         style = MaterialTheme.typography.labelLarge,
-        color = LocalAndroidColorScheme.current.deprecated.colorAccentPrimaryVariant,
+        color = colorAttr(com.android.internal.R.attr.colorAccentPrimaryVariant),
     )
 
     Spacer(Modifier.height(10.dp))
@@ -167,7 +156,7 @@
     tiles.forEachIndexed { index, tile ->
         if (index > 0) {
             Divider(
-                color = LocalAndroidColorScheme.current.deprecated.colorBackground,
+                color = colorAttr(com.android.internal.R.attr.colorBackground),
                 thickness = 2.dp,
             )
         }
@@ -190,14 +179,13 @@
     withTopCornerRadius: Boolean,
     withBottomCornerRadius: Boolean,
 ) {
-    val androidColors = LocalAndroidColorScheme.current.deprecated
     val cornerRadius = dimensionResource(R.dimen.people_space_widget_radius)
     val topCornerRadius = if (withTopCornerRadius) cornerRadius else 0.dp
     val bottomCornerRadius = if (withBottomCornerRadius) cornerRadius else 0.dp
 
     Surface(
-        color = androidColors.colorSurface,
-        contentColor = androidColors.textColorPrimary,
+        color = colorAttr(com.android.internal.R.attr.colorSurface),
+        contentColor = colorAttr(com.android.internal.R.attr.textColorPrimary),
         shape =
             RoundedCornerShape(
                 topStart = topCornerRadius,
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreenEmpty.kt b/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreenEmpty.kt
index 26cc9b9..d483f88 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreenEmpty.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/people/ui/compose/PeopleScreenEmpty.kt
@@ -41,13 +41,11 @@
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.text.style.TextOverflow
 import androidx.compose.ui.unit.dp
-import com.android.compose.theme.LocalAndroidColorScheme
+import com.android.compose.theme.colorAttr
 import com.android.systemui.res.R
 
 @Composable
-internal fun PeopleScreenEmpty(
-    onGotItClicked: () -> Unit,
-) {
+internal fun PeopleScreenEmpty(onGotItClicked: () -> Unit) {
     Column(
         Modifier.fillMaxSize().padding(PeopleSpacePadding),
         horizontalAlignment = Alignment.CenterHorizontally,
@@ -70,15 +68,14 @@
         ExampleTile()
         Spacer(Modifier.weight(1f))
 
-        val androidColors = LocalAndroidColorScheme.current
         Button(
             onGotItClicked,
             Modifier.fillMaxWidth().defaultMinSize(minHeight = 56.dp),
             colors =
                 ButtonDefaults.buttonColors(
-                    containerColor = androidColors.deprecated.colorAccentPrimary,
-                    contentColor = androidColors.deprecated.textColorOnAccent,
-                )
+                    containerColor = colorAttr(com.android.internal.R.attr.colorAccentPrimary),
+                    contentColor = colorAttr(com.android.internal.R.attr.textColorOnAccent),
+                ),
         ) {
             Text(stringResource(R.string.got_it))
         }
@@ -87,11 +84,10 @@
 
 @Composable
 private fun ExampleTile() {
-    val androidColors = LocalAndroidColorScheme.current
     Surface(
         shape = RoundedCornerShape(28.dp),
-        color = androidColors.deprecated.colorSurface,
-        contentColor = androidColors.deprecated.textColorPrimary,
+        color = colorAttr(com.android.internal.R.attr.colorSurface),
+        contentColor = colorAttr(com.android.internal.R.attr.textColorPrimary),
     ) {
         Row(
             Modifier.padding(vertical = 20.dp, horizontal = 16.dp),
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 e9c5c03..edc4cba 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
@@ -199,7 +199,7 @@
     QuickSettingsTheme {
         val context = LocalContext.current
 
-        LaunchedEffect(key1 = context) {
+        LaunchedEffect(context) {
             if (qsView == null) {
                 qsSceneAdapter.inflate(context)
             }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
index 417f2b8..2d58c8c 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
@@ -24,6 +24,7 @@
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateMapOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
@@ -31,6 +32,8 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.input.pointer.pointerInput
+import androidx.compose.ui.platform.LocalContext
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
 import com.android.compose.animation.scene.OverlayKey
@@ -39,9 +42,13 @@
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.compose.animation.scene.observableTransitionState
+import com.android.systemui.qs.ui.adapter.QSSceneAdapter
+import com.android.systemui.qs.ui.composable.QuickSettingsTheme
 import com.android.systemui.ribbon.ui.composable.BottomRightCornerRibbon
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
+import javax.inject.Provider
 import kotlinx.coroutines.flow.collectLatest
 
 /**
@@ -73,6 +80,7 @@
     overlayByKey: Map<OverlayKey, Overlay>,
     initialSceneKey: SceneKey,
     dataSourceDelegator: SceneDataSourceDelegator,
+    qsSceneAdapter: Provider<QSSceneAdapter>,
     modifier: Modifier = Modifier,
 ) {
     val coroutineScope = rememberCoroutineScope()
@@ -118,6 +126,24 @@
         }
     }
 
+    // Inflate qsView here so that shade has the correct qqs height in the first measure pass after
+    // rebooting
+    if (
+        viewModel.allContentKeys.contains(Scenes.QuickSettings) ||
+            viewModel.allContentKeys.contains(Scenes.Shade)
+    ) {
+        val qsAdapter = qsSceneAdapter.get()
+        QuickSettingsTheme {
+            val context = LocalContext.current
+            val qsView by qsAdapter.qsView.collectAsStateWithLifecycle()
+            LaunchedEffect(context) {
+                if (qsView == null) {
+                    qsAdapter.inflate(context)
+                }
+            }
+        }
+    }
+
     Box(
         modifier =
             Modifier.fillMaxSize().pointerInput(Unit) {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
index 2e8fc14..2657d7c 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
@@ -234,7 +234,6 @@
     canHideOverlay: (OverlayKey) -> Boolean = { true },
     canReplaceOverlay: (from: OverlayKey, to: OverlayKey) -> Boolean = { _, _ -> true },
     stateLinks: List<StateLink> = emptyList(),
-    enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
 ): MutableSceneTransitionLayoutState {
     return MutableSceneTransitionLayoutStateImpl(
         initialScene,
@@ -245,7 +244,6 @@
         canHideOverlay,
         canReplaceOverlay,
         stateLinks,
-        enableInterruptions,
     )
 }
 
@@ -261,9 +259,6 @@
         true
     },
     private val stateLinks: List<StateLink> = emptyList(),
-
-    // TODO(b/290930950): Remove this flag.
-    internal val enableInterruptions: Boolean = DEFAULT_INTERRUPTIONS_ENABLED,
 ) : MutableSceneTransitionLayoutState {
     private val creationThread: Thread = Thread.currentThread()
 
@@ -406,13 +401,6 @@
             transition.updateOverscrollSpecs(fromSpec = null, toSpec = null)
         }
 
-        if (!enableInterruptions) {
-            // Set the current transition.
-            check(transitionStates.size == 1)
-            transitionStates = listOf(transition)
-            return
-        }
-
         when (val currentState = transitionStates.last()) {
             is TransitionState.Idle -> {
                 // Replace [Idle] by [transition].
@@ -755,9 +743,6 @@
 
 private const val TAG = "SceneTransitionLayoutState"
 
-/** Whether support for interruptions in enabled by default. */
-internal const val DEFAULT_INTERRUPTIONS_ENABLED = true
-
 /**
  * The max number of concurrent transitions. If the number of transitions goes past this number,
  * this probably means that there is a leak and we will Log.wtf before clearing the list of
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
index d6751ae..9d3f25e 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
@@ -337,10 +337,6 @@
         }
 
         internal open fun interruptionProgress(layoutImpl: SceneTransitionLayoutImpl): Float {
-            if (!layoutImpl.state.enableInterruptions) {
-                return 0f
-            }
-
             if (replacedTransition != null) {
                 return replacedTransition.interruptionProgress(layoutImpl)
             }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
index 1eed54e..a2c2729 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
@@ -213,10 +213,6 @@
                         from(SceneA, to = SceneB) { spec = tween }
                         from(SceneB, to = SceneC) { spec = tween }
                     },
-
-                    // Disable interruptions so that the current transition is directly removed
-                    // when starting a new one.
-                    enableInterruptions = false,
                 )
             }
 
@@ -243,7 +239,12 @@
                 onElement(TestElements.Bar).assertExists()
 
                 // Start transition from SceneB to SceneC
-                rule.runOnUiThread { state.setTargetScene(SceneC, coroutineScope) }
+                rule.runOnUiThread {
+                    // We snap to scene B so that the transition A => B is removed from the list of
+                    // transitions.
+                    state.snapToScene(SceneB)
+                    state.setTargetScene(SceneC, coroutineScope)
+                }
             }
 
             at(3 * frameDuration) { onElement(TestElements.Bar).assertIsNotDisplayed() }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
index 400f0b3..b00b894b 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
@@ -85,15 +85,9 @@
 
     /** The content under test. */
     @Composable
-    private fun TestContent(enableInterruptions: Boolean = true) {
+    private fun TestContent() {
         coroutineScope = rememberCoroutineScope()
-        layoutState = remember {
-            MutableSceneTransitionLayoutState(
-                SceneA,
-                EmptyTestTransitions,
-                enableInterruptions = enableInterruptions,
-            )
-        }
+        layoutState = remember { MutableSceneTransitionLayoutState(SceneA, EmptyTestTransitions) }
 
         SceneTransitionLayout(state = layoutState, modifier = Modifier.size(LayoutSize)) {
             scene(SceneA, userActions = mapOf(Back to SceneB)) {
@@ -205,7 +199,7 @@
 
     @Test
     fun testSharedElement() {
-        rule.setContent { TestContent(enableInterruptions = false) }
+        rule.setContent { TestContent() }
 
         // In scene A, the shared element SharedFoo() is at the top end of the layout and has a size
         // of 50.dp.
@@ -253,6 +247,9 @@
             .isWithin(DpOffsetSubject.DefaultTolerance)
             .of(DpOffset(25.dp, 25.dp))
 
+        // Finish the transition.
+        rule.mainClock.advanceTimeBy(TestTransitionDuration / 2)
+
         // Animate to scene C, let the animation start then go to the middle of the transition.
         currentScene = SceneC
         rule.mainClock.advanceTimeByFrame()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
index 43c7ed6..9c58e2b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.time.FakeSystemClock
+import com.android.wm.shell.keyguard.KeyguardTransitions
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -46,6 +47,7 @@
     @Mock private lateinit var keyguardSurfaceBehindAnimator: KeyguardSurfaceBehindParamsApplier
     @Mock
     private lateinit var keyguardDismissTransitionInteractor: KeyguardDismissTransitionInteractor
+    @Mock private lateinit var keyguardTransitions: KeyguardTransitions
 
     @Before
     fun setUp() {
@@ -59,6 +61,7 @@
                 keyguardStateController = keyguardStateController,
                 keyguardSurfaceBehindAnimator = keyguardSurfaceBehindAnimator,
                 keyguardDismissTransitionInteractor = keyguardDismissTransitionInteractor,
+                keyguardTransitions = keyguardTransitions,
             )
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
index 5186536..12eadfc 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelTest.kt
@@ -184,17 +184,6 @@
         }
 
     @Test
-    fun iconContainer_isVisible_bypassEnabled() =
-        testScope.runTest {
-            val isVisible by collectLastValue(underTest.isNotifIconContainerVisible)
-            runCurrent()
-            deviceEntryRepository.setBypassEnabled(true)
-            runCurrent()
-
-            assertThat(isVisible?.value).isTrue()
-        }
-
-    @Test
     fun iconContainer_isNotVisible_pulseExpanding_notBypassing() =
         testScope.runTest {
             val isVisible by collectLastValue(underTest.isNotifIconContainerVisible)
@@ -283,22 +272,23 @@
         }
 
     @Test
-    fun iconContainer_isNotVisible_bypassDisabled_onLockscreen() =
+    fun iconContainer_isNotVisible_notifsFullyHiddenThenVisible_bypassEnabled() =
         testScope.runTest {
             val isVisible by collectLastValue(underTest.isNotifIconContainerVisible)
             runCurrent()
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.AOD,
-                to = KeyguardState.LOCKSCREEN,
-                testScope,
-            )
             notificationsKeyguardInteractor.setPulseExpanding(false)
-            deviceEntryRepository.setBypassEnabled(false)
+            deviceEntryRepository.setBypassEnabled(true)
             whenever(dozeParameters.alwaysOn).thenReturn(true)
             whenever(dozeParameters.displayNeedsBlanking).thenReturn(false)
             notificationsKeyguardInteractor.setNotificationsFullyHidden(true)
             runCurrent()
 
+            assertThat(isVisible?.value).isTrue()
+            assertThat(isVisible?.isAnimating).isTrue()
+
+            notificationsKeyguardInteractor.setNotificationsFullyHidden(false)
+            runCurrent()
+
             assertThat(isVisible?.value).isFalse()
             assertThat(isVisible?.isAnimating).isTrue()
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/privacy/PrivacyDialogControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/privacy/PrivacyDialogControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/privacy/PrivacyDialogControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/privacy/PrivacyDialogControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/privacy/PrivacyDialogControllerV2Test.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/privacy/PrivacyDialogControllerV2Test.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/privacy/PrivacyDialogControllerV2Test.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/privacy/PrivacyDialogControllerV2Test.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qrcodescanner/controller/QRCodeScannerControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/qrcodescanner/controller/QRCodeScannerControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qrcodescanner/controller/QRCodeScannerControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qrcodescanner/controller/QRCodeScannerControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/TileLayoutTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/TileLayoutTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qs/TileLayoutTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/TileLayoutTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/compose/DragAndDropTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/compose/ResizingTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/HearingDevicesTileTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/HearingDevicesTileTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qs/tiles/HearingDevicesTileTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/HearingDevicesTileTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/InternetTileNewImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/InternetTileNewImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qs/tiles/InternetTileNewImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/InternetTileNewImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenrecord/RecordingServiceTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/RecordingServiceTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenrecord/RecordingServiceTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/RecordingServiceTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialogDelegateTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialogDelegateTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialogDelegateTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/ScreenRecordPermissionDialogDelegateTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenrecord/data/model/ScreenRecordModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/data/model/ScreenRecordModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenrecord/data/model/ScreenRecordModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/data/model/ScreenRecordModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenrecord/data/repository/ScreenRecordRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/data/repository/ScreenRecordRepositoryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenrecord/data/repository/ScreenRecordRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenrecord/data/repository/ScreenRecordRepositoryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ActionIntentExecutorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ActionIntentExecutorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ActionIntentExecutorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ActionIntentExecutorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/AnnouncementResolverTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/AnnouncementResolverTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/AnnouncementResolverTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/AnnouncementResolverTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/DraggableConstraintLayoutTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/DraggableConstraintLayoutTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/DraggableConstraintLayoutTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/DraggableConstraintLayoutTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeImageCapture.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/FakeImageCapture.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeImageCapture.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/FakeImageCapture.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeScreenshotPolicy.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/FakeScreenshotPolicy.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeScreenshotPolicy.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/FakeScreenshotPolicy.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/RecyclerViewActivity.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/RecyclerViewActivity.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/RecyclerViewActivity.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/RecyclerViewActivity.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotActionsControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotActionsControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotActionsControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotActionsControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDetectionControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotDetectionControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDetectionControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotDetectionControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotNotificationSmartActionsTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotNotificationSmartActionsTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotNotificationSmartActionsTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotNotificationSmartActionsTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotSoundControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotSoundControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotSoundControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ScreenshotSoundControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/SmartActionsReceiverTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/WorkProfileMessageControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/WorkProfileMessageControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/WorkProfileMessageControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/WorkProfileMessageControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/appclips/AppClipsScreenshotHelperServiceTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsServiceTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/appclips/AppClipsServiceTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsServiceTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/appclips/AppClipsServiceTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/data/repository/ProfileTypeRepositoryKosmos.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/repository/ProfileTypeRepositoryKosmos.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/data/repository/ProfileTypeRepositoryKosmos.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/repository/ProfileTypeRepositoryKosmos.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/message/ProfileMessageControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/message/ProfileMessageControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/message/ProfileMessageControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/message/ProfileMessageControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/TestUserIds.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/TestUserIds.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/TestUserIds.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/TestUserIds.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/FakeSessionTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/FakeSessionTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/FakeSessionTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/FakeSessionTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/ScrollCaptureControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/ScrollCaptureControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureFrameworkSmokeTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/ScrollCaptureFrameworkSmokeTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureFrameworkSmokeTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/ScrollCaptureFrameworkSmokeTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollViewActivity.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/ScrollViewActivity.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollViewActivity.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/scroll/ScrollViewActivity.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/scrim/ScrimViewTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/scrim/ScrimViewTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/scrim/ScrimViewTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/scrim/ScrimViewTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/sensorprivacy/SensorUseStartedActivityTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/sensorprivacy/SensorUseStartedActivityTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/sensorprivacy/SensorUseStartedActivityTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/sensorprivacy/SensorUseStartedActivityTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/settings/brightness/BrightnessControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/settings/brightness/BrightnessControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/settings/brightness/BrightnessControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/settings/brightness/BrightnessControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/settings/brightness/BrightnessDialogTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/settings/brightness/BrightnessDialogTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/settings/brightness/BrightnessDialogTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/settings/brightness/BrightnessDialogTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/settings/brightness/BrightnessSliderControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/settings/brightness/BrightnessSliderControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/settings/brightness/BrightnessSliderControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/settings/brightness/BrightnessSliderControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/CombinedShadeHeaderConstraintsTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/CombinedShadeHeaderConstraintsTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/CombinedShadeHeaderConstraintsTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/CombinedShadeHeaderConstraintsTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/ConstraintChangeTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ConstraintChangeTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/ConstraintChangeTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ConstraintChangeTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/ConstraintChangesTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ConstraintChangesTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/ConstraintChangesTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ConstraintChangesTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/LockscreenHostedDreamGestureListenerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/LockscreenHostedDreamGestureListenerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/LockscreenHostedDreamGestureListenerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/LockscreenHostedDreamGestureListenerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelUnfoldAnimationControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelUnfoldAnimationControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelUnfoldAnimationControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelUnfoldAnimationControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowViewControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationShadeWindowViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationsQuickSettingsContainerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationsQuickSettingsContainerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/NotificationsQuickSettingsContainerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/NotificationsQuickSettingsContainerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/PulsingGestureListenerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/PulsingGestureListenerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/PulsingGestureListenerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/PulsingGestureListenerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QsBatteryModeControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QsBatteryModeControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/QsBatteryModeControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QsBatteryModeControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QuickSettingsControllerImplTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QuickSettingsControllerImplTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplWithCoroutinesTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QuickSettingsControllerImplWithCoroutinesTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplWithCoroutinesTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/QuickSettingsControllerImplWithCoroutinesTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/ShadeControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ShadeControllerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/ShadeControllerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ShadeControllerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/carrier/CellSignalStateTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/carrier/CellSignalStateTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/carrier/CellSignalStateTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/carrier/CellSignalStateTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/carrier/ShadeCarrierTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/carrier/ShadeCarrierTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/carrier/ShadeCarrierTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/carrier/ShadeCarrierTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/data/repository/ShadeRepositoryImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/transition/ScrimShadeTransitionControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shadow/DoubleShadowTextClockTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/animation/DisableSubpixelTextTransitionListenerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/animation/DisableSubpixelTextTransitionListenerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/animation/DisableSubpixelTextTransitionListenerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/animation/DisableSubpixelTextTransitionListenerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/animation/UnfoldConstantTranslateAnimatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/animation/UnfoldConstantTranslateAnimatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/animation/UnfoldConstantTranslateAnimatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/animation/UnfoldConstantTranslateAnimatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/animation/UnfoldMoveFromCenterAnimatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/animation/UnfoldMoveFromCenterAnimatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/animation/UnfoldMoveFromCenterAnimatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/animation/UnfoldMoveFromCenterAnimatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/clocks/AnimatableClockViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/clocks/AnimatableClockViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/clocks/AnimatableClockViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/clocks/AnimatableClockViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/clocks/ClockRegistryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/clocks/ClockRegistryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/clocks/ClockRegistryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/clocks/ClockRegistryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/clocks/DefaultClockProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/clocks/DefaultClockProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/clocks/DefaultClockProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/clocks/DefaultClockProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/condition/CombinedConditionTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/CombinedConditionTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/condition/CombinedConditionTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/CombinedConditionTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/condition/ConditionExtensionsTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/ConditionExtensionsTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/condition/ConditionExtensionsTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/ConditionExtensionsTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/condition/ConditionMonitorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/ConditionMonitorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/condition/ConditionMonitorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/ConditionMonitorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/condition/ConditionTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/ConditionTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/condition/ConditionTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/ConditionTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/condition/FakeCondition.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/FakeCondition.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/condition/FakeCondition.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/condition/FakeCondition.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepositoryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/plugins/PluginInstanceTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/plugins/PluginInstanceTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/plugins/PluginInstanceTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/plugins/PluginInstanceTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/plugins/PluginManagerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/plugins/PluginManagerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/plugins/PluginManagerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/plugins/PluginManagerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/plugins/VersionInfoTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/plugins/VersionInfoTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/plugins/VersionInfoTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/plugins/VersionInfoTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/regionsampling/RegionSamplerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/regionsampling/RegionSamplerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/regionsampling/RegionSamplerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/regionsampling/RegionSamplerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/rotation/RotationButtonControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/rotation/RotationButtonControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/rotation/RotationButtonControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/rotation/RotationButtonControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/system/QuickStepContractTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/system/QuickStepContractTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/system/QuickStepContractTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/system/QuickStepContractTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/system/UncaughtExceptionPreHandlerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shared/system/UncaughtExceptionPreHandlerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/system/UncaughtExceptionPreHandlerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shared/system/UncaughtExceptionPreHandlerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/BlurUtilsTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/BlurUtilsTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/BlurUtilsTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/BlurUtilsTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/DragDownHelperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/DragDownHelperTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/DragDownHelperTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/DragDownHelperTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyguardIndicationControllerBaseTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/KeyguardIndicationControllerBaseTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/KeyguardIndicationControllerBaseTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/KeyguardIndicationControllerBaseTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/LSShadeTransitionLoggerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/LSShadeTransitionLoggerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/LSShadeTransitionLoggerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/LSShadeTransitionLoggerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/LockscreenShadeQsTransitionControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/LockscreenShadeQsTransitionControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/LockscreenShadeQsTransitionControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/LockscreenShadeQsTransitionControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationListenerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationListenerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationListenerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationShadeDepthControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationShadeDepthControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationShadeDepthControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationShadeDepthControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationUiAdjustmentTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationUiAdjustmentTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/NotificationUiAdjustmentTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationUiAdjustmentTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/PulseExpansionHandlerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/PulseExpansionHandlerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/PulseExpansionHandlerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/PulseExpansionHandlerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilderTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilderTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilderTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/RemoteInputNotificationRebuilderTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SingleShadeLockScreenOverScrollerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/SingleShadeLockScreenOverScrollerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/SingleShadeLockScreenOverScrollerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/SingleShadeLockScreenOverScrollerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateEventTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarStateEventTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateEventTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarStateEventTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/VibratorHelperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/VibratorHelperTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/VibratorHelperTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/VibratorHelperTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/view/ChipBackgroundContainerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/chips/ui/view/ChipBackgroundContainerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/view/ChipBackgroundContainerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/chips/ui/view/ChipBackgroundContainerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/view/ChipChronometerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/chips/ui/view/ChipChronometerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/view/ChipChronometerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/chips/ui/view/ChipChronometerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ParametersTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/commandline/ParametersTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ParametersTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/commandline/ParametersTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ParseableCommandTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/commandline/ParseableCommandTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ParseableCommandTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/commandline/ParseableCommandTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/AccessPointControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/AccessPointControllerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/AccessPointControllerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/AccessPointControllerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/MobileIconCarrierIdOverridesFake.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/MobileIconCarrierIdOverridesFake.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/MobileIconCarrierIdOverridesFake.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/MobileIconCarrierIdOverridesFake.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkTypeResIdCacheTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/NetworkTypeResIdCacheTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkTypeResIdCacheTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/NetworkTypeResIdCacheTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/ui/MobileContextProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/ui/MobileContextProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/ui/MobileContextProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/connectivity/ui/MobileContextProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/core/CommandQueueInitializerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/CommandQueueInitializerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/core/CommandQueueInitializerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/CommandQueueInitializerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarInitializerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/core/StatusBarOrchestratorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarOrchestratorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/core/StatusBarOrchestratorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/core/StatusBarOrchestratorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/data/repository/KeyguardStatusBarRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/data/repository/KeyguardStatusBarRepositoryImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/data/repository/KeyguardStatusBarRepositoryImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/data/repository/KeyguardStatusBarRepositoryImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/data/repository/StatusBarModeRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/data/repository/StatusBarModeRepositoryImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/data/repository/StatusBarModeRepositoryImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/data/repository/StatusBarModeRepositoryImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/disableflags/DisableFlagsLoggerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/disableflags/DisableFlagsLoggerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/disableflags/DisableFlagsLoggerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/disableflags/DisableFlagsLoggerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/disableflags/DisableStateTrackerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/disableflags/DisableStateTrackerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/disableflags/DisableStateTrackerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/disableflags/DisableStateTrackerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/disableflags/data/repository/DisableFlagsRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/disableflags/data/repository/DisableFlagsRepositoryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/disableflags/data/repository/DisableFlagsRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/disableflags/data/repository/DisableFlagsRepositoryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/events/FakeStatusEvent.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/FakeStatusEvent.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/events/FakeStatusEvent.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/FakeStatusEvent.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventChipAnimationControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemEventCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/events/SystemEventCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/events/SystemEventCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/DynamicChildBindControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicPrivacyControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/DynamicPrivacyControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/DynamicPrivacyControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/DynamicPrivacyControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorLoggerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorLoggerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorLoggerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorLoggerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/NotificationWakeUpCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/HighPriorityProviderTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/HighPriorityProviderTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/HighPriorityProviderTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/HighPriorityProviderTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NoManSimulator.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NoManSimulator.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NoManSimulator.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NoManSimulator.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreMocks.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreMocks.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreMocks.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifLiveDataStoreMocks.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifPipelineChoreographerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifPipelineChoreographerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/NotifPipelineChoreographerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotifPipelineChoreographerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/SectionStyleProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/SectionStyleProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/SectionStyleProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/SectionStyleProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coalescer/GroupCoalescerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coalescer/GroupCoalescerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coalescer/GroupCoalescerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coalescer/GroupCoalescerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/ColorizedFgsCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/ColorizedFgsCoordinatorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/ColorizedFgsCoordinatorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/ColorizedFgsCoordinatorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/DeviceProvisionedCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/DeviceProvisionedCoordinatorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/DeviceProvisionedCoordinatorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/DeviceProvisionedCoordinatorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/DismissibilityCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/DismissibilityCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/DismissibilityCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/DismissibilityCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/DreamCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/DreamCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/DreamCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/DreamCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupCountCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupCountCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupCountCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupCountCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupWhenCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupWhenCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupWhenCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/GroupWhenCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/GutsCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/GutsCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/GutsCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/GutsCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/HideNotifsForOtherUsersCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HideNotifsForOtherUsersCoordinatorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/HideNotifsForOtherUsersCoordinatorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HideNotifsForOtherUsersCoordinatorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/MediaCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/MediaCoordinatorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/MediaCoordinatorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/MediaCoordinatorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/NotificationStatsLoggerCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/NotificationStatsLoggerCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/NotificationStatsLoggerCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/NotificationStatsLoggerCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/PreparationCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/PreparationCoordinatorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/PreparationCoordinatorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/PreparationCoordinatorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinatorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinatorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinatorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/RankingCoordinatorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RowAlertTimeCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/RowAlertTimeCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/RowAlertTimeCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/RowAlertTimeCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/SensitiveContentCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/SensitiveContentCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/SensitiveContentCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/SensitiveContentCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/SmartspaceDedupingCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/SmartspaceDedupingCoordinatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/SmartspaceDedupingCoordinatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/SmartspaceDedupingCoordinatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/inflation/NotifUiAdjustmentProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/inflation/NotifUiAdjustmentProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/inflation/NotifUiAdjustmentProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/inflation/NotifUiAdjustmentProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/listbuilder/SemiStableSortTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/listbuilder/SemiStableSortTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/listbuilder/SemiStableSortTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/listbuilder/SemiStableSortTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/listbuilder/ShadeListBuilderHelperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/listbuilder/ShadeListBuilderHelperTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/listbuilder/ShadeListBuilderHelperTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/listbuilder/ShadeListBuilderHelperTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionInconsistencyTrackerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionInconsistencyTrackerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionInconsistencyTrackerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/notifcollection/NotifCollectionInconsistencyTrackerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/notifcollection/SelfTrackingLifetimeExtenderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/provider/VisualStabilityProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/provider/VisualStabilityProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/provider/VisualStabilityProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/provider/VisualStabilityProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/FakeNodeController.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/FakeNodeController.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/FakeNodeController.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/FakeNodeController.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/RenderStageManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/RenderStageManagerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/RenderStageManagerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/RenderStageManagerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/ShadeViewDifferTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/ShadeViewDifferTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/ShadeViewDifferTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/render/ShadeViewDifferTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationAlertsInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationAlertsInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationAlertsInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationAlertsInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationLaunchAnimationInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationLaunchAnimationInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationLaunchAnimationInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationLaunchAnimationInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/NotificationsKeyguardInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/RenderNotificationsListInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/footer/ui/viewmodel/FooterViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/footer/ui/viewmodel/FooterViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/footer/ui/viewmodel/FooterViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/footer/ui/viewmodel/FooterViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/icon/domain/interactor/NotificationIconsInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/icon/ui/viewmodel/NotificationIconContainerAlwaysOnDisplayViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/HeadsUpViewBinderTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/HeadsUpViewBinderTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/HeadsUpViewBinderTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/HeadsUpViewBinderTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/KeyguardNotificationVisibilityProviderTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/KeyguardNotificationVisibilityProviderTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/KeyguardNotificationVisibilityProviderTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/KeyguardNotificationVisibilityProviderTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestBase.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/interruption/VisualInterruptionDecisionProviderTestUtil.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/ExpansionStateLoggerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/logging/ExpansionStateLoggerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/ExpansionStateLoggerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/logging/ExpansionStateLoggerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLoggerFake.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLoggerFake.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLoggerFake.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLoggerFake.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/FeedbackInfoTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/FeedbackInfoTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/FeedbackInfoTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/FeedbackInfoTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/HeadsUpStyleProviderImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/HeadsUpStyleProviderImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/HeadsUpStyleProviderImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/HeadsUpStyleProviderImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotifBindPipelineTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotifBindPipelineTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotifBindPipelineTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotifBindPipelineTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotifInflationErrorManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotifInflationErrorManagerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotifInflationErrorManagerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotifInflationErrorManagerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewCacheImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewCacheImplTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewCacheImplTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotifRemoteViewCacheImplTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationGutsTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationGutsTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationGutsTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationInfoTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationInfoTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationInfoTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationInfoTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolverTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolverTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolverTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolverTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/PartialConversationInfoTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/PartialConversationInfoTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/PartialConversationInfoTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/PartialConversationInfoTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/RowContentBindStageTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/RowContentBindStageTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/RowContentBindStageTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/RowContentBindStageTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/ActivatableNotificationViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/ActivatableNotificationViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/ActivatableNotificationViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/ActivatableNotificationViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/NotificationViewFlipperViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/NotificationViewFlipperViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/NotificationViewFlipperViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/ui/viewmodel/NotificationViewFlipperViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/shared/TestActiveNotificationModel.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/shared/TestActiveNotificationModel.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/shared/TestActiveNotificationModel.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/shared/TestActiveNotificationModel.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/shelf/domain/interactor/NotificationShelfInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/shelf/domain/interactor/NotificationShelfInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/shelf/domain/interactor/NotificationShelfInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/shelf/domain/interactor/NotificationShelfInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/shelf/ui/viewmodel/NotificationShelfViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/shelf/ui/viewmodel/NotificationShelfViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/shelf/ui/viewmodel/NotificationShelfViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/shelf/ui/viewmodel/NotificationShelfViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/AmbientStateTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/DisplaySwitchNotificationsHiderTrackerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/DisplaySwitchNotificationsHiderTrackerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/DisplaySwitchNotificationsHiderTrackerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/DisplaySwitchNotificationsHiderTrackerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/MediaContainerViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/MediaContainerViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/MediaContainerViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/MediaContainerViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/NotificationStackSizeCalculatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackStateAnimatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/StackStateAnimatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackStateAnimatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/StackStateAnimatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ViewStateTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ViewStateTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ViewStateTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ViewStateTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/HideNotificationsInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/domain/interactor/NotificationStackInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationLoggerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationLoggerViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationLoggerViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationLoggerViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/BiometricsUnlockControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/BiometricsUnlockControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/BiometricsUnlockControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/BiometricsUnlockControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/CentralSurfacesCommandQueueCallbacksTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ConfigurationControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ConfigurationControllerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ConfigurationControllerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ConfigurationControllerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardBypassControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/KeyguardBypassControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardBypassControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/KeyguardBypassControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardDismissUtilTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/KeyguardDismissUtilTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardDismissUtilTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/KeyguardDismissUtilTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardIndicationTextViewTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/KeyguardIndicationTextViewTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardIndicationTextViewTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/KeyguardIndicationTextViewTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LegacyLightsOutNotifControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyLightsOutNotifControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LegacyLightsOutNotifControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LegacyLightsOutNotifControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LetterboxAppearanceCalculatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LetterboxAppearanceCalculatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LetterboxAppearanceCalculatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LetterboxAppearanceCalculatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LetterboxBackgroundProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LetterboxBackgroundProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LetterboxBackgroundProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LetterboxBackgroundProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LightBarControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LightBarControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/LightBarControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/LightBarControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationGroupTestHelper.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/NotificationGroupTestHelper.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationGroupTestHelper.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/NotificationGroupTestHelper.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationIconContainerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/NotificationIconContainerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationIconContainerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/NotificationIconContainerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationTapHelperTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/NotificationTapHelperTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/NotificationTapHelperTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/NotificationTapHelperTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarNotificationPresenterTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarRemoteInputCallbackTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarRemoteInputCallbackTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarRemoteInputCallbackTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarRemoteInputCallbackTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManagerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManagerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusBarTouchableRegionManagerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusOverlayHoverListenerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusOverlayHoverListenerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusOverlayHoverListenerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/StatusOverlayHoverListenerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/domain/interactor/LightsOutInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/domain/interactor/LightsOutInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/domain/interactor/LightsOutInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/domain/interactor/LightsOutInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragmentLoggerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragmentLoggerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragmentLoggerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/fragment/CollapsedStatusBarFragmentLoggerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/fragment/StatusBarVisibilityModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/fragment/StatusBarVisibilityModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/fragment/StatusBarVisibilityModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/fragment/StatusBarVisibilityModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaListenerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaListenerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaListenerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaListenerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaRepoTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaRepoTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaRepoTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ongoingcall/OngoingCallControllerViaRepoTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ongoingcall/data/repository/OngoingCallRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ongoingcall/data/repository/OngoingCallRepositoryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ongoingcall/data/repository/OngoingCallRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ongoingcall/data/repository/OngoingCallRepositoryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconListTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconListTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconListTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconListTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/airplane/data/repository/AirplaneModeRepositoryImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/ethernet/domain/EthernetInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/ethernet/domain/EthernetInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/ethernet/domain/EthernetInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/ethernet/domain/EthernetInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/CarrierMergedConnectionRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/CarrierMergedConnectionRepositoryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/CarrierMergedConnectionRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/CarrierMergedConnectionRepositoryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileViewLoggerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileViewLoggerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileViewLoggerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileViewLoggerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/model/SignalIconModelParameterizedTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/model/SignalIconModelParameterizedTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/model/SignalIconModelParameterizedTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/model/SignalIconModelParameterizedTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/data/DeviceBasedSatelliteRepositorySwitcherTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/data/DeviceBasedSatelliteRepositorySwitcherTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/data/DeviceBasedSatelliteRepositorySwitcherTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/data/DeviceBasedSatelliteRepositorySwitcherTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/data/demo/DemoDeviceBasedSatelliteRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/data/demo/DemoDeviceBasedSatelliteRepositoryTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/data/demo/DemoDeviceBasedSatelliteRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/data/demo/DemoDeviceBasedSatelliteRepositoryTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/data/prod/FakeDeviceBasedSatelliteRepository.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/data/prod/FakeDeviceBasedSatelliteRepository.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/data/prod/FakeDeviceBasedSatelliteRepository.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/data/prod/FakeDeviceBasedSatelliteRepository.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/domain/interactor/DeviceBasedSatelliteInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/domain/interactor/DeviceBasedSatelliteInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/domain/interactor/DeviceBasedSatelliteInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/domain/interactor/DeviceBasedSatelliteInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/DeviceBasedSatelliteViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/DeviceBasedSatelliteViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/DeviceBasedSatelliteViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/DeviceBasedSatelliteViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/FakeDeviceBasedSatelliteViewModel.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/FakeDeviceBasedSatelliteViewModel.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/FakeDeviceBasedSatelliteViewModel.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/satellite/ui/viewmodel/FakeDeviceBasedSatelliteViewModel.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/data/model/DefaultConnectionModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/data/model/DefaultConnectionModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/data/model/DefaultConnectionModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/data/model/DefaultConnectionModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/domain/interactor/CollapsedStatusBarInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/domain/interactor/CollapsedStatusBarInteractorTest.kt
similarity index 96%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/domain/interactor/CollapsedStatusBarInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/domain/interactor/CollapsedStatusBarInteractorTest.kt
index 5036e77..46f822a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/domain/interactor/CollapsedStatusBarInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/domain/interactor/CollapsedStatusBarInteractorTest.kt
@@ -21,6 +21,7 @@
 import android.app.StatusBarManager.DISABLE_NONE
 import android.app.StatusBarManager.DISABLE_NOTIFICATION_ICONS
 import android.app.StatusBarManager.DISABLE_SYSTEM_INFO
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -31,8 +32,10 @@
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
 import kotlinx.coroutines.test.runTest
+import org.junit.runner.RunWith;
 
 @SmallTest
+@RunWith(AndroidJUnit4::class)
 class CollapsedStatusBarInteractorTest : SysuiTestCase() {
     val kosmos = testKosmos()
     val testScope = kosmos.testScope
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/ModernStatusBarViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/ModernStatusBarViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/ModernStatusBarViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/ModernStatusBarViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/SingleBindableStatusBarIconViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/SingleBindableStatusBarIconViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/SingleBindableStatusBarIconViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/view/SingleBindableStatusBarIconViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewBinder.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewBinder.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewBinder.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewBinder.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewModel.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewModel.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewModel.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/FakeCollapsedStatusBarViewModel.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BaseUserSwitcherAdapterTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BaseUserSwitcherAdapterTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BaseUserSwitcherAdapterTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BaseUserSwitcherAdapterTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BlockingQueueIntentReceiver.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BlockingQueueIntentReceiver.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BlockingQueueIntentReceiver.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BlockingQueueIntentReceiver.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
similarity index 99%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
index 2588f1f..e3bd885 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/BluetoothControllerImplTest.java
@@ -61,7 +61,7 @@
 import java.util.List;
 import java.util.concurrent.Executor;
 
-@RunWith(AndroidTestingRunner.class)
+@RunWith(AndroidJUnit4.class)
 @RunWithLooper
 @SmallTest
 public class BluetoothControllerImplTest extends SysuiTestCase {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/CastControllerImplTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ClockTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ClockTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ClockTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ClockTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/DevicePostureControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/DevicePostureControllerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/DevicePostureControllerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/DevicePostureControllerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/DeviceProvisionedControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/DeviceProvisionedControllerImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/DeviceProvisionedControllerImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/DeviceProvisionedControllerImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ExtensionControllerImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ExtensionControllerImplTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ExtensionControllerImplTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ExtensionControllerImplTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/HotspotControllerImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/HotspotControllerImplTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/HotspotControllerImplTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/HotspotControllerImplTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyguardQsUserSwitchControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/KeyguardQsUserSwitchControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyguardQsUserSwitchControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/KeyguardQsUserSwitchControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyguardStateControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/KeyguardStateControllerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyguardStateControllerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/KeyguardStateControllerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyguardUserSwitcherAdapterTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/KeyguardUserSwitcherAdapterTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyguardUserSwitcherAdapterTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/KeyguardUserSwitcherAdapterTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputQuickSettingsDisablerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/RemoteInputQuickSettingsDisablerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputQuickSettingsDisablerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/RemoteInputQuickSettingsDisablerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/SafetyControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyConstantsTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/SmartReplyConstantsTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyConstantsTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/SmartReplyConstantsTest.java
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.kt
new file mode 100644
index 0000000..9abdf42
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.kt
@@ -0,0 +1,207 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+package com.android.systemui.statusbar.policy
+
+import android.app.NotificationManager
+import android.os.Handler
+import android.provider.Settings
+import android.service.notification.ZenModeConfig
+import android.testing.TestableLooper
+import android.testing.TestableLooper.RunWithLooper
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.broadcast.broadcastDispatcher
+import com.android.systemui.dump.dumpManager
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.settings.userTracker
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeGlobalSettings
+import com.google.common.truth.Truth.assertThat
+import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicInteger
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers
+import org.mockito.Mockito
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@RunWithLooper
+class ZenModeControllerImplTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private lateinit var testableLooper: TestableLooper
+    private lateinit var controller: ZenModeControllerImpl
+
+    private val globalSettings = kosmos.fakeGlobalSettings
+    private val config: ZenModeConfig = mock<ZenModeConfig>()
+    private val mNm: NotificationManager = mock<NotificationManager>()
+
+    @Before
+    fun setUp() {
+        testableLooper = TestableLooper.get(this)
+        mContext.addMockSystemService(NotificationManager::class.java, mNm)
+        whenever(mNm.zenModeConfig).thenReturn(config)
+
+        controller =
+            ZenModeControllerImpl(
+                mContext,
+                Handler.createAsync(testableLooper.looper),
+                kosmos.broadcastDispatcher,
+                kosmos.dumpManager,
+                globalSettings,
+                kosmos.userTracker,
+            )
+    }
+
+    @Test
+    fun testRemoveDuringCallback() {
+        val callback =
+            object : ZenModeController.Callback {
+                override fun onConfigChanged(config: ZenModeConfig) {
+                    controller.removeCallback(this)
+                }
+            }
+
+        controller.addCallback(callback)
+        val mockCallback = Mockito.mock(ZenModeController.Callback::class.java)
+        controller.addCallback(mockCallback)
+        controller.fireConfigChanged(config)
+        Mockito.verify(mockCallback).onConfigChanged(ArgumentMatchers.eq(config))
+    }
+
+    @Test
+    fun testAreNotificationsHiddenInShade_zenOffShadeSuppressed() {
+        config.suppressedVisualEffects =
+            NotificationManager.Policy.SUPPRESSED_EFFECT_NOTIFICATION_LIST
+        controller.updateZenMode(Settings.Global.ZEN_MODE_OFF)
+        controller.updateZenModeConfig()
+        assertThat(controller.areNotificationsHiddenInShade()).isFalse()
+    }
+
+    @Test
+    fun testAreNotificationsHiddenInShade_zenOnShadeNotSuppressed() {
+        val policy =
+            NotificationManager.Policy(
+                0,
+                0,
+                0,
+                NotificationManager.Policy.SUPPRESSED_EFFECT_STATUS_BAR,
+            )
+        whenever(mNm.consolidatedNotificationPolicy).thenReturn(policy)
+        controller.updateConsolidatedNotificationPolicy()
+        controller.updateZenMode(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS)
+        assertThat(controller.areNotificationsHiddenInShade()).isFalse()
+    }
+
+    @Test
+    fun testAreNotificationsHiddenInShade_zenOnShadeSuppressed() {
+        val policy =
+            NotificationManager.Policy(
+                0,
+                0,
+                0,
+                NotificationManager.Policy.SUPPRESSED_EFFECT_NOTIFICATION_LIST,
+            )
+        whenever(mNm.consolidatedNotificationPolicy).thenReturn(policy)
+        controller.updateConsolidatedNotificationPolicy()
+        controller.updateZenMode(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS)
+        assertThat(controller.areNotificationsHiddenInShade()).isTrue()
+    }
+
+    @Test
+    fun testModeChange() =
+        testScope.runTest {
+            val states =
+                listOf(
+                    Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS,
+                    Settings.Global.ZEN_MODE_NO_INTERRUPTIONS,
+                    Settings.Global.ZEN_MODE_ALARMS,
+                    Settings.Global.ZEN_MODE_ALARMS,
+                )
+
+            for (state in states) {
+                globalSettings.putInt(Settings.Global.ZEN_MODE, state)
+                testScope.runCurrent()
+                testableLooper.processAllMessages()
+                assertThat(controller.zen).isEqualTo(state)
+            }
+        }
+
+    @Test
+    fun testModeChange_callbackNotified() =
+        testScope.runTest {
+            val currentState = AtomicInteger(-1)
+
+            val callback: ZenModeController.Callback =
+                object : ZenModeController.Callback {
+                    override fun onZenChanged(zen: Int) {
+                        currentState.set(zen)
+                    }
+                }
+
+            controller.addCallback(callback)
+
+            val states =
+                listOf(
+                    Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS,
+                    Settings.Global.ZEN_MODE_NO_INTERRUPTIONS,
+                    Settings.Global.ZEN_MODE_ALARMS,
+                    Settings.Global.ZEN_MODE_ALARMS,
+                )
+
+            for (state in states) {
+                globalSettings.putInt(Settings.Global.ZEN_MODE, state)
+                testScope.runCurrent()
+                testableLooper.processAllMessages()
+                assertThat(currentState.get()).isEqualTo(state)
+            }
+        }
+
+    @Test
+    fun testCallbackRemovedWhileDispatching_doesntCrash() =
+        testScope.runTest {
+            val remove = AtomicBoolean(false)
+            globalSettings.putInt(Settings.Global.ZEN_MODE, Settings.Global.ZEN_MODE_OFF)
+            testableLooper.processAllMessages()
+            val callback: ZenModeController.Callback =
+                object : ZenModeController.Callback {
+                    override fun onZenChanged(zen: Int) {
+                        if (remove.get()) {
+                            controller.removeCallback(this)
+                        }
+                    }
+                }
+            controller.addCallback(callback)
+            controller.addCallback(object : ZenModeController.Callback {})
+
+            remove.set(true)
+
+            globalSettings.putInt(
+                Settings.Global.ZEN_MODE,
+                Settings.Global.ZEN_MODE_NO_INTERRUPTIONS,
+            )
+            testScope.runCurrent()
+            testableLooper.processAllMessages()
+        }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/bluetooth/BluetoothRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/bluetooth/BluetoothRepositoryImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/bluetooth/BluetoothRepositoryImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/bluetooth/BluetoothRepositoryImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/bluetooth/FakeBluetoothRepository.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/bluetooth/FakeBluetoothRepository.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/bluetooth/FakeBluetoothRepository.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/bluetooth/FakeBluetoothRepository.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/data/repository/DeviceProvisioningRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/data/repository/DeviceProvisioningRepositoryImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/data/repository/DeviceProvisioningRepositoryImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/data/repository/DeviceProvisioningRepositoryImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/StatusBarWindowStateControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/StatusBarWindowStateControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/window/StatusBarWindowStateControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/StatusBarWindowStateControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStatePerDisplayRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStatePerDisplayRepositoryTest.kt
similarity index 97%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStatePerDisplayRepositoryTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStatePerDisplayRepositoryTest.kt
index 0c27e58..c7c7fdc 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStatePerDisplayRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStatePerDisplayRepositoryTest.kt
@@ -21,6 +21,7 @@
 import android.app.StatusBarManager.WINDOW_STATE_HIDING
 import android.app.StatusBarManager.WINDOW_STATE_SHOWING
 import android.app.StatusBarManager.WINDOW_STATUS_BAR
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -34,11 +35,13 @@
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
+import org.junit.runner.RunWith
 import org.mockito.Mockito.verify
 import org.mockito.kotlin.argumentCaptor
 
 @SmallTest
 @OptIn(ExperimentalCoroutinesApi::class)
+@RunWith(AndroidJUnit4::class)
 class StatusBarWindowStatePerDisplayRepositoryTest : SysuiTestCase() {
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryStoreTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryStoreTest.kt
similarity index 97%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryStoreTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryStoreTest.kt
index b6a3f36..e23e88c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryStoreTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryStoreTest.kt
@@ -19,6 +19,7 @@
 import android.app.StatusBarManager.WINDOW_STATE_HIDDEN
 import android.app.StatusBarManager.WINDOW_STATE_SHOWING
 import android.app.StatusBarManager.WINDOW_STATUS_BAR
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -33,12 +34,14 @@
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import org.junit.runner.RunWith
 import org.mockito.Mockito.verify
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.reset
 
 @SmallTest
 @OptIn(ExperimentalCoroutinesApi::class)
+@RunWith(AndroidJUnit4::class)
 class StatusBarWindowStateRepositoryStoreTest : SysuiTestCase() {
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/glowboxeffect/GlowBoxEffectTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/glowboxeffect/GlowBoxEffectTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/glowboxeffect/GlowBoxEffectTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/glowboxeffect/GlowBoxEffectTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/loadingeffect/LoadingEffectTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/loadingeffect/LoadingEffectTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/loadingeffect/LoadingEffectTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/loadingeffect/LoadingEffectTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/MultiRippleControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/MultiRippleControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/MultiRippleControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/MultiRippleControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/RippleAnimationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/RippleAnimationTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/RippleAnimationTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/RippleAnimationTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/RippleShaderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/RippleShaderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/RippleShaderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/RippleShaderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/RippleViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/RippleViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/ripple/RippleViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/ripple/RippleViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/shaders/SolidColorShaderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/shaders/SolidColorShaderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/shaders/SolidColorShaderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/shaders/SolidColorShaderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/shaders/SparkleShaderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/shaders/SparkleShaderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/shaders/SparkleShaderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/shaders/SparkleShaderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseShaderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseViewTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseViewTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseViewTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/surfaceeffects/turbulencenoise/TurbulenceNoiseViewTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/TouchableRegionViewControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/temporarydisplay/TouchableRegionViewControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/temporarydisplay/TouchableRegionViewControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/temporarydisplay/TouchableRegionViewControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/UnfoldLatencyTrackerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/UnfoldLatencyTrackerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/unfold/UnfoldLatencyTrackerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/UnfoldLatencyTrackerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/progress/PhysicsBasedUnfoldTransitionProgressProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/progress/UnfoldRemoteFilterTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/progress/UnfoldRemoteFilterTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/unfold/progress/UnfoldRemoteFilterTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/progress/UnfoldRemoteFilterTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/util/ScaleAwareUnfoldProgressProviderTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/util/ScaleAwareUnfoldProgressProviderTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/unfold/util/ScaleAwareUnfoldProgressProviderTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/unfold/util/ScaleAwareUnfoldProgressProviderTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/usb/UsbPermissionActivityTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/usb/UsbPermissionActivityTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/usb/UsbPermissionActivityTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/usb/UsbPermissionActivityTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/CreateUserActivityTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/user/CreateUserActivityTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/user/CreateUserActivityTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/user/CreateUserActivityTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/UserCreatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/user/UserCreatorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/user/UserCreatorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/user/UserCreatorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/GuestUserInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/user/domain/interactor/GuestUserInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/GuestUserInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/user/domain/interactor/GuestUserInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/RefreshUsersSchedulerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/user/domain/interactor/RefreshUsersSchedulerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/RefreshUsersSchedulerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/user/domain/interactor/RefreshUsersSchedulerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/user/domain/interactor/UserSwitcherInteractorTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/NotificationChannelsTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/NotificationChannelsTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/util/NotificationChannelsTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/util/NotificationChannelsTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/kotlin/FlowUtilTests.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/kotlin/FlowUtilTests.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/util/kotlin/FlowUtilTests.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/util/kotlin/FlowUtilTests.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/kotlin/PackageManagerExtComponentEnabledTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/kotlin/PackageManagerExtComponentEnabledTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/util/kotlin/PackageManagerExtComponentEnabledTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/util/kotlin/PackageManagerExtComponentEnabledTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/leak/LeakDetectorTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/leak/LeakDetectorTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/util/leak/LeakDetectorTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/util/leak/LeakDetectorTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/leak/ReferenceTestUtilsTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/leak/ReferenceTestUtilsTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/util/leak/ReferenceTestUtilsTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/util/leak/ReferenceTestUtilsTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/sensors/AsyncSensorManagerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/util/sensors/AsyncSensorManagerTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/util/sensors/AsyncSensorManagerTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/util/sensors/AsyncSensorManagerTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/CsdWarningDialogTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/CsdWarningDialogTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/volume/CsdWarningDialogTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/volume/CsdWarningDialogTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTestKt.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/VolumeDialogControllerImplTestKt.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTestKt.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/volume/VolumeDialogControllerImplTestKt.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/wallpapers/data/repository/FakeWallpaperRepository.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/wallpapers/data/repository/FakeWallpaperRepository.kt
index fe5024f..59676ce 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/wallpapers/data/repository/FakeWallpaperRepository.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/wallpapers/data/repository/FakeWallpaperRepository.kt
@@ -17,10 +17,17 @@
 package com.android.systemui.wallpapers.data.repository
 
 import android.app.WallpaperInfo
+import android.view.View
 import kotlinx.coroutines.flow.MutableStateFlow
 
 /** Fake implementation of the wallpaper repository. */
 class FakeWallpaperRepository : WallpaperRepository {
     override val wallpaperInfo = MutableStateFlow<WallpaperInfo?>(null)
     override val wallpaperSupportsAmbientMode = MutableStateFlow(false)
+    override var rootView: View? = null
+    private val _notificationStackAbsoluteBottom = MutableStateFlow(0F)
+
+    override fun setNotificationStackAbsoluteBottom(bottom: Float) {
+        _notificationStackAbsoluteBottom.value = bottom
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubbleEducationControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/wmshell/BubbleEducationControllerTest.kt
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/wmshell/BubbleEducationControllerTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/wmshell/BubbleEducationControllerTest.kt
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTestActivity.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/wmshell/BubblesTestActivity.java
similarity index 100%
rename from packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTestActivity.java
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/wmshell/BubblesTestActivity.java
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
index 8e01e04..431f048 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
@@ -315,6 +315,7 @@
     };
     private final FaceWakeUpTriggersConfig mFaceWakeUpTriggersConfig;
 
+    private final Object mSimDataLockObject = new Object();
     HashMap<Integer, SimData> mSimDatas = new HashMap<>();
     HashMap<Integer, ServiceState> mServiceStates = new HashMap<>();
 
@@ -610,14 +611,26 @@
 
         // It is possible for active subscriptions to become invalid (-1), and these will not be
         // present in the subscriptionInfo list
-        Iterator<Map.Entry<Integer, SimData>> iter = mSimDatas.entrySet().iterator();
-        while (iter.hasNext()) {
-            Map.Entry<Integer, SimData> simData = iter.next();
-            if (!activeSubIds.contains(simData.getKey())) {
-                mSimLogger.logInvalidSubId(simData.getKey());
-                iter.remove();
+        synchronized (mSimDataLockObject) {
+            Iterator<Map.Entry<Integer, SimData>> iter = mSimDatas.entrySet().iterator();
+            while (iter.hasNext()) {
+                Map.Entry<Integer, SimData> simData = iter.next();
+                if (!activeSubIds.contains(simData.getKey())) {
+                    mSimLogger.logInvalidSubId(simData.getKey());
+                    iter.remove();
 
-                SimData data = simData.getValue();
+                    SimData data = simData.getValue();
+                    for (int j = 0; j < mCallbacks.size(); j++) {
+                        KeyguardUpdateMonitorCallback cb = mCallbacks.get(j).get();
+                        if (cb != null) {
+                            cb.onSimStateChanged(data.subId, data.slotId, data.simState);
+                        }
+                    }
+                }
+            }
+
+            for (int i = 0; i < changedSubscriptions.size(); i++) {
+                SimData data = mSimDatas.get(changedSubscriptions.get(i).getSubscriptionId());
                 for (int j = 0; j < mCallbacks.size(); j++) {
                     KeyguardUpdateMonitorCallback cb = mCallbacks.get(j).get();
                     if (cb != null) {
@@ -625,18 +638,8 @@
                     }
                 }
             }
+            callbacksRefreshCarrierInfo();
         }
-
-        for (int i = 0; i < changedSubscriptions.size(); i++) {
-            SimData data = mSimDatas.get(changedSubscriptions.get(i).getSubscriptionId());
-            for (int j = 0; j < mCallbacks.size(); j++) {
-                KeyguardUpdateMonitorCallback cb = mCallbacks.get(j).get();
-                if (cb != null) {
-                    cb.onSimStateChanged(data.subId, data.slotId, data.simState);
-                }
-            }
-        }
-        callbacksRefreshCarrierInfo();
     }
 
     private void handleAirplaneModeChanged() {
@@ -3376,12 +3379,15 @@
      * Removes all valid subscription info from the map for the given slotId.
      */
     private void invalidateSlot(int slotId) {
-        Iterator<Map.Entry<Integer, SimData>> iter = mSimDatas.entrySet().iterator();
-        while (iter.hasNext()) {
-            SimData data = iter.next().getValue();
-            if (data.slotId == slotId && SubscriptionManager.isValidSubscriptionId(data.subId)) {
-                mSimLogger.logInvalidSubId(data.subId);
-                iter.remove();
+        synchronized (mSimDataLockObject) {
+            Iterator<Map.Entry<Integer, SimData>> iter = mSimDatas.entrySet().iterator();
+            while (iter.hasNext()) {
+                SimData data = iter.next().getValue();
+                if (data.slotId == slotId
+                        && SubscriptionManager.isValidSubscriptionId(data.subId)) {
+                    mSimLogger.logInvalidSubId(data.subId);
+                    iter.remove();
+                }
             }
         }
     }
@@ -3408,23 +3414,25 @@
         }
 
         // TODO(b/327476182): Preserve SIM_STATE_CARD_IO_ERROR sims in a separate data source.
-        SimData data = mSimDatas.get(subId);
-        final boolean changed;
-        if (data == null) {
-            data = new SimData(state, slotId, subId);
-            mSimDatas.put(subId, data);
-            changed = true; // no data yet; force update
-        } else {
-            changed = (data.simState != state || data.subId != subId || data.slotId != slotId);
-            data.simState = state;
-            data.subId = subId;
-            data.slotId = slotId;
-        }
-        if ((changed || becameAbsent)) {
-            for (int i = 0; i < mCallbacks.size(); i++) {
-                KeyguardUpdateMonitorCallback cb = mCallbacks.get(i).get();
-                if (cb != null) {
-                    cb.onSimStateChanged(subId, slotId, state);
+        synchronized (mSimDataLockObject) {
+            SimData data = mSimDatas.get(subId);
+            final boolean changed;
+            if (data == null) {
+                data = new SimData(state, slotId, subId);
+                mSimDatas.put(subId, data);
+                changed = true; // no data yet; force update
+            } else {
+                changed = (data.simState != state || data.subId != subId || data.slotId != slotId);
+                data.simState = state;
+                data.subId = subId;
+                data.slotId = slotId;
+            }
+            if ((changed || becameAbsent)) {
+                for (int i = 0; i < mCallbacks.size(); i++) {
+                    KeyguardUpdateMonitorCallback cb = mCallbacks.get(i).get();
+                    if (cb != null) {
+                        cb.onSimStateChanged(subId, slotId, state);
+                    }
                 }
             }
         }
@@ -3684,9 +3692,11 @@
         callback.onKeyguardVisibilityChanged(isKeyguardVisible());
         callback.onTelephonyCapable(mTelephonyCapable);
 
-        for (Entry<Integer, SimData> data : mSimDatas.entrySet()) {
-            final SimData state = data.getValue();
-            callback.onSimStateChanged(state.subId, state.slotId, state.simState);
+        synchronized (mSimDataLockObject) {
+            for (Entry<Integer, SimData> data : mSimDatas.entrySet()) {
+                final SimData state = data.getValue();
+                callback.onSimStateChanged(state.subId, state.slotId, state.simState);
+            }
         }
     }
 
@@ -3823,19 +3833,23 @@
     }
 
     public int getSimState(int subId) {
-        if (mSimDatas.containsKey(subId)) {
-            return mSimDatas.get(subId).simState;
-        } else {
-            return TelephonyManager.SIM_STATE_UNKNOWN;
+        synchronized (mSimDataLockObject) {
+            if (mSimDatas.containsKey(subId)) {
+                return mSimDatas.get(subId).simState;
+            } else {
+                return TelephonyManager.SIM_STATE_UNKNOWN;
+            }
         }
     }
 
     private int getSlotId(int subId) {
-        if (!mSimDatas.containsKey(subId)) {
-            refreshSimState(subId, SubscriptionManager.getSlotIndex(subId));
+        synchronized (mSimDataLockObject) {
+            if (!mSimDatas.containsKey(subId)) {
+                refreshSimState(subId, SubscriptionManager.getSlotIndex(subId));
+            }
+            SimData simData = mSimDatas.get(subId);
+            return simData != null ? simData.slotId : SubscriptionManager.INVALID_SUBSCRIPTION_ID;
         }
-        SimData simData = mSimDatas.get(subId);
-        return simData != null ? simData.slotId : SubscriptionManager.INVALID_SUBSCRIPTION_ID;
     }
 
     private final TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
@@ -3876,22 +3890,24 @@
      */
     private boolean refreshSimState(int subId, int slotId) {
         int state = mTelephonyManager.getSimState(slotId);
-        SimData data = mSimDatas.get(subId);
+        synchronized (mSimDataLockObject) {
+            SimData data = mSimDatas.get(subId);
 
-        if (!SubscriptionManager.isValidSubscriptionId(subId)) {
-            invalidateSlot(slotId);
-        }
+            if (!SubscriptionManager.isValidSubscriptionId(subId)) {
+                invalidateSlot(slotId);
+            }
 
-        final boolean changed;
-        if (data == null) {
-            data = new SimData(state, slotId, subId);
-            mSimDatas.put(subId, data);
-            changed = true; // no data yet; force update
-        } else {
-            changed = data.simState != state;
-            data.simState = state;
+            final boolean changed;
+            if (data == null) {
+                data = new SimData(state, slotId, subId);
+                mSimDatas.put(subId, data);
+                changed = true; // no data yet; force update
+            } else {
+                changed = data.simState != state;
+                data.simState = state;
+            }
+            return changed;
         }
-        return changed;
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt b/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
index 462e820..b56ed8c 100644
--- a/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/DisplayModule.kt
@@ -16,16 +16,25 @@
 
 package com.android.systemui.display
 
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.display.data.repository.DeviceStateRepository
 import com.android.systemui.display.data.repository.DeviceStateRepositoryImpl
 import com.android.systemui.display.data.repository.DisplayRepository
 import com.android.systemui.display.data.repository.DisplayRepositoryImpl
+import com.android.systemui.display.data.repository.DisplayScopeRepository
+import com.android.systemui.display.data.repository.DisplayScopeRepositoryImpl
 import com.android.systemui.display.data.repository.FocusedDisplayRepository
 import com.android.systemui.display.data.repository.FocusedDisplayRepositoryImpl
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractorImpl
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
 import dagger.Binds
+import dagger.Lazy
 import dagger.Module
+import dagger.Provides
+import dagger.multibindings.ClassKey
+import dagger.multibindings.IntoMap
 
 /** Module binding display related classes. */
 @Module
@@ -46,4 +55,22 @@
     fun bindsFocusedDisplayRepository(
         focusedDisplayRepository: FocusedDisplayRepositoryImpl
     ): FocusedDisplayRepository
+
+    @Binds fun displayScopeRepository(impl: DisplayScopeRepositoryImpl): DisplayScopeRepository
+
+    companion object {
+        @Provides
+        @SysUISingleton
+        @IntoMap
+        @ClassKey(DisplayScopeRepositoryImpl::class)
+        fun displayScopeRepoCoreStartable(
+            repoImplLazy: Lazy<DisplayScopeRepositoryImpl>
+        ): CoreStartable {
+            return if (StatusBarConnectedDisplays.isEnabled) {
+                repoImplLazy.get()
+            } else {
+                CoreStartable.NOP
+            }
+        }
+    }
 }
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 e68aba5..6a69136 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
@@ -61,8 +61,11 @@
     /** Display addition event indicating a new display has been added. */
     val displayAdditionEvent: Flow<Display?>
 
+    /** Display removal event indicating a display has been removed. */
+    val displayRemovalEvent: Flow<Int>
+
     /** Provides the current set of displays. */
-    val displays: Flow<Set<Display>>
+    val displays: StateFlow<Set<Display>>
 
     /**
      * Pending display id that can be enabled/disabled.
@@ -79,8 +82,8 @@
      *
      * This method is guaranteed to not result in any binder call.
      */
-    suspend fun getDisplay(displayId: Int): Display? =
-        displays.first().firstOrNull { it.displayId == displayId }
+    fun getDisplay(displayId: Int): Display? =
+        displays.value.firstOrNull { it.displayId == displayId }
 
     /** Represents a connected display that has not been enabled yet. */
     interface PendingDisplay {
@@ -148,6 +151,9 @@
             getDisplayFromDisplayManager(it.displayId)
         }
 
+    override val displayRemovalEvent: Flow<Int> =
+        allDisplayEvents.filterIsInstance<DisplayEvent.Removed>().map { it.displayId }
+
     // 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).
@@ -180,7 +186,7 @@
      *
      * Those are commonly the ones provided by [DisplayManager.getDisplays] by default.
      */
-    private val enabledDisplays: Flow<Set<Display>> =
+    private val enabledDisplays: StateFlow<Set<Display>> =
         enabledDisplayIds
             .mapElementsLazily { displayId -> getDisplayFromDisplayManager(displayId) }
             .onEach {
@@ -204,7 +210,7 @@
      *
      * Those are commonly the ones provided by [DisplayManager.getDisplays] by default.
      */
-    override val displays: Flow<Set<Display>> = enabledDisplays
+    override val displays: StateFlow<Set<Display>> = enabledDisplays
 
     val _ignoredDisplayIds = MutableStateFlow<Set<Int>>(emptySet())
     private val ignoredDisplayIds: Flow<Set<Int>> = _ignoredDisplayIds.debugLog("ignoredDisplayIds")
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayScopeRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayScopeRepository.kt
new file mode 100644
index 0000000..3062475
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayScopeRepository.kt
@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.display.data.repository
+
+import android.view.Display
+import com.android.app.tracing.coroutines.createCoroutineTracingContext
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
+import java.util.concurrent.ConcurrentHashMap
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
+
+/**
+ * Provides per display instances of [CoroutineScope]. These will remain active as long as the
+ * display is connected, and automatically cancelled when the display is removed.
+ */
+interface DisplayScopeRepository {
+    fun scopeForDisplay(displayId: Int): CoroutineScope
+}
+
+@SysUISingleton
+class DisplayScopeRepositoryImpl
+@Inject
+constructor(
+    @Background private val backgroundApplicationScope: CoroutineScope,
+    @Background private val backgroundDispatcher: CoroutineDispatcher,
+    private val displayRepository: DisplayRepository,
+) : DisplayScopeRepository, CoreStartable {
+
+    private val perDisplayScopes = ConcurrentHashMap<Int, CoroutineScope>()
+
+    override fun scopeForDisplay(displayId: Int): CoroutineScope {
+        return perDisplayScopes.computeIfAbsent(displayId) { createScopeForDisplay(displayId) }
+    }
+
+    override fun start() {
+        StatusBarConnectedDisplays.assertInNewMode()
+        backgroundApplicationScope.launch {
+            displayRepository.displayRemovalEvent.collect { displayId ->
+                val scope = perDisplayScopes.remove(displayId)
+                scope?.cancel("Display $displayId has been removed.")
+            }
+        }
+    }
+
+    private fun createScopeForDisplay(displayId: Int): CoroutineScope {
+        return if (displayId == Display.DEFAULT_DISPLAY) {
+            // The default display is connected all the time, therefore we can optimise by reusing
+            // the application scope, and don't need to create a new scope.
+            backgroundApplicationScope
+        } else {
+            CoroutineScope(
+                backgroundDispatcher + createCoroutineTracingContext("DisplayScope$displayId")
+            )
+        }
+    }
+}
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 5e05dab..b2acc2a 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
@@ -24,6 +24,7 @@
 import androidx.compose.animation.core.animateFloatAsState
 import androidx.compose.foundation.Image
 import androidx.compose.foundation.background
+import androidx.compose.foundation.border
 import androidx.compose.foundation.focusable
 import androidx.compose.foundation.interaction.MutableInteractionSource
 import androidx.compose.foundation.interaction.collectIsFocusedAsState
@@ -80,25 +81,18 @@
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.draw.drawWithContent
 import androidx.compose.ui.focus.FocusDirection
 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.Rect
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.RectangleShape
 import androidx.compose.ui.graphics.Shape
-import androidx.compose.ui.graphics.drawscope.Stroke
 import androidx.compose.ui.graphics.graphicsLayer
 import androidx.compose.ui.input.key.Key
 import androidx.compose.ui.input.key.key
 import androidx.compose.ui.input.key.onKeyEvent
-import androidx.compose.ui.input.nestedscroll.nestedScroll
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalFocusManager
-import androidx.compose.ui.platform.rememberNestedScrollInteropConnection
 import androidx.compose.ui.res.painterResource
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.semantics.Role
@@ -113,9 +107,9 @@
 import androidx.compose.ui.util.fastFirstOrNull
 import androidx.compose.ui.util.fastForEach
 import androidx.compose.ui.util.fastForEachIndexed
-import androidx.compose.ui.zIndex
+import com.android.compose.modifiers.thenIf
 import com.android.compose.ui.graphics.painter.rememberDrawablePainter
-import com.android.systemui.keyboard.shortcut.shared.model.Shortcut
+import com.android.systemui.keyboard.shortcut.shared.model.Shortcut as ShortcutModel
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategory
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCategoryType
 import com.android.systemui.keyboard.shortcut.shared.model.ShortcutCommand
@@ -405,7 +399,7 @@
         if (index > 0) {
             HorizontalDivider(color = MaterialTheme.colorScheme.surfaceContainerHigh)
         }
-        ShortcutView(Modifier.padding(vertical = 24.dp), searchQuery, shortcut)
+        Shortcut(Modifier.padding(vertical = 24.dp), searchQuery, shortcut)
     }
 }
 
@@ -444,9 +438,9 @@
         NoSearchResultsText(horizontalPadding = 24.dp, fillHeight = false)
         return
     }
-    LazyColumn(modifier.nestedScroll(rememberNestedScrollInteropConnection())) {
-        items(items = category.subCategories, key = { item -> item.label }) {
-            SubCategoryContainerDualPane(searchQuery, it)
+    LazyColumn(modifier = modifier) {
+        items(category.subCategories) { subcategory ->
+            SubCategoryContainerDualPane(searchQuery = searchQuery, subCategory = subcategory)
             Spacer(modifier = Modifier.height(8.dp))
         }
     }
@@ -477,14 +471,21 @@
         shape = RoundedCornerShape(28.dp),
         color = MaterialTheme.colorScheme.surfaceBright,
     ) {
-        Column(Modifier.padding(24.dp)) {
+        Column(Modifier.padding(16.dp)) {
             SubCategoryTitle(subCategory.label)
             Spacer(Modifier.height(8.dp))
             subCategory.shortcuts.fastForEachIndexed { index, shortcut ->
                 if (index > 0) {
-                    HorizontalDivider(color = MaterialTheme.colorScheme.surfaceContainerHigh)
+                    HorizontalDivider(
+                        modifier = Modifier.padding(horizontal = 8.dp),
+                        color = MaterialTheme.colorScheme.surfaceContainerHigh,
+                    )
                 }
-                ShortcutView(Modifier.padding(vertical = 16.dp), searchQuery, shortcut)
+                Shortcut(
+                    modifier = Modifier.padding(vertical = 8.dp),
+                    searchQuery = searchQuery,
+                    shortcut = shortcut,
+                )
             }
         }
     }
@@ -500,18 +501,17 @@
 }
 
 @Composable
-private fun ShortcutView(modifier: Modifier, searchQuery: String, shortcut: Shortcut) {
+private fun Shortcut(modifier: Modifier, searchQuery: String, shortcut: ShortcutModel) {
     val interactionSource = remember { MutableInteractionSource() }
     val isFocused by interactionSource.collectIsFocusedAsState()
+    val focusColor = MaterialTheme.colorScheme.secondary
     Row(
         modifier
+            .thenIf(isFocused) {
+                Modifier.border(width = 3.dp, color = focusColor, shape = RoundedCornerShape(16.dp))
+            }
             .focusable(interactionSource = interactionSource)
-            .outlineFocusModifier(
-                isFocused = isFocused,
-                focusColor = MaterialTheme.colorScheme.secondary,
-                padding = 8.dp,
-                cornerRadius = 16.dp,
-            )
+            .padding(8.dp)
     ) {
         Row(
             modifier = Modifier.width(128.dp).align(Alignment.CenterVertically),
@@ -523,7 +523,7 @@
             }
             ShortcutDescriptionText(searchQuery = searchQuery, shortcut = shortcut)
         }
-        Spacer(modifier = Modifier.width(16.dp))
+        Spacer(modifier = Modifier.width(24.dp))
         ShortcutKeyCombinations(modifier = Modifier.weight(1f), shortcut = shortcut)
     }
 }
@@ -548,7 +548,7 @@
 
 @OptIn(ExperimentalLayoutApi::class)
 @Composable
-private fun ShortcutKeyCombinations(modifier: Modifier = Modifier, shortcut: Shortcut) {
+private fun ShortcutKeyCombinations(modifier: Modifier = Modifier, shortcut: ShortcutModel) {
     FlowRow(
         modifier = modifier,
         verticalArrangement = Arrangement.spacedBy(8.dp),
@@ -628,7 +628,7 @@
 @Composable
 private fun ShortcutDescriptionText(
     searchQuery: String,
-    shortcut: Shortcut,
+    shortcut: ShortcutModel,
     modifier: Modifier = Modifier,
 ) {
     Text(
@@ -751,39 +751,6 @@
     }
 }
 
-private fun Modifier.outlineFocusModifier(
-    isFocused: Boolean,
-    focusColor: Color,
-    padding: Dp,
-    cornerRadius: Dp,
-): Modifier {
-    if (isFocused) {
-        return this.drawWithContent {
-                val focusOutline =
-                    Rect(Offset.Zero, size).let {
-                        if (padding > 0.dp) {
-                            it.inflate(padding.toPx())
-                        } else {
-                            it.deflate(padding.unaryMinus().toPx())
-                        }
-                    }
-                drawContent()
-                drawRoundRect(
-                    color = focusColor,
-                    style = Stroke(width = 3.dp.toPx()),
-                    topLeft = focusOutline.topLeft,
-                    size = focusOutline.size,
-                    cornerRadius = CornerRadius(cornerRadius.toPx()),
-                )
-            }
-            // Increasing Z-Index so focus outline is drawn on top of "selected" category
-            // background.
-            .zIndex(1f)
-    } else {
-        return this
-    }
-}
-
 @Composable
 @OptIn(ExperimentalMaterial3Api::class)
 private fun TitleBar() {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 9c7cc81..2052459 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -178,6 +178,7 @@
 import com.android.systemui.util.settings.SystemSettings;
 import com.android.systemui.util.time.SystemClock;
 import com.android.systemui.wallpapers.data.repository.WallpaperRepository;
+import com.android.window.flags.Flags;
 import com.android.wm.shell.keyguard.KeyguardTransitions;
 
 import dagger.Lazy;
@@ -236,6 +237,9 @@
  */
 public class KeyguardViewMediator implements CoreStartable, Dumpable,
         StatusBarStateController.StateListener {
+
+    private static final boolean ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS =
+            Flags.ensureKeyguardDoesTransitionStarting();
     private static final int KEYGUARD_DISPLAY_TIMEOUT_DELAY_DEFAULT = 30000;
     private static final long KEYGUARD_DONE_PENDING_TIMEOUT_MS = 3000;
 
@@ -2865,9 +2869,14 @@
                 return;
             }
 
-            try {
-                mActivityTaskManagerService.setLockScreenShown(showing, aodShowing);
-            } catch (RemoteException ignored) {
+            if (ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS) {
+                mKeyguardTransitions.startKeyguardTransition(showing, aodShowing);
+            } else {
+                try {
+
+                    mActivityTaskManagerService.setLockScreenShown(showing, aodShowing);
+                } catch (RemoteException ignored) {
+                }
             }
         });
     }
@@ -2998,18 +3007,23 @@
 
             // Handled in WmLockscreenVisibilityManager if flag is enabled.
             if (!KeyguardWmStateRefactor.isEnabled()) {
-                    // Don't actually hide the Keyguard at the moment, wait for window manager 
-                    // until it tells us it's safe to do so with startKeyguardExitAnimation.
-		    // Posting to mUiOffloadThread to ensure that calls to ActivityTaskManager 
-		    // will be in order.
-		    final int keyguardFlag = flags;
-		    mUiBgExecutor.execute(() -> {
-		        try {
-		            mActivityTaskManagerService.keyguardGoingAway(keyguardFlag);
-		        } catch (RemoteException e) {
-		            Log.e(TAG, "Error while calling WindowManager", e);
-		        }
-		    });
+                // Don't actually hide the Keyguard at the moment, wait for window manager
+                // until it tells us it's safe to do so with startKeyguardExitAnimation.
+                // Posting to mUiOffloadThread to ensure that calls to ActivityTaskManager
+                // will be in order.
+                final int keyguardFlag = flags;
+                mUiBgExecutor.execute(() -> {
+                    if (ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS) {
+                        mKeyguardTransitions.startKeyguardTransition(
+                                false /* keyguardShowing */, false /* aodShowing */);
+                        return;
+                    }
+                    try {
+                        mActivityTaskManagerService.keyguardGoingAway(keyguardFlag);
+                    } catch (RemoteException e) {
+                        Log.e(TAG, "Error while calling WindowManager", e);
+                    }
+                });
             }
 
             Trace.endSection();
@@ -3464,6 +3478,12 @@
     public void showSurfaceBehindKeyguard() {
         mSurfaceBehindRemoteAnimationRequested = true;
 
+        if (ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS) {
+            mKeyguardTransitions.startKeyguardTransition(
+                    false /* keyguardShowing */, false /* aodShowing */);
+            return;
+        }
+
         try {
             int flags = KEYGUARD_GOING_AWAY_FLAG_NO_WINDOW_ANIMATIONS
                     | KEYGUARD_GOING_AWAY_FLAG_WITH_WALLPAPER;
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt b/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt
index e89594e..032af94 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt
@@ -26,6 +26,8 @@
 import com.android.systemui.keyguard.domain.interactor.KeyguardDismissTransitionInteractor
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindParamsApplier
 import com.android.systemui.statusbar.policy.KeyguardStateController
+import com.android.window.flags.Flags
+import com.android.wm.shell.keyguard.KeyguardTransitions
 import java.util.concurrent.Executor
 import javax.inject.Inject
 
@@ -42,6 +44,7 @@
     private val keyguardStateController: KeyguardStateController,
     private val keyguardSurfaceBehindAnimator: KeyguardSurfaceBehindParamsApplier,
     private val keyguardDismissTransitionInteractor: KeyguardDismissTransitionInteractor,
+    private val keyguardTransitions: KeyguardTransitions
 ) {
 
     /**
@@ -97,6 +100,9 @@
     /** Callback provided by WM to call once we're done with the going away animation. */
     private var goingAwayRemoteAnimationFinishedCallback: IRemoteAnimationFinishedCallback? = null
 
+    private val enableNewKeyguardShellTransitions: Boolean =
+        Flags.ensureKeyguardDoesTransitionStarting()
+
     /**
      * Set the visibility of the surface behind the keyguard, making the appropriate calls to Window
      * Manager to effect the change.
@@ -114,7 +120,14 @@
             return
         }
 
+
+
         if (visible) {
+            if (enableNewKeyguardShellTransitions) {
+                keyguardTransitions.startKeyguardTransition(false /* keyguardShowing */, false /* aodShowing */)
+                isKeyguardGoingAway = true
+                return
+            }
             // Make the surface visible behind the keyguard by calling keyguardGoingAway. The
             // lockscreen is still showing as well, allowing us to animate unlocked.
             Log.d(TAG, "ActivityTaskManagerService#keyguardGoingAway()")
@@ -220,7 +233,11 @@
                 "isLockscreenShowing=$lockscreenShowing, " +
                 "aodVisible=$aodVisible)."
         )
-        activityTaskManagerService.setLockScreenShown(lockscreenShowing, aodVisible)
+        if (enableNewKeyguardShellTransitions) {
+            keyguardTransitions.startKeyguardTransition(lockscreenShowing, aodVisible)
+        } else {
+            activityTaskManagerService.setLockScreenShown(lockscreenShowing, aodVisible)
+        }
         this.isLockscreenShowing = lockscreenShowing
         this.isAodVisible = aodVisible
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt
index ec52055..95d1b5d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardClockRepository.kt
@@ -68,11 +68,16 @@
 
     val previewClock: Flow<ClockController>
 
+    /** top of notifications without bcsmartspace in small clock settings */
+    val notificationDefaultTop: StateFlow<Float>
+
     val clockEventController: ClockEventController
 
     val shouldForceSmallClock: Boolean
 
     fun setClockSize(size: ClockSize)
+
+    fun setNotificationDefaultTop(top: Float)
 }
 
 @SysUISingleton
@@ -108,7 +113,7 @@
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue = getClockSize()
+                initialValue = getClockSize(),
             )
 
     override val currentClockId: Flow<ClockId> =
@@ -138,7 +143,7 @@
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.WhileSubscribed(),
-                initialValue = clockRegistry.createCurrentClock()
+                initialValue = clockRegistry.createCurrentClock(),
             )
 
     override val previewClock: Flow<ClockController> =
@@ -149,6 +154,14 @@
             clockRegistry.createCurrentClock()
         }
 
+    private val _notificationDefaultTop: MutableStateFlow<Float> = MutableStateFlow(0F)
+
+    override val notificationDefaultTop: StateFlow<Float> = _notificationDefaultTop.asStateFlow()
+
+    override fun setNotificationDefaultTop(top: Float) {
+        _notificationDefaultTop.value = top
+    }
+
     override val shouldForceSmallClock: Boolean
         get() =
             featureFlags.isEnabled(Flags.LOCKSCREEN_ENABLE_LANDSCAPE) &&
@@ -160,7 +173,7 @@
             secureSettings.getIntForUser(
                 Settings.Secure.LOCKSCREEN_USE_DOUBLE_LINE_CLOCK,
                 /* defaultValue= */ 1,
-                UserHandle.USER_CURRENT
+                UserHandle.USER_CURRENT,
             )
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
index 130242f..8210174 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardRepository.kt
@@ -269,6 +269,9 @@
      */
     val isEncryptedOrLockdown: Flow<Boolean>
 
+    /** The top of shortcut in screen, used by wallpaper to find remaining space in lockscreen */
+    val shortcutAbsoluteTop: StateFlow<Float>
+
     /**
      * Returns `true` if the keyguard is showing; `false` otherwise.
      *
@@ -339,6 +342,8 @@
      * otherwise.
      */
     fun isShowKeyguardWhenReenabled(): Boolean
+
+    fun setShortcutAbsoluteTop(top: Float)
 }
 
 /** Encapsulates application state for the keyguard. */
@@ -503,7 +508,7 @@
                 trySendWithFailureLogging(
                     statusBarStateController.dozeAmount,
                     TAG,
-                    "initial dozeAmount"
+                    "initial dozeAmount",
                 )
 
                 awaitClose { statusBarStateController.removeCallback(callback) }
@@ -521,7 +526,7 @@
             object : DozeTransitionCallback {
                 override fun onDozeTransition(
                     oldState: DozeMachine.State,
-                    newState: DozeMachine.State
+                    newState: DozeMachine.State,
                 ) {
                     trySendWithFailureLogging(
                         DozeTransitionModel(
@@ -529,7 +534,7 @@
                             to = dozeMachineStateToModel(newState),
                         ),
                         TAG,
-                        "doze transition model"
+                        "doze transition model",
                     )
                 }
             }
@@ -541,7 +546,7 @@
                 to = dozeMachineStateToModel(dozeTransitionListener.newState),
             ),
             TAG,
-            "initial doze transition model"
+            "initial doze transition model",
         )
 
         awaitClose { dozeTransitionListener.removeCallback(callback) }
@@ -579,7 +584,7 @@
                             trySendWithFailureLogging(
                                 statusBarStateIntToObject(state),
                                 TAG,
-                                "state"
+                                "state",
                             )
                         }
                     }
@@ -590,7 +595,7 @@
             .stateIn(
                 scope,
                 SharingStarted.Eagerly,
-                statusBarStateIntToObject(statusBarStateController.state)
+                statusBarStateIntToObject(statusBarStateController.state),
             )
 
     private val _biometricUnlockState: MutableStateFlow<BiometricUnlockModel> =
@@ -610,7 +615,7 @@
             trySendWithFailureLogging(
                 authController.fingerprintSensorLocation,
                 TAG,
-                "AuthController.Callback#onFingerprintLocationChanged"
+                "AuthController.Callback#onFingerprintLocationChanged",
             )
         }
 
@@ -635,6 +640,9 @@
     private val _isActiveDreamLockscreenHosted = MutableStateFlow(false)
     override val isActiveDreamLockscreenHosted = _isActiveDreamLockscreenHosted.asStateFlow()
 
+    private val _shortcutAbsoluteTop = MutableStateFlow(0F)
+    override val shortcutAbsoluteTop = _shortcutAbsoluteTop.asStateFlow()
+
     init {
         val callback =
             object : KeyguardStateController.Callback {
@@ -721,6 +729,10 @@
         }
     }
 
+    override fun setShortcutAbsoluteTop(top: Float) {
+        _shortcutAbsoluteTop.value = top
+    }
+
     private fun dozeMachineStateToModel(state: DozeMachine.State): DozeStateModel {
         return when (state) {
             DozeMachine.State.UNINITIALIZED -> DozeStateModel.UNINITIALIZED
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt
index c0049d4..5b7eedd 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractor.kt
@@ -94,7 +94,7 @@
                 .stateIn(
                     scope = applicationScope,
                     started = SharingStarted.WhileSubscribed(),
-                    initialValue = ClockSize.LARGE
+                    initialValue = ClockSize.LARGE,
                 )
         } else {
             keyguardClockRepository.clockSize
@@ -152,4 +152,8 @@
             clock.largeClock.animations.fold(foldFraction)
         }
     }
+
+    fun setNotificationStackDefaultTop(top: Float) {
+        keyguardClockRepository.setNotificationDefaultTop(top)
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
index e6ee112..d7f96b5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -526,6 +526,10 @@
         repository.showDismissibleKeyguard()
     }
 
+    fun setShortcutAbsoluteTop(top: Float) {
+        repository.setShortcutAbsoluteTop(top)
+    }
+
     companion object {
         private const val TAG = "KeyguardInteractor"
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt
index a1c963b..d745522 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt
@@ -44,6 +44,7 @@
 import com.android.systemui.plugins.clocks.ClockController
 import com.android.systemui.plugins.clocks.ClockFaceLayout
 import com.android.systemui.res.R
+import com.android.systemui.shade.LargeScreenHeaderHelper
 import com.android.systemui.shared.R as sharedR
 import com.android.systemui.util.ui.value
 import dagger.Lazy
@@ -53,6 +54,9 @@
 internal fun ConstraintSet.setVisibility(views: Iterable<View>, visibility: Int) =
     views.forEach { view -> this.setVisibility(view.id, visibility) }
 
+internal fun ConstraintSet.setAlpha(views: Iterable<View>, alpha: Float) =
+    views.forEach { view -> this.setAlpha(view.id, alpha) }
+
 internal fun ConstraintSet.setScaleX(views: Iterable<View>, scaleX: Float) =
     views.forEach { view -> this.setScaleX(view.id, scaleX) }
 
@@ -70,6 +74,7 @@
     val blueprintInteractor: Lazy<KeyguardBlueprintInteractor>,
     private val rootViewModel: KeyguardRootViewModel,
     private val aodBurnInViewModel: AodBurnInViewModel,
+    private val largeScreenHeaderHelperLazy: Lazy<LargeScreenHeaderHelper>,
 ) : KeyguardSection() {
     private var disposableHandle: DisposableHandle? = null
 
@@ -123,6 +128,8 @@
         return constraintSet.apply {
             setVisibility(getTargetClockFace(clock).views, VISIBLE)
             setVisibility(getNonTargetClockFace(clock).views, GONE)
+            setAlpha(getTargetClockFace(clock).views, 1F)
+            setAlpha(getNonTargetClockFace(clock).views, 0F)
             if (!keyguardClockViewModel.isLargeClockVisible.value) {
                 connect(sharedR.id.bc_smartspace_view, TOP, sharedR.id.date_smartspace_view, BOTTOM)
             } else {
@@ -172,7 +179,7 @@
         }
     }
 
-    open fun applyDefaultConstraints(constraints: ConstraintSet) {
+    fun applyDefaultConstraints(constraints: ConstraintSet) {
         val guideline =
             if (keyguardClockViewModel.clockShouldBeCentered.value) PARENT_ID
             else R.id.split_shade_guideline
@@ -211,6 +218,28 @@
 
             // Explicitly clear pivot to force recalculate pivot instead of using legacy value
             setTransformPivot(R.id.lockscreen_clock_view_large, Float.NaN, Float.NaN)
+
+            val smallClockBottom =
+                keyguardClockViewModel.getSmallClockTopMargin() +
+                    context.resources.getDimensionPixelSize(
+                        com.android.systemui.customization.R.dimen.small_clock_height
+                    )
+            val dateWeatherSmartspaceHeight = getDimen(context, DATE_WEATHER_VIEW_HEIGHT).toFloat()
+            val marginBetweenSmartspaceAndNotification =
+                context.resources.getDimensionPixelSize(
+                    R.dimen.keyguard_status_view_bottom_margin
+                ) +
+                    if (context.resources.getBoolean(R.bool.config_use_large_screen_shade_header)) {
+                        largeScreenHeaderHelperLazy.get().getLargeScreenHeaderHeight()
+                    } else {
+                        0
+                    }
+
+            clockInteractor.setNotificationStackDefaultTop(
+                smallClockBottom +
+                    dateWeatherSmartspaceHeight +
+                    marginBetweenSmartspaceAndNotification
+            )
         }
 
         constrainWeatherClockDateIconsBarrier(constraints)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt
index 6c6e14c..d3895de 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultShortcutsSection.kt
@@ -30,6 +30,7 @@
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.keyguard.KeyguardBottomAreaRefactor
 import com.android.systemui.keyguard.domain.interactor.KeyguardBlueprintInteractor
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.ui.binder.KeyguardQuickAffordanceViewBinder
 import com.android.systemui.keyguard.ui.view.layout.blueprints.transitions.IntraBlueprintTransition
 import com.android.systemui.keyguard.ui.viewmodel.KeyguardQuickAffordancesCombinedViewModel
@@ -52,6 +53,7 @@
     private val indicationController: KeyguardIndicationController,
     private val keyguardBlueprintInteractor: Lazy<KeyguardBlueprintInteractor>,
     private val keyguardQuickAffordanceViewBinder: KeyguardQuickAffordanceViewBinder,
+    private val keyguardInteractor: KeyguardInteractor,
 ) : BaseShortcutSection() {
 
     // Amount to increase the bottom margin by to avoid colliding with inset
@@ -117,7 +119,7 @@
                 BOTTOM,
                 PARENT_ID,
                 BOTTOM,
-                verticalOffsetMargin + safeInsetBottom
+                verticalOffsetMargin + safeInsetBottom,
             )
 
             constrainWidth(R.id.end_button, width)
@@ -128,7 +130,7 @@
                 BOTTOM,
                 PARENT_ID,
                 BOTTOM,
-                verticalOffsetMargin + safeInsetBottom
+                verticalOffsetMargin + safeInsetBottom,
             )
 
             // The constraint set visibility for start and end button are default visible, set to
@@ -136,5 +138,13 @@
             setVisibilityMode(R.id.start_button, VISIBILITY_MODE_IGNORE)
             setVisibilityMode(R.id.end_button, VISIBILITY_MODE_IGNORE)
         }
+
+        val shortcutAbsoluteTopInScreen =
+            (resources.displayMetrics.heightPixels -
+                    (verticalOffsetMargin + safeInsetBottom) -
+                    height)
+                .toFloat()
+
+        keyguardInteractor.setShortcutAbsoluteTop(shortcutAbsoluteTopInScreen)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
index 3705c2c..40d4193 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModel.kt
@@ -347,12 +347,8 @@
                             when {
                                 // If there are no notification icons to show, then it can be hidden
                                 !hasAodIcons -> false
-                                // If we're bypassing, then we're visible
-                                isBypassEnabled -> true
                                 // If we are pulsing (and not bypassing), then we are hidden
                                 isPulseExpanding -> false
-                                // Besides bypass above, they should not be visible on lockscreen
-                                isOnLockscreen -> false
                                 // If notifs are fully gone, then we're visible
                                 areNotifsFullyHidden -> true
                                 // Otherwise, we're hidden
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSContainerImpl.java b/packages/SystemUI/src/com/android/systemui/qs/QSContainerImpl.java
index 1511f31..fa987dd 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSContainerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSContainerImpl.java
@@ -255,7 +255,8 @@
      * @return size in pixels of QQS
      */
     public int getQqsHeight() {
-        return mHeader.getHeight();
+        SceneContainerFlag.assertInNewMode();
+        return mHeader.getMeasuredHeight();
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
index 8a2e274..cd38ca6 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
@@ -6,6 +6,7 @@
 import android.view.View
 import android.view.WindowInsets
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerDependencies
+import com.android.systemui.qs.ui.adapter.QSSceneAdapter
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
 import com.android.systemui.scene.ui.composable.Overlay
@@ -13,19 +14,13 @@
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
 import com.android.systemui.shade.TouchLogger
 import com.android.systemui.statusbar.notification.stack.ui.view.SharedNotificationContainer
+import javax.inject.Provider
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
 
 /** A root view of the main SysUI window that supports scenes. */
 @ExperimentalCoroutinesApi
-class SceneWindowRootView(
-    context: Context,
-    attrs: AttributeSet?,
-) :
-    WindowRootView(
-        context,
-        attrs,
-    ) {
+class SceneWindowRootView(context: Context, attrs: AttributeSet?) : WindowRootView(context, attrs) {
 
     private var motionEventHandler: SceneContainerViewModel.MotionEventHandler? = null
     // TODO(b/298525212): remove once Compose exposes window inset bounds.
@@ -39,6 +34,7 @@
         overlays: Set<Overlay>,
         layoutInsetController: LayoutInsetsController,
         sceneDataSourceDelegator: SceneDataSourceDelegator,
+        qsSceneAdapter: Provider<QSSceneAdapter>,
         alternateBouncerDependencies: AlternateBouncerDependencies,
     ) {
         setLayoutInsetsController(layoutInsetController)
@@ -57,6 +53,7 @@
                 super.setVisibility(if (isVisible) View.VISIBLE else View.INVISIBLE)
             },
             dataSourceDelegator = sceneDataSourceDelegator,
+            qsSceneAdapter = qsSceneAdapter,
             alternateBouncerDependencies = alternateBouncerDependencies,
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
index 38f4e73..1e3a233 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
@@ -42,6 +42,7 @@
 import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.lifecycle.setSnapshotBinding
 import com.android.systemui.lifecycle.viewModel
+import com.android.systemui.qs.ui.adapter.QSSceneAdapter
 import com.android.systemui.res.R
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.SceneContainerConfig
@@ -51,6 +52,7 @@
 import com.android.systemui.scene.ui.composable.SceneContainer
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
 import com.android.systemui.statusbar.notification.stack.ui.view.SharedNotificationContainer
+import javax.inject.Provider
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.awaitCancellation
@@ -74,6 +76,7 @@
         overlays: Set<Overlay>,
         onVisibilityChangedInternal: (isVisible: Boolean) -> Unit,
         dataSourceDelegator: SceneDataSourceDelegator,
+        qsSceneAdapter: Provider<QSSceneAdapter>,
         alternateBouncerDependencies: AlternateBouncerDependencies,
     ) {
         val unsortedSceneByKey: Map<SceneKey, Scene> = scenes.associateBy { scene -> scene.key }
@@ -132,6 +135,7 @@
                                 sceneByKey = sortedSceneByKey,
                                 overlayByKey = sortedOverlayByKey,
                                 dataSourceDelegator = dataSourceDelegator,
+                                qsSceneAdapter = qsSceneAdapter,
                                 containerConfig = containerConfig,
                             )
                             .also { it.id = R.id.scene_container_root_composable }
@@ -177,6 +181,7 @@
         sceneByKey: Map<SceneKey, Scene>,
         overlayByKey: Map<OverlayKey, Overlay>,
         dataSourceDelegator: SceneDataSourceDelegator,
+        qsSceneAdapter: Provider<QSSceneAdapter>,
         containerConfig: SceneContainerConfig,
     ): View {
         return ComposeView(context).apply {
@@ -192,6 +197,7 @@
                             overlayByKey = overlayByKey,
                             initialSceneKey = containerConfig.initialSceneKey,
                             dataSourceDelegator = dataSourceDelegator,
+                            qsSceneAdapter = qsSceneAdapter,
                         )
                     }
                 }
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 889380a..ce942fe 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
@@ -73,6 +73,8 @@
     /** Whether the container is visible. */
     val isVisible: Boolean by hydrator.hydratedStateOf("isVisible", sceneInteractor.isVisible)
 
+    val allContentKeys: List<ContentKey> = sceneInteractor.allContentKeys
+
     private val hapticsViewModel = hapticsViewModelFactory.create(view)
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt b/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt
index 7b56688..6d3b9aa 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/HeadlessScreenshotHandler.kt
@@ -19,7 +19,6 @@
 import android.net.Uri
 import android.os.UserManager
 import android.util.Log
-import android.view.WindowManager
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.res.R
@@ -51,10 +50,6 @@
         finisher: Consumer<Uri?>,
         requestCallback: TakeScreenshotService.RequestCallback,
     ) {
-        if (screenshot.type == WindowManager.TAKE_SCREENSHOT_FULLSCREEN) {
-            screenshot.bitmap = imageCapture.captureDisplay(screenshot.displayId, crop = null)
-        }
-
         if (screenshot.bitmap == null) {
             Log.e(TAG, "handleScreenshot: Screenshot bitmap was null")
             notificationsControllerFactory
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
index e589600..acfcd13 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
@@ -255,12 +255,6 @@
         Assert.isMainThread();
 
         mCurrentRequestCallback = requestCallback;
-        if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_FULLSCREEN
-                && screenshot.getBitmap() == null) {
-            Rect bounds = getFullScreenRect();
-            screenshot.setBitmap(mImageCapture.captureDisplay(mDisplay.getDisplayId(), bounds));
-            screenshot.setScreenBounds(bounds);
-        }
 
         if (screenshot.getBitmap() == null) {
             Log.e(TAG, "handleScreenshot: Screenshot bitmap was null");
@@ -323,24 +317,25 @@
 
         attachWindow();
 
+        Rect bounds = screenshot.getOriginalScreenBounds();
         boolean showFlash;
         if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-            if (screenshot.getScreenBounds() != null
-                    && aspectRatiosMatch(screenshot.getBitmap(), screenshot.getInsets(),
-                    screenshot.getScreenBounds())) {
+            if (bounds != null
+                    && aspectRatiosMatch(screenshot.getBitmap(), screenshot.getOriginalInsets(),
+                    bounds)) {
                 showFlash = false;
             } else {
                 showFlash = true;
-                screenshot.setInsets(Insets.NONE);
-                screenshot.setScreenBounds(new Rect(0, 0, screenshot.getBitmap().getWidth(),
-                        screenshot.getBitmap().getHeight()));
+                bounds = new Rect(0, 0, screenshot.getBitmap().getWidth(),
+                        screenshot.getBitmap().getHeight());
             }
         } else {
             showFlash = true;
         }
 
+        final Rect animationBounds = bounds;
         mViewProxy.prepareEntranceAnimation(
-                () -> startAnimation(screenshot.getScreenBounds(), showFlash,
+                () -> startAnimation(animationBounds, showFlash,
                         () -> mMessageContainerController.onScreenshotTaken(screenshot)));
 
         mViewProxy.setScreenshot(screenshot);
@@ -646,7 +641,7 @@
     private void saveScreenshotInBackground(ScreenshotData screenshot, UUID requestId,
             Consumer<Uri> finisher, Consumer<ImageExporter.Result> onResult) {
         ListenableFuture<ImageExporter.Result> future = mImageExporter.export(mBgExecutor,
-                requestId, screenshot.getBitmap(), screenshot.getUserOrDefault(),
+                requestId, screenshot.getBitmap(), screenshot.getUserHandle(),
                 mDisplay.getDisplayId());
         future.addListener(() -> {
             try {
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
index a762d84..f5c6052 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
@@ -35,7 +35,6 @@
 import android.view.Display
 import android.view.ScrollCaptureResponse
 import android.view.ViewRootImpl.ActivityConfigCallback
-import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
 import android.widget.Toast
 import android.window.WindowContext
@@ -60,7 +59,6 @@
 import java.util.concurrent.ExecutorService
 import java.util.concurrent.Executors
 import java.util.function.Consumer
-import javax.inject.Provider
 import kotlin.math.abs
 
 /** Controls the state and flow for screenshots. */
@@ -73,7 +71,7 @@
     screenshotNotificationsControllerFactory: ScreenshotNotificationsController.Factory,
     screenshotActionsControllerFactory: ScreenshotActionsController.Factory,
     actionExecutorFactory: ActionExecutor.Factory,
-    screenshotSoundControllerProvider: Provider<ScreenshotSoundController?>,
+    private val screenshotSoundController: ScreenshotSoundController,
     private val uiEventLogger: UiEventLogger,
     private val imageExporter: ImageExporter,
     private val imageCapture: ImageCapture,
@@ -100,7 +98,6 @@
     private val currentRequestCallbacks: MutableList<TakeScreenshotService.RequestCallback> =
         mutableListOf()
 
-    private var screenshotSoundController: ScreenshotSoundController? = null
     private var screenBitmap: Bitmap? = null
     private var screenshotTakenInPortrait = false
     private var screenshotAnimation: Animator? = null
@@ -138,14 +135,6 @@
         actionExecutor = actionExecutorFactory.create(window.window, viewProxy) { finishDismiss() }
         actionsController = screenshotActionsControllerFactory.getController(actionExecutor)
 
-        // Sound is only reproduced from the controller of the default display.
-        screenshotSoundController =
-            if (display.displayId == Display.DEFAULT_DISPLAY) {
-                screenshotSoundControllerProvider.get()
-            } else {
-                null
-            }
-
         copyBroadcastReceiver =
             object : BroadcastReceiver() {
                 override fun onReceive(context: Context, intent: Intent) {
@@ -172,12 +161,6 @@
         Assert.isMainThread()
         screenshotHandler.resetTimeout()
 
-        if (screenshot.type == TAKE_SCREENSHOT_FULLSCREEN && screenshot.bitmap == null) {
-            val bounds = fullScreenRect
-            screenshot.bitmap = imageCapture.captureDisplay(display.displayId, bounds)
-            screenshot.screenBounds = bounds
-        }
-
         val currentBitmap = screenshot.bitmap
         if (currentBitmap == null) {
             Log.e(TAG, "handleScreenshot: Screenshot bitmap was null")
@@ -238,23 +221,27 @@
 
         window.attachWindow()
 
+        var bounds =
+            screenshot.originalScreenBounds ?: Rect(0, 0, currentBitmap.width, currentBitmap.height)
+
         val showFlash: Boolean
         if (screenshot.type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-            if (aspectRatiosMatch(currentBitmap, screenshot.insets, screenshot.screenBounds)) {
+            if (
+                aspectRatiosMatch(
+                    currentBitmap,
+                    screenshot.originalInsets,
+                    screenshot.originalScreenBounds,
+                )
+            ) {
                 showFlash = false
             } else {
                 showFlash = true
-                screenshot.insets = Insets.NONE
-                screenshot.screenBounds = Rect(0, 0, currentBitmap.width, currentBitmap.height)
+                bounds = Rect(0, 0, currentBitmap.width, currentBitmap.height)
             }
         } else {
             showFlash = true
         }
 
-        // screenshot.screenBounds is expected to be non-null in all cases at this point
-        val bounds =
-            screenshot.screenBounds ?: Rect(0, 0, currentBitmap.width, currentBitmap.height)
-
         viewProxy.prepareEntranceAnimation {
             startAnimation(bounds, showFlash) {
                 messageContainerController.onScreenshotTaken(screenshot)
@@ -305,7 +292,7 @@
     // Any cleanup needed when the service is being destroyed.
     override fun onDestroy() {
         removeWindow()
-        releaseMediaPlayer()
+        screenshotSoundController.releaseScreenshotSoundAsync()
         releaseContext()
         bgExecutor.shutdown()
     }
@@ -316,10 +303,6 @@
         context.release()
     }
 
-    private fun releaseMediaPlayer() {
-        screenshotSoundController?.releaseScreenshotSoundAsync()
-    }
-
     /** Update resources on configuration change. Reinflate for theme/color changes. */
     private fun reloadAssets() {
         if (LogConfig.DEBUG_UI) {
@@ -445,18 +428,13 @@
         viewProxy.stopInputListening()
     }
 
-    private fun playCameraSoundIfNeeded() {
-        // the controller is not-null only on the default display controller
-        screenshotSoundController?.playScreenshotSoundAsync()
-    }
-
     /**
      * Save the bitmap but don't show the normal screenshot UI.. just a toast (or notification on
      * failure).
      */
     private fun saveScreenshotAndToast(screenshot: ScreenshotData, finisher: Consumer<Uri?>) {
         // Play the shutter sound to notify that we've taken a screenshot
-        playCameraSoundIfNeeded()
+        screenshotSoundController.playScreenshotSoundAsync()
 
         saveScreenshotInBackground(screenshot, UUID.randomUUID(), finisher) {
             result: ImageExporter.Result ->
@@ -485,7 +463,7 @@
             viewProxy.createScreenshotDropInAnimation(screenRect, showFlash).apply {
                 doOnEnd { onAnimationComplete?.run() }
                 // Play the shutter sound to notify that we've taken a screenshot
-                playCameraSoundIfNeeded()
+                screenshotSoundController.playScreenshotSoundAsync()
                 if (LogConfig.DEBUG_ANIM) {
                     Log.d(TAG, "starting post-screenshot animation")
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
index 2df1e8a..b5b15a9 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotData.kt
@@ -4,7 +4,6 @@
 import android.graphics.Bitmap
 import android.graphics.Insets
 import android.graphics.Rect
-import android.os.Process
 import android.os.UserHandle
 import android.view.Display
 import android.view.WindowManager
@@ -21,19 +20,15 @@
     val userHandle: UserHandle,
     /** ComponentName of the top-most app in the screenshot. */
     val topComponent: ComponentName?,
-    var screenBounds: Rect?,
     val taskId: Int,
-    var insets: Insets,
+    val originalScreenBounds: Rect?,
+    val originalInsets: Insets,
     var bitmap: Bitmap?,
     val displayId: Int,
 ) {
     val packageNameString
         get() = topComponent?.packageName ?: ""
 
-    fun getUserOrDefault(): UserHandle {
-        return userHandle ?: Process.myUserHandle()
-    }
-
     companion object {
         @JvmStatic
         fun fromRequest(request: ScreenshotRequest, displayId: Int = Display.DEFAULT_DISPLAY) =
@@ -42,9 +37,9 @@
                 source = request.source,
                 userHandle = UserHandle.of(request.userId),
                 topComponent = request.topComponent,
-                screenBounds = request.boundsInScreen,
+                originalScreenBounds = request.boundsInScreen,
                 taskId = request.taskId,
-                insets = request.insets,
+                originalInsets = request.insets,
                 bitmap = request.bitmap,
                 displayId = displayId,
             )
@@ -61,9 +56,9 @@
                 source = source,
                 userHandle = userHandle,
                 topComponent = topComponent,
-                screenBounds = null,
+                originalScreenBounds = null,
                 taskId = 0,
-                insets = Insets.NONE,
+                originalInsets = Insets.NONE,
                 bitmap = bitmap,
                 displayId = Display.DEFAULT_DISPLAY,
             )
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
index b67ad8a..039143a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
@@ -126,7 +126,7 @@
         )
     }
 
-    suspend fun replaceWithTaskSnapshot(
+    private suspend fun replaceWithTaskSnapshot(
         original: ScreenshotData,
         componentName: ComponentName?,
         owner: UserHandle,
@@ -134,14 +134,14 @@
         taskBounds: Rect?,
     ): ScreenshotData {
         Log.i(TAG, "Capturing task snapshot: $componentName / $owner")
-        val taskSnapshot = capture.captureTask(taskId)
+        val taskSnapshot = capture.captureTask(taskId) ?: error("Failed to capture task")
         return original.copy(
             type = TAKE_SCREENSHOT_PROVIDED_IMAGE,
             bitmap = taskSnapshot,
             userHandle = owner,
             taskId = taskId,
             topComponent = componentName,
-            screenBounds = taskBounds,
+            originalScreenBounds = taskBounds,
         )
     }
 
@@ -153,13 +153,13 @@
         taskId: Int? = null,
     ): ScreenshotData {
         Log.i(TAG, "Capturing screenshot: $componentName / $owner")
-        val screenshot = captureDisplay(displayId)
+        val screenshot = captureDisplay(displayId) ?: error("Failed to capture screenshot")
         return original.copy(
             type = TAKE_SCREENSHOT_FULLSCREEN,
             bitmap = screenshot,
             userHandle = owner,
             topComponent = componentName,
-            screenBounds = Rect(0, 0, screenshot?.width ?: 0, screenshot?.height ?: 0),
+            originalScreenBounds = Rect(0, 0, screenshot.width, screenshot.height),
             taskId = taskId ?: -1,
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
index 2bff7c86..757e37e 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
@@ -2728,6 +2728,11 @@
     }
 
     private int calculatePanelHeightShade() {
+        // Bypass should always occupy the full height
+        if (mBarState == KEYGUARD && mKeyguardBypassController.getBypassEnabled()) {
+            return mNotificationStackScrollLayoutController.getHeight();
+        }
+
         int emptyBottomMargin = mNotificationStackScrollLayoutController.getEmptyBottomMargin();
         int maxHeight = mNotificationStackScrollLayoutController.getHeight() - emptyBottomMargin;
 
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
index fc8a593..5afc539 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
@@ -35,6 +35,7 @@
 import com.android.systemui.keyguard.ui.view.KeyguardRootView
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerDependencies
 import com.android.systemui.privacy.OngoingPrivacyChip
+import com.android.systemui.qs.ui.adapter.QSSceneAdapter
 import com.android.systemui.res.R
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.SceneContainerConfig
@@ -89,6 +90,7 @@
             overlaysProvider: Provider<Set<@JvmSuppressWildcards Overlay>>,
             layoutInsetController: NotificationInsetsController,
             sceneDataSourceDelegator: Provider<SceneDataSourceDelegator>,
+            qsSceneAdapter: Provider<QSSceneAdapter>,
             alternateBouncerDependencies: Provider<AlternateBouncerDependencies>,
         ): WindowRootView {
             return if (SceneContainerFlag.isEnabled) {
@@ -104,6 +106,7 @@
                     overlays = overlaysProvider.get(),
                     layoutInsetController = layoutInsetController,
                     sceneDataSourceDelegator = sceneDataSourceDelegator.get(),
+                    qsSceneAdapter = qsSceneAdapter,
                     alternateBouncerDependencies = alternateBouncerDependencies.get(),
                 )
                 sceneWindowRootView
@@ -119,9 +122,7 @@
         //  {@link NotificationShadeWindowViewController} can inject this view.
         @Provides
         @SysUISingleton
-        fun providesNotificationShadeWindowView(
-            root: WindowRootView,
-        ): NotificationShadeWindowView {
+        fun providesNotificationShadeWindowView(root: WindowRootView): NotificationShadeWindowView {
             if (SceneContainerFlag.isEnabled) {
                 return root.requireViewById(R.id.legacy_window_root)
             }
@@ -133,7 +134,7 @@
         @Provides
         @SysUISingleton
         fun providesNotificationStackScrollLayout(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            notificationShadeWindowView: NotificationShadeWindowView
         ): NotificationStackScrollLayout {
             return notificationShadeWindowView.requireViewById(R.id.notification_stack_scroller)
         }
@@ -142,7 +143,7 @@
         @Provides
         @SysUISingleton
         fun providesNotificationPanelView(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            notificationShadeWindowView: NotificationShadeWindowView
         ): NotificationPanelView {
             return notificationShadeWindowView.requireViewById(R.id.notification_panel)
         }
@@ -178,7 +179,7 @@
         @Provides
         @SysUISingleton
         fun providesKeyguardRootView(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            notificationShadeWindowView: NotificationShadeWindowView
         ): KeyguardRootView {
             return notificationShadeWindowView.requireViewById(R.id.keyguard_root_view)
         }
@@ -186,7 +187,7 @@
         @Provides
         @SysUISingleton
         fun providesSharedNotificationContainer(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            notificationShadeWindowView: NotificationShadeWindowView
         ): SharedNotificationContainer {
             return notificationShadeWindowView.requireViewById(R.id.shared_notification_container)
         }
@@ -195,7 +196,7 @@
         @Provides
         @SysUISingleton
         fun providesAuthRippleView(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            notificationShadeWindowView: NotificationShadeWindowView
         ): AuthRippleView? {
             return notificationShadeWindowView.requireViewById(R.id.auth_ripple)
         }
@@ -203,9 +204,7 @@
         // TODO(b/277762009): Only allow this view's controller to inject the view. See above.
         @Provides
         @SysUISingleton
-        fun providesTapAgainView(
-            notificationPanelView: NotificationPanelView,
-        ): TapAgainView {
+        fun providesTapAgainView(notificationPanelView: NotificationPanelView): TapAgainView {
             return notificationPanelView.requireViewById(R.id.shade_falsing_tap_again)
         }
 
@@ -213,7 +212,7 @@
         @Provides
         @SysUISingleton
         fun providesNotificationsQuickSettingsContainer(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            notificationShadeWindowView: NotificationShadeWindowView
         ): NotificationsQuickSettingsContainer {
             return notificationShadeWindowView.requireViewById(R.id.notification_container_parent)
         }
@@ -223,7 +222,7 @@
         @SysUISingleton
         @Named(SHADE_HEADER)
         fun providesShadeHeaderView(
-            notificationShadeWindowView: NotificationShadeWindowView,
+            notificationShadeWindowView: NotificationShadeWindowView
         ): MotionLayout {
             val stub = notificationShadeWindowView.requireViewById<ViewStub>(R.id.qs_header_stub)
             val layoutId = R.layout.combined_qs_header
@@ -275,7 +274,7 @@
         @SysUISingleton
         @Named(SHADE_HEADER)
         fun providesOngoingPrivacyChip(
-            @Named(SHADE_HEADER) header: MotionLayout,
+            @Named(SHADE_HEADER) header: MotionLayout
         ): OngoingPrivacyChip {
             return header.requireViewById(R.id.privacy_chip)
         }
@@ -284,7 +283,7 @@
         @SysUISingleton
         @Named(SHADE_HEADER)
         fun providesStatusIconContainer(
-            @Named(SHADE_HEADER) header: MotionLayout,
+            @Named(SHADE_HEADER) header: MotionLayout
         ): StatusIconContainer {
             return header.requireViewById(R.id.statusIcons)
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index 2c6b09c..adb3352 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -3884,6 +3884,7 @@
             }
             showingLayout.dump(pw, args);
             dumpCustomOutline(pw, args);
+            dumpClipping(pw, args);
             if (getViewState() != null) {
                 getViewState().dump(pw, args);
                 pw.println();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java
index afda426..ef6cad1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableView.java
@@ -856,17 +856,23 @@
                 pw.println();
             }
             if (DUMP_VERBOSE) {
-                pw.println("mInRemovalAnimation: " + mInRemovalAnimation);
-                pw.println("mClipTopAmount: " + mClipTopAmount);
-                pw.println("mClipBottomAmount " + mClipBottomAmount);
-                pw.println("mClipToActualHeight: " + mClipToActualHeight);
-                pw.println("mExtraWidthForClipping: " + mExtraWidthForClipping);
-                pw.println("mMinimumHeightForClipping: " + mMinimumHeightForClipping);
-                pw.println("getClipBounds(): " + getClipBounds());
+                dumpClipping(pw, args);
             }
         });
     }
 
+    protected void dumpClipping(IndentingPrintWriter pw, String[] args) {
+        pw.print("Clipping: ");
+        pw.print("mInRemovalAnimation", mInRemovalAnimation);
+        pw.print("mClipTopAmount", mClipTopAmount);
+        pw.print("mClipBottomAmount", mClipBottomAmount);
+        pw.print("mClipToActualHeight", mClipToActualHeight);
+        pw.print("mExtraWidthForClipping", mExtraWidthForClipping);
+        pw.print("mMinimumHeightForClipping", mMinimumHeightForClipping);
+        pw.print("getClipBounds()", getClipBounds());
+        pw.println();
+    }
+
     /**
      * return the amount that the content is translated
      */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
index 36e3e92..69f45db 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
@@ -360,11 +360,11 @@
         if ((contentViews & FLAG_CONTENT_VIEW_PUBLIC) != 0) {
             row.getPublicLayout().removeContentInactiveRunnable(VISIBLE_TYPE_CONTRACTED);
         }
-        if (AsyncHybridViewInflation.isEnabled()
+        if (LockscreenOtpRedaction.isEnabled()
                 && (contentViews & FLAG_CONTENT_VIEW_PUBLIC_SINGLE_LINE) != 0) {
             row.getPublicLayout().removeContentInactiveRunnable(VISIBLE_TYPE_SINGLELINE);
         }
-        if (LockscreenOtpRedaction.isEnabled()
+        if (AsyncHybridViewInflation.isEnabled()
                 && (contentViews & FLAG_CONTENT_VIEW_SINGLE_LINE) != 0) {
             row.getPrivateLayout().removeContentInactiveRunnable(VISIBLE_TYPE_SINGLELINE);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index 8ae6b79..d828a67 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -130,6 +130,7 @@
 import com.android.systemui.util.ColorUtilKt;
 import com.android.systemui.util.DumpUtilsKt;
 import com.android.systemui.util.ListenerSet;
+import com.android.systemui.wallpapers.domain.interactor.WallpaperInteractor;
 
 import com.google.errorprone.annotations.CompileTimeConstant;
 
@@ -628,6 +629,9 @@
     @Nullable
     private OnClickListener mManageButtonClickListener;
 
+    @Nullable
+    private WallpaperInteractor mWallpaperInteractor;
+
     public NotificationStackScrollLayout(Context context, AttributeSet attrs) {
         super(context, attrs, 0, 0);
         Resources res = getResources();
@@ -1192,6 +1196,7 @@
         if (!SceneContainerFlag.isEnabled()) {
             setMaxLayoutHeight(getHeight());
             updateContentHeight();
+            mWallpaperInteractor.setNotificationStackAbsoluteBottom(mContentHeight);
         }
         clampScrollPosition();
         requestChildrenUpdate();
@@ -1253,6 +1258,7 @@
         if (mAmbientState.getStackTop() != stackTop) {
             mAmbientState.setStackTop(stackTop);
             onTopPaddingChanged(/* animate = */ isAddOrRemoveAnimationPending());
+            mWallpaperInteractor.setNotificationStackAbsoluteBottom((int) stackTop);
         }
     }
 
@@ -5898,6 +5904,10 @@
         mController.getNotificationRoundnessManager().setAnimatedChildren(mChildrenToAddAnimated);
     }
 
+    public void setWallpaperInteractor(WallpaperInteractor wallpaperInteractor) {
+        mWallpaperInteractor = wallpaperInteractor;
+    }
+
     void addSwipedOutView(View v) {
         mSwipedOutViews.add(v);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index 7b02d0c..00c5c40 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -145,6 +145,7 @@
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.util.Compile;
 import com.android.systemui.util.settings.SecureSettings;
+import com.android.systemui.wallpapers.domain.interactor.WallpaperInteractor;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -226,6 +227,8 @@
     private final SensitiveNotificationProtectionController
             mSensitiveNotificationProtectionController;
 
+    private final WallpaperInteractor mWallpaperInteractor;
+
     private View mLongPressedView;
 
     private final NotificationListContainerImpl mNotificationListContainer =
@@ -756,7 +759,8 @@
             NotificationDismissibilityProvider dismissibilityProvider,
             ActivityStarter activityStarter,
             SplitShadeStateController splitShadeStateController,
-            SensitiveNotificationProtectionController sensitiveNotificationProtectionController) {
+            SensitiveNotificationProtectionController sensitiveNotificationProtectionController,
+            WallpaperInteractor wallpaperInteractor) {
         mView = view;
         mKeyguardTransitionRepo = keyguardTransitionRepo;
         mViewBinder = viewBinder;
@@ -812,6 +816,7 @@
         mDismissibilityProvider = dismissibilityProvider;
         mActivityStarter = activityStarter;
         mSensitiveNotificationProtectionController = sensitiveNotificationProtectionController;
+        mWallpaperInteractor = wallpaperInteractor;
         mView.passSplitShadeStateController(splitShadeStateController);
         if (SceneContainerFlag.isEnabled()) {
             mWakeUpCoordinator.setStackScroller(this);
@@ -948,6 +953,8 @@
             collectFlow(mView, mKeyguardTransitionRepo.getTransitions(),
                     this::onKeyguardTransitionChanged);
         }
+
+        mView.setWallpaperInteractor(mWallpaperInteractor);
     }
 
     private boolean isInVisibleLocation(NotificationEntry entry) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardStatusBarView.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardStatusBarView.java
index 178c318..6a77988 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardStatusBarView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/KeyguardStatusBarView.java
@@ -182,14 +182,8 @@
         mCarrierLabel.setTextSize(TypedValue.COMPLEX_UNIT_PX,
                 getResources().getDimensionPixelSize(
                         com.android.internal.R.dimen.text_size_small_material));
-        lp = (MarginLayoutParams) mCarrierLabel.getLayoutParams();
+        updateCarrierLabelMargin();
 
-        int marginStart = calculateMargin(
-                getResources().getDimensionPixelSize(R.dimen.keyguard_carrier_text_margin),
-                mPadding.left);
-        lp.setMarginStart(marginStart);
-
-        mCarrierLabel.setLayoutParams(lp);
         updateKeyguardStatusBarHeight();
     }
 
@@ -203,6 +197,15 @@
         setLayoutParams(lp);
     }
 
+    private void updateCarrierLabelMargin() {
+        MarginLayoutParams lp = (MarginLayoutParams) mCarrierLabel.getLayoutParams();
+        int marginStart = calculateMargin(
+                getResources().getDimensionPixelSize(R.dimen.keyguard_carrier_text_margin),
+                mPadding.left);
+        lp.setMarginStart(marginStart);
+        mCarrierLabel.setLayoutParams(lp);
+    }
+
     void loadDimens() {
         Resources res = getResources();
         mSystemIconsSwitcherHiddenExpandedMargin = res.getDimensionPixelSize(
@@ -334,6 +337,7 @@
 
         RelativeLayout.LayoutParams lp = (LayoutParams) mCarrierLabel.getLayoutParams();
         lp.addRule(RelativeLayout.START_OF, R.id.status_icon_area);
+        updateCarrierLabelMargin();
 
         lp = (LayoutParams) mStatusIconArea.getLayoutParams();
         lp.removeRule(RelativeLayout.RIGHT_OF);
@@ -366,6 +370,7 @@
 
         lp = (LayoutParams) mCarrierLabel.getLayoutParams();
         lp.addRule(RelativeLayout.START_OF, R.id.cutout_space_view);
+        updateCarrierLabelMargin();
 
         lp = (LayoutParams) mStatusIconArea.getLayoutParams();
         lp.addRule(RelativeLayout.RIGHT_OF, R.id.cutout_space_view);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java
index e09e577..0cba940 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ZenModeControllerImpl.java
@@ -16,8 +16,6 @@
 
 package com.android.systemui.statusbar.policy;
 
-import static com.android.systemui.Flags.registerZenModeContentObserverBackground;
-
 import android.app.AlarmManager;
 import android.app.Flags;
 import android.app.NotificationManager;
@@ -47,7 +45,6 @@
 import com.android.systemui.Dumpable;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.SysUISingleton;
-import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.settings.UserTracker;
@@ -107,7 +104,6 @@
     public ZenModeControllerImpl(
             Context context,
             @Main Handler handler,
-            @Background Handler bgHandler,
             BroadcastDispatcher broadcastDispatcher,
             DumpManager dumpManager,
             GlobalSettings globalSettings,
@@ -138,17 +134,9 @@
             }
         };
         mNoMan = (NotificationManager) context.getSystemService(Context.NOTIFICATION_SERVICE);
-        if (registerZenModeContentObserverBackground()) {
-            bgHandler.post(() -> {
-                globalSettings.registerContentObserverSync(Global.ZEN_MODE, modeContentObserver);
-                globalSettings.registerContentObserverSync(Global.ZEN_MODE_CONFIG_ETAG,
-                        configContentObserver);
-            });
-        } else {
-            globalSettings.registerContentObserverSync(Global.ZEN_MODE, modeContentObserver);
-            globalSettings.registerContentObserverSync(Global.ZEN_MODE_CONFIG_ETAG,
-                    configContentObserver);
-        }
+        globalSettings.registerContentObserverAsync(Global.ZEN_MODE, modeContentObserver);
+        globalSettings.registerContentObserverAsync(Global.ZEN_MODE_CONFIG_ETAG,
+                configContentObserver);
         updateZenMode(getModeSettingValueFromProvider());
         updateZenModeConfig();
         updateConsolidatedNotificationPolicy();
diff --git a/packages/SystemUI/src/com/android/systemui/user/legacyhelper/ui/LegacyUserUiHelper.kt b/packages/SystemUI/src/com/android/systemui/user/legacyhelper/ui/LegacyUserUiHelper.kt
index 8957fe1..09cef1e 100644
--- a/packages/SystemUI/src/com/android/systemui/user/legacyhelper/ui/LegacyUserUiHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/user/legacyhelper/ui/LegacyUserUiHelper.kt
@@ -18,6 +18,7 @@
 package com.android.systemui.user.legacyhelper.ui
 
 import android.content.Context
+import android.util.Log
 import androidx.annotation.DrawableRes
 import androidx.annotation.StringRes
 import com.android.systemui.res.R
@@ -32,6 +33,8 @@
  */
 object LegacyUserUiHelper {
 
+    private const val TAG = "LegacyUserUiHelper"
+
     @JvmStatic
     @DrawableRes
     fun getUserSwitcherActionIconResourceId(
@@ -67,7 +70,9 @@
         val resourceId: Int? = getGuestUserRecordNameResourceId(record)
         return when {
             resourceId != null -> context.getString(resourceId)
-            record.info != null -> checkNotNull(record.info.name)
+            record.info != null ->
+                record.info.name
+                    ?: "".also { Log.i(TAG, "Expected display name for: ${record.info}") }
             else ->
                 context.getString(
                     getUserSwitcherActionTextResourceId(
diff --git a/packages/SystemUI/src/com/android/systemui/util/AlphaTintDrawableWrapper.java b/packages/SystemUI/src/com/android/systemui/util/AlphaTintDrawableWrapper.java
index 947746c..4b2fe49 100644
--- a/packages/SystemUI/src/com/android/systemui/util/AlphaTintDrawableWrapper.java
+++ b/packages/SystemUI/src/com/android/systemui/util/AlphaTintDrawableWrapper.java
@@ -45,8 +45,8 @@
  * This class should only be used in XML.
  *
  * @attr ref android.R.styleable#DrawableWrapper_drawable
- * @attr ref R.styleable#AlphaTintDrawableWrapper_tint
- * @attr ref R.styleable#AlphaTintDrawableWrapper_alpha
+ * @attr ref android.R.styleable#AlphaTintDrawableWrapper_tint
+ * @attr ref android.R.styleable#AlphaTintDrawableWrapper_alpha
  */
 public class AlphaTintDrawableWrapper extends InsetDrawable {
     private ColorStateList mTint;
diff --git a/packages/SystemUI/src/com/android/systemui/util/WallpaperController.kt b/packages/SystemUI/src/com/android/systemui/util/WallpaperController.kt
index 65a0218..0744b5a 100644
--- a/packages/SystemUI/src/com/android/systemui/util/WallpaperController.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/WallpaperController.kt
@@ -32,19 +32,24 @@
  * Note: New logic should be added to [WallpaperRepository], not this class.
  */
 @SysUISingleton
-class WallpaperController @Inject constructor(
+class WallpaperController
+@Inject
+constructor(
     private val wallpaperManager: WallpaperManager,
     private val wallpaperRepository: WallpaperRepository,
 ) {
 
     var rootView: View? = null
+        set(value) {
+            field = value
+            wallpaperRepository.rootView = value
+        }
 
     private var notificationShadeZoomOut: Float = 0f
     private var unfoldTransitionZoomOut: Float = 0f
 
     private val shouldUseDefaultUnfoldTransition: Boolean
-        get() = wallpaperRepository.wallpaperInfo.value?.shouldUseDefaultUnfoldTransition()
-            ?: true
+        get() = wallpaperRepository.wallpaperInfo.value?.shouldUseDefaultUnfoldTransition() ?: true
 
     fun setNotificationShadeZoom(zoomOut: Float) {
         notificationShadeZoomOut = zoomOut
diff --git a/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/NoopWallpaperRepository.kt b/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/NoopWallpaperRepository.kt
index b45b8cd..54953c9 100644
--- a/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/NoopWallpaperRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/NoopWallpaperRepository.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.wallpapers.data.repository
 
 import android.app.WallpaperInfo
+import android.view.View
 import com.android.systemui.dagger.SysUISingleton
 import javax.inject.Inject
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -33,4 +34,7 @@
 class NoopWallpaperRepository @Inject constructor() : WallpaperRepository {
     override val wallpaperInfo: StateFlow<WallpaperInfo?> = MutableStateFlow(null).asStateFlow()
     override val wallpaperSupportsAmbientMode = MutableStateFlow(false).asStateFlow()
+    override var rootView: View? = null
+
+    override fun setNotificationStackAbsoluteBottom(bottom: Float) {}
 }
diff --git a/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/WallpaperRepository.kt b/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/WallpaperRepository.kt
index 041b6f9..203e1da 100644
--- a/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/WallpaperRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/wallpapers/data/repository/WallpaperRepository.kt
@@ -21,10 +21,16 @@
 import android.content.Context
 import android.content.Intent
 import android.content.IntentFilter
+import android.os.Bundle
 import android.os.UserHandle
+import android.view.View
+import androidx.annotation.VisibleForTesting
+import com.android.systemui.Flags
 import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.keyguard.data.repository.KeyguardClockRepository
+import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.user.data.model.SelectedUserModel
 import com.android.systemui.user.data.model.SelectionStatus
 import com.android.systemui.user.data.repository.UserRepository
@@ -32,16 +38,19 @@
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 
 /** A repository storing information about the current wallpaper. */
@@ -51,6 +60,15 @@
 
     /** Emits true if the current user's current wallpaper supports ambient mode. */
     val wallpaperSupportsAmbientMode: StateFlow<Boolean>
+
+    /** Set rootView to get its windowToken afterwards */
+    var rootView: View?
+
+    /**
+     * Set bottom of notifications from notification stack, and Magic Portrait will layout base on
+     * this value
+     */
+    fun setNotificationStackAbsoluteBottom(bottom: Float)
 }
 
 @SysUISingleton
@@ -61,6 +79,8 @@
     @Background private val bgDispatcher: CoroutineDispatcher,
     broadcastDispatcher: BroadcastDispatcher,
     userRepository: UserRepository,
+    keyguardRepository: KeyguardRepository,
+    keyguardClockRepository: KeyguardClockRepository,
     private val wallpaperManager: WallpaperManager,
     context: Context,
 ) : WallpaperRepository {
@@ -69,10 +89,7 @@
 
     private val wallpaperChanged: Flow<Unit> =
         broadcastDispatcher
-            .broadcastFlow(
-                IntentFilter(Intent.ACTION_WALLPAPER_CHANGED),
-                user = UserHandle.ALL,
-            )
+            .broadcastFlow(IntentFilter(Intent.ACTION_WALLPAPER_CHANGED), user = UserHandle.ALL)
             // The `combine` defining `wallpaperSupportsAmbientMode` will not run until both of the
             // input flows emit at least once. Since this flow is an input flow, it needs to emit
             // when it starts up to ensure that the `combine` will run if the user changes before we
@@ -87,6 +104,27 @@
             // Only update the wallpaper status once the user selection has finished.
             .filter { it.selectionStatus == SelectionStatus.SELECTION_COMPLETE }
 
+    /** The bottom of notification stack respect to the top of screen. */
+    private val notificationStackAbsoluteBottom: MutableStateFlow<Float> = MutableStateFlow(0F)
+
+    /** The top of shortcut respect to the top of screen. */
+    private val shortcutAbsoluteTop: StateFlow<Float> = keyguardRepository.shortcutAbsoluteTop
+
+    /**
+     * The top of notification stack to give a default state of lockscreen remaining space for
+     * states with notifications to compare with. It's the bottom of smartspace date and weather
+     * smartspace in small clock state, plus proper bottom margin.
+     */
+    private val notificationStackDefaultTop = keyguardClockRepository.notificationDefaultTop
+    @VisibleForTesting var sendLockscreenLayoutJob: Job? = null
+    private val lockscreenRemainingSpaceWithNotification: Flow<Triple<Float, Float, Float>> =
+        combine(
+            notificationStackAbsoluteBottom,
+            notificationStackDefaultTop,
+            shortcutAbsoluteTop,
+            ::Triple,
+        )
+
     override val wallpaperInfo: StateFlow<WallpaperInfo?> =
         if (!wallpaperManager.isWallpaperSupported || !deviceSupportsAodWallpaper) {
             MutableStateFlow(null).asStateFlow()
@@ -116,9 +154,70 @@
                 initialValue = wallpaperInfo.value?.supportsAmbientMode() == true,
             )
 
+    override var rootView: View? = null
+
+    val shouldSendNotificationLayout =
+        wallpaperInfo
+            .map {
+                val shouldSendNotificationLayout = shouldSendNotificationLayout(it)
+                if (shouldSendNotificationLayout) {
+                    sendLockscreenLayoutJob =
+                        scope.launch {
+                            lockscreenRemainingSpaceWithNotification.collect {
+                                (notificationBottom, notificationDefaultTop, shortcutTop) ->
+                                wallpaperManager.sendWallpaperCommand(
+                                    /* windowToken = */ rootView?.windowToken,
+                                    /* action = */ WallpaperManager
+                                        .COMMAND_LOCKSCREEN_LAYOUT_CHANGED,
+                                    /* x = */ 0,
+                                    /* y = */ 0,
+                                    /* z = */ 0,
+                                    /* extras = */ Bundle().apply {
+                                        putFloat("screenLeft", 0F)
+                                        putFloat("smartspaceBottom", notificationDefaultTop)
+                                        putFloat("notificationBottom", notificationBottom)
+                                        putFloat(
+                                            "screenRight",
+                                            context.resources.displayMetrics.widthPixels.toFloat(),
+                                        )
+                                        putFloat("shortCutTop", shortcutTop)
+                                    },
+                                )
+                            }
+                        }
+                } else {
+                    sendLockscreenLayoutJob?.cancel()
+                }
+                shouldSendNotificationLayout
+            }
+            .stateIn(
+                scope,
+                // Always be listening for wallpaper changes.
+                if (Flags.magicPortraitWallpapers()) SharingStarted.Eagerly
+                else SharingStarted.Lazily,
+                initialValue = false,
+            )
+
+    override fun setNotificationStackAbsoluteBottom(bottom: Float) {
+        notificationStackAbsoluteBottom.value = bottom
+    }
+
     private suspend fun getWallpaper(selectedUser: SelectedUserModel): WallpaperInfo? {
         return withContext(bgDispatcher) {
             wallpaperManager.getWallpaperInfoForUser(selectedUser.userInfo.id)
         }
     }
+
+    private fun shouldSendNotificationLayout(wallpaperInfo: WallpaperInfo?): Boolean {
+        return if (wallpaperInfo != null && wallpaperInfo.component != null) {
+            wallpaperInfo.component!!.className == MAGIC_PORTRAIT_CLASSNAME
+        } else {
+            false
+        }
+    }
+
+    companion object {
+        const val MAGIC_PORTRAIT_CLASSNAME =
+            "com.google.android.apps.magicportrait.service.MagicPortraitWallpaperService"
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/wallpapers/domain/interactor/WallpaperInteractor.kt b/packages/SystemUI/src/com/android/systemui/wallpapers/domain/interactor/WallpaperInteractor.kt
new file mode 100644
index 0000000..79ebf01
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/wallpapers/domain/interactor/WallpaperInteractor.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.wallpapers.domain.interactor
+
+import com.android.systemui.wallpapers.data.repository.WallpaperRepository
+import javax.inject.Inject
+
+class WallpaperInteractor @Inject constructor(val wallpaperRepository: WallpaperRepository) {
+    fun setNotificationStackAbsoluteBottom(bottom: Float) {
+        wallpaperRepository.setNotificationStackAbsoluteBottom(bottom)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryImplTest.kt
new file mode 100644
index 0000000..5d5c120
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayScopeRepositoryImplTest.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.systemui.display.data.repository
+
+import android.platform.test.annotations.EnableFlags
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.unconfinedTestDispatcher
+import com.android.systemui.statusbar.core.StatusBarConnectedDisplays
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.isActive
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@EnableFlags(StatusBarConnectedDisplays.FLAG_NAME)
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class DisplayScopeRepositoryImplTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos().also { it.testDispatcher = it.unconfinedTestDispatcher }
+    private val testScope = kosmos.testScope
+    private val fakeDisplayRepository = kosmos.displayRepository
+
+    private val repo =
+        DisplayScopeRepositoryImpl(
+            kosmos.applicationCoroutineScope,
+            kosmos.testDispatcher,
+            fakeDisplayRepository,
+        )
+
+    @Before
+    fun setUp() {
+        repo.start()
+    }
+
+    @Test
+    fun scopeForDisplay_multipleCallsForSameDisplayId_returnsSameInstance() {
+        val scopeForDisplay = repo.scopeForDisplay(displayId = 1)
+
+        assertThat(repo.scopeForDisplay(displayId = 1)).isSameInstanceAs(scopeForDisplay)
+    }
+
+    @Test
+    fun scopeForDisplay_differentDisplayId_returnsNewInstance() {
+        val scopeForDisplay1 = repo.scopeForDisplay(displayId = 1)
+        val scopeForDisplay2 = repo.scopeForDisplay(displayId = 2)
+
+        assertThat(scopeForDisplay1).isNotSameInstanceAs(scopeForDisplay2)
+    }
+
+    @Test
+    fun scopeForDisplay_activeByDefault() =
+        testScope.runTest {
+            val scopeForDisplay = repo.scopeForDisplay(displayId = 1)
+
+            assertThat(scopeForDisplay.isActive).isTrue()
+        }
+
+    @Test
+    fun scopeForDisplay_afterDisplayRemoved_scopeIsCancelled() =
+        testScope.runTest {
+            val scopeForDisplay = repo.scopeForDisplay(displayId = 1)
+
+            fakeDisplayRepository.removeDisplay(displayId = 1)
+
+            assertThat(scopeForDisplay.isActive).isFalse()
+        }
+
+    @Test
+    fun scopeForDisplay_afterDisplayRemoved_returnsNewInstance() =
+        testScope.runTest {
+            val initialScope = repo.scopeForDisplay(displayId = 1)
+
+            fakeDisplayRepository.removeDisplay(displayId = 1)
+
+            val newScope = repo.scopeForDisplay(displayId = 1)
+            assertThat(newScope).isNotSameInstanceAs(initialScope)
+        }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSectionTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSectionTest.kt
index 0b944f0..96a0aad 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSectionTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSectionTest.kt
@@ -39,13 +39,13 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
+import com.android.systemui.shade.LargeScreenHeaderHelper
 import com.android.systemui.shade.data.repository.shadeRepository
 import com.android.systemui.statusbar.notification.stack.domain.interactor.notificationsKeyguardInteractor
 import com.android.systemui.statusbar.policy.fakeConfigurationController
 import com.android.systemui.statusbar.ui.fakeSystemBarUtilsProxy
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.eq
-import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -57,6 +57,7 @@
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyString
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.mock
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @RunWith(AndroidJUnit4::class)
@@ -122,6 +123,7 @@
                     { keyguardBlueprintInteractor },
                     keyguardRootViewModel,
                     aodBurnInViewModel,
+                    largeScreenHeaderHelperLazy = { mock<LargeScreenHeaderHelper>() },
                 )
         }
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
index 1d74e8b..4ede90e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScreenshotDataTest.kt
@@ -53,8 +53,8 @@
 
         assertThat(data.source).isEqualTo(source)
         assertThat(data.type).isEqualTo(type)
-        assertThat(data.screenBounds).isEqualTo(bounds)
-        assertThat(data.insets).isEqualTo(insets)
+        assertThat(data.originalScreenBounds).isEqualTo(bounds)
+        assertThat(data.originalInsets).isEqualTo(insets)
         assertThat(data.taskId).isEqualTo(taskId)
         assertThat(data.userHandle).isEqualTo(UserHandle.of(userId))
         assertThat(data.topComponent).isEqualTo(component)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
index 2fcacb9..0d4cb4c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
@@ -16,6 +16,7 @@
 package com.android.systemui.screenshot.policy
 
 import android.content.ComponentName
+import android.graphics.Bitmap
 import android.graphics.Insets
 import android.graphics.Rect
 import android.os.UserHandle
@@ -27,10 +28,12 @@
 import com.android.systemui.screenshot.ScreenshotData
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.launcherOnly
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.repository.DisplayContentRepository
 import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
 import com.android.systemui.screenshot.policy.TestUserIds.WORK
+import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
@@ -39,14 +42,6 @@
 
 @RunWith(AndroidJUnit4::class)
 class PolicyRequestProcessorTest {
-
-    val imageCapture =
-        object : ImageCapture {
-            override fun captureDisplay(displayId: Int, crop: Rect?) = null
-
-            override suspend fun captureTask(taskId: Int) = null
-        }
-
     /** Tests behavior when no policies are applied */
     @Test
     fun testProcess_defaultOwner_whenNoPolicyApplied() {
@@ -60,9 +55,9 @@
                 SCREENSHOT_KEY_CHORD,
                 UserHandle.CURRENT,
                 topComponent = null,
-                screenBounds = Rect(0, 0, 1, 1),
+                originalScreenBounds = Rect(0, 0, 1, 1),
                 taskId = -1,
-                insets = Insets.NONE,
+                originalInsets = Insets.NONE,
                 bitmap = null,
                 displayId = DEFAULT_DISPLAY,
             )
@@ -71,7 +66,7 @@
         val requestProcessor =
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
-                imageCapture,
+                createImageCapture(),
                 policies = emptyList(),
                 defaultOwner = UserHandle.of(PERSONAL),
                 defaultComponent = ComponentName("default", "Component"),
@@ -91,6 +86,70 @@
         assertWithMessage("Task ID").that(result.taskId).isEqualTo(TASK_ID)
     }
 
+    @Test
+    fun testProcess_throwsWhenCaptureFails() {
+        val request = ScreenshotData.forTesting()
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(display = null),
+                policies = emptyList(),
+                defaultComponent = ComponentName("default", "Component"),
+                displayTasks = DisplayContentRepository { launcherOnly() },
+            )
+
+        val result = runCatching { runBlocking { requestProcessor.process(request) } }
+
+        assertThat(result.isFailure).isTrue()
+    }
+
+    @Test
+    fun testProcess_throwsWhenTaskCaptureFails() {
+        val request = ScreenshotData.forTesting()
+        val fullScreenWork = DisplayContentRepository {
+            singleFullScreen(TaskSpec(taskId = TASK_ID, name = FILES, userId = WORK))
+        }
+
+        val captureTaskPolicy = CapturePolicy {
+            CapturePolicy.PolicyResult.Matched(
+                policy = "",
+                reason = "",
+                parameters =
+                    CaptureParameters(
+                        CaptureType.IsolatedTask(taskId = 0, taskBounds = null),
+                        null,
+                        UserHandle.CURRENT,
+                    ),
+            )
+        }
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(task = null),
+                policies = listOf(captureTaskPolicy),
+                defaultComponent = ComponentName("default", "Component"),
+                displayTasks = fullScreenWork,
+            )
+
+        val result = runCatching { runBlocking { requestProcessor.process(request) } }
+
+        assertThat(result.isFailure).isTrue()
+    }
+
+    private fun createImageCapture(
+        display: Bitmap? = Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888),
+        task: Bitmap? = Bitmap.createBitmap(100, 100, Bitmap.Config.ARGB_8888),
+    ) =
+        object : ImageCapture {
+            override fun captureDisplay(displayId: Int, crop: Rect?) = display
+
+            override suspend fun captureTask(taskId: Int) = task
+        }
+
     companion object {
         const val TASK_ID = 1001
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
index 7cd306e..6425da4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutControllerTest.java
@@ -83,6 +83,7 @@
 import com.android.systemui.statusbar.SysuiStatusBarStateController;
 import com.android.systemui.statusbar.notification.ColorUpdateLogger;
 import com.android.systemui.statusbar.notification.DynamicPrivacyController;
+import com.android.systemui.statusbar.notification.HeadsUpTouchHelper;
 import com.android.systemui.statusbar.notification.NotificationWakeUpCoordinator;
 import com.android.systemui.statusbar.notification.collection.NotifCollection;
 import com.android.systemui.statusbar.notification.collection.NotifPipeline;
@@ -103,7 +104,6 @@
 import com.android.systemui.statusbar.notification.stack.NotificationSwipeHelper.NotificationCallback;
 import com.android.systemui.statusbar.notification.stack.ui.viewbinder.NotificationListViewBinder;
 import com.android.systemui.statusbar.phone.HeadsUpAppearanceController;
-import com.android.systemui.statusbar.notification.HeadsUpTouchHelper;
 import com.android.systemui.statusbar.phone.KeyguardBypassController;
 import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
@@ -113,6 +113,7 @@
 import com.android.systemui.statusbar.policy.ZenModeController;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.util.settings.SecureSettings;
+import com.android.systemui.wallpapers.domain.interactor.WallpaperInteractor;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -154,6 +155,7 @@
     @Mock private KeyguardBypassController mKeyguardBypassController;
     @Mock private PowerInteractor mPowerInteractor;
     @Mock private PrimaryBouncerInteractor mPrimaryBouncerInteractor;
+    @Mock private WallpaperInteractor mWallpaperInteractor;
     @Mock private NotificationLockscreenUserManager mNotificationLockscreenUserManager;
     @Mock private MetricsLogger mMetricsLogger;
     @Mock private ColorUpdateLogger mColorUpdateLogger;
@@ -1070,7 +1072,8 @@
                 mock(NotificationDismissibilityProvider.class),
                 mActivityStarter,
                 new ResourcesSplitShadeStateController(),
-                mSensitiveNotificationProtectionController);
+                mSensitiveNotificationProtectionController,
+                mWallpaperInteractor);
     }
 
     static class LogMatcher implements ArgumentMatcher<LogMaker> {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
index 8a3e551..59fc0d1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java
@@ -102,6 +102,7 @@
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
 import com.android.systemui.statusbar.policy.AvalancheController;
 import com.android.systemui.statusbar.policy.ResourcesSplitShadeStateController;
+import com.android.systemui.wallpapers.domain.interactor.WallpaperInteractor;
 
 import kotlin.Unit;
 
@@ -146,6 +147,7 @@
     @Mock private NotificationStackScrollLayoutController mStackScrollLayoutController;
     @Mock private ScreenOffAnimationController mScreenOffAnimationController;
     @Mock private NotificationShelf mNotificationShelf;
+    @Mock private WallpaperInteractor mWallpaperInteractor;
     @Mock private NotificationStackSizeCalculator mStackSizeCalculator;
     @Mock private StatusBarKeyguardViewManager mStatusBarKeyguardViewManager;
     @Mock private LargeScreenShadeInterpolator mLargeScreenShadeInterpolator;
@@ -208,6 +210,7 @@
                 .thenReturn(mNotificationRoundnessManager);
         mStackScroller.setController(mStackScrollLayoutController);
         mStackScroller.setShelf(mNotificationShelf);
+        mStackScroller.setWallpaperInteractor(mWallpaperInteractor);
         when(mStackScroller.getExpandHelper()).thenReturn(mExpandHelper);
 
         doNothing().when(mGroupExpansionManager).collapseGroups();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.java
deleted file mode 100644
index 637a0f1..0000000
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/ZenModeControllerImplTest.java
+++ /dev/null
@@ -1,201 +0,0 @@
-/*
- * Copyright (C) 2017 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
- * except in compliance with the License. You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software distributed under the
- * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the specific language governing
- * permissions and limitations under the License.
- */
-
-package com.android.systemui.statusbar.policy;
-
-import static junit.framework.Assert.assertEquals;
-import static junit.framework.Assert.assertFalse;
-import static junit.framework.Assert.assertTrue;
-
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import android.app.NotificationManager;
-import android.os.Handler;
-import android.provider.Settings;
-import android.service.notification.ZenModeConfig;
-import android.testing.TestableLooper;
-import android.testing.TestableLooper.RunWithLooper;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
-import com.android.systemui.SysuiTestCase;
-import com.android.systemui.broadcast.BroadcastDispatcher;
-import com.android.systemui.dump.DumpManager;
-import com.android.systemui.settings.UserTracker;
-import com.android.systemui.statusbar.policy.ZenModeController.Callback;
-import com.android.systemui.util.settings.FakeGlobalSettings;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
-import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
-
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-@RunWithLooper
-public class ZenModeControllerImplTest extends SysuiTestCase {
-
-    private Callback mCallback;
-    @Mock
-    NotificationManager mNm;
-    @Mock
-    ZenModeConfig mConfig;
-    @Mock
-    BroadcastDispatcher mBroadcastDispatcher;
-    @Mock
-    DumpManager mDumpManager;
-    @Mock
-    UserTracker mUserTracker;
-    private ZenModeControllerImpl mController;
-
-    private final FakeGlobalSettings mGlobalSettings = new FakeGlobalSettings();
-
-    @Before
-    public void setUp() {
-        MockitoAnnotations.initMocks(this);
-        mContext.addMockSystemService(NotificationManager.class, mNm);
-        when(mNm.getZenModeConfig()).thenReturn(mConfig);
-
-        mController = new ZenModeControllerImpl(
-                mContext,
-                Handler.createAsync(TestableLooper.get(this).getLooper()),
-                Handler.createAsync(TestableLooper.get(this).getLooper()),
-                mBroadcastDispatcher,
-                mDumpManager,
-                mGlobalSettings,
-                mUserTracker);
-    }
-
-    @Test
-    public void testRemoveDuringCallback() {
-        mCallback = new Callback() {
-            @Override
-            public void onConfigChanged(ZenModeConfig config) {
-                mController.removeCallback(mCallback);
-            }
-        };
-        mController.addCallback(mCallback);
-        Callback mockCallback = mock(Callback.class);
-        mController.addCallback(mockCallback);
-        mController.fireConfigChanged(null);
-        verify(mockCallback).onConfigChanged(eq(null));
-    }
-
-    @Test
-    public void testAreNotificationsHiddenInShade_zenOffShadeSuppressed() {
-        mConfig.suppressedVisualEffects =
-                NotificationManager.Policy.SUPPRESSED_EFFECT_NOTIFICATION_LIST;
-        mController.updateZenMode(Settings.Global.ZEN_MODE_OFF);
-        mController.updateZenModeConfig();
-
-        assertFalse(mController.areNotificationsHiddenInShade());
-    }
-
-    @Test
-    public void testAreNotificationsHiddenInShade_zenOnShadeNotSuppressed() {
-        NotificationManager.Policy policy = new NotificationManager.Policy(0, 0, 0,
-                NotificationManager.Policy.SUPPRESSED_EFFECT_STATUS_BAR);
-        when(mNm.getConsolidatedNotificationPolicy()).thenReturn(policy);
-        mController.updateConsolidatedNotificationPolicy();
-        mController.updateZenMode(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS);
-
-        assertFalse(mController.areNotificationsHiddenInShade());
-    }
-
-    @Test
-    public void testAreNotificationsHiddenInShade_zenOnShadeSuppressed() {
-        NotificationManager.Policy policy = new NotificationManager.Policy(0, 0, 0,
-                NotificationManager.Policy.SUPPRESSED_EFFECT_NOTIFICATION_LIST);
-        when(mNm.getConsolidatedNotificationPolicy()).thenReturn(policy);
-        mController.updateConsolidatedNotificationPolicy();
-        mController.updateZenMode(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS);
-
-        assertTrue(mController.areNotificationsHiddenInShade());
-    }
-
-    @Test
-    public void testModeChange() {
-        List<Integer> states = List.of(
-                Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS,
-                Settings.Global.ZEN_MODE_NO_INTERRUPTIONS,
-                Settings.Global.ZEN_MODE_ALARMS,
-                Settings.Global.ZEN_MODE_ALARMS
-        );
-
-        for (Integer state : states) {
-            mGlobalSettings.putInt(Settings.Global.ZEN_MODE, state);
-            TestableLooper.get(this).processAllMessages();
-            assertEquals(state.intValue(), mController.getZen());
-        }
-    }
-
-    @Test
-    public void testModeChange_callbackNotified() {
-        final AtomicInteger currentState = new AtomicInteger(-1);
-
-        ZenModeController.Callback callback = new Callback() {
-            @Override
-            public void onZenChanged(int zen) {
-                currentState.set(zen);
-            }
-        };
-
-        mController.addCallback(callback);
-
-        List<Integer> states = List.of(
-                Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS,
-                Settings.Global.ZEN_MODE_NO_INTERRUPTIONS,
-                Settings.Global.ZEN_MODE_ALARMS,
-                Settings.Global.ZEN_MODE_ALARMS
-        );
-
-        for (Integer state : states) {
-            mGlobalSettings.putInt(Settings.Global.ZEN_MODE, state);
-            TestableLooper.get(this).processAllMessages();
-            assertEquals(state.intValue(), currentState.get());
-        }
-
-    }
-
-    @Test
-    public void testCallbackRemovedWhileDispatching_doesntCrash() {
-        final AtomicBoolean remove = new AtomicBoolean(false);
-        mGlobalSettings.putInt(Settings.Global.ZEN_MODE, Settings.Global.ZEN_MODE_OFF);
-        TestableLooper.get(this).processAllMessages();
-        final ZenModeController.Callback callback = new ZenModeController.Callback() {
-            @Override
-            public void onZenChanged(int zen) {
-                if (remove.get()) {
-                    mController.removeCallback(this);
-                }
-            }
-        };
-        mController.addCallback(callback);
-        mController.addCallback(new ZenModeController.Callback() {});
-
-        remove.set(true);
-
-        mGlobalSettings.putInt(Settings.Global.ZEN_MODE, Settings.Global.ZEN_MODE_NO_INTERRUPTIONS);
-        TestableLooper.get(this).processAllMessages();
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wallpapers/data/repository/WallpaperRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/wallpapers/data/repository/WallpaperRepositoryImplTest.kt
index bdecf2b..b8dd334 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wallpapers/data/repository/WallpaperRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/wallpapers/data/repository/WallpaperRepositoryImplTest.kt
@@ -18,18 +18,21 @@
 
 import android.app.WallpaperInfo
 import android.app.WallpaperManager
+import android.content.ComponentName
 import android.content.Intent
 import android.content.pm.UserInfo
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.keyguard.data.repository.FakeKeyguardClockRepository
+import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
 import com.android.systemui.user.data.model.SelectedUserModel
 import com.android.systemui.user.data.model.SelectionStatus
 import com.android.systemui.user.data.repository.FakeUserRepository
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
+import com.android.systemui.wallpapers.data.repository.WallpaperRepositoryImpl.Companion.MAGIC_PORTRAIT_CLASSNAME
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
@@ -39,6 +42,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -48,6 +54,8 @@
     private val testDispatcher = StandardTestDispatcher()
     private val testScope = TestScope(testDispatcher)
     private val userRepository = FakeUserRepository()
+    private val keyguardClockRepository = FakeKeyguardClockRepository()
+    private val keyguardRepository = FakeKeyguardRepository()
     private val wallpaperManager: WallpaperManager = mock()
 
     private val underTest: WallpaperRepositoryImpl by lazy {
@@ -56,6 +64,8 @@
             testDispatcher,
             fakeBroadcastDispatcher,
             userRepository,
+            keyguardRepository,
+            keyguardClockRepository,
             wallpaperManager,
             context,
         )
@@ -219,7 +229,7 @@
         testScope.runTest {
             context.orCreateTestableResources.addOverride(
                 com.android.internal.R.bool.config_dozeSupportsAodWallpaper,
-                false
+                false,
             )
 
             val latest by collectLastValue(underTest.wallpaperInfo)
@@ -407,7 +417,7 @@
         testScope.runTest {
             context.orCreateTestableResources.addOverride(
                 com.android.internal.R.bool.config_dozeSupportsAodWallpaper,
-                false
+                false,
             )
 
             val latest by collectLastValue(underTest.wallpaperSupportsAmbientMode)
@@ -425,6 +435,54 @@
             assertThat(latest).isFalse()
         }
 
+    @Test
+    @EnableFlags(Flags.FLAG_MAGIC_PORTRAIT_WALLPAPERS)
+    fun shouldSendNotificationLayout_setMagicPortraitWallpaper_launchSendLayoutJob() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.shouldSendNotificationLayout)
+            val magicPortraitWallpaper =
+                mock<WallpaperInfo>().apply {
+                    whenever(this.component)
+                        .thenReturn(ComponentName(context, MAGIC_PORTRAIT_CLASSNAME))
+                }
+            whenever(wallpaperManager.getWallpaperInfoForUser(any()))
+                .thenReturn(magicPortraitWallpaper)
+            fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
+                context,
+                Intent(Intent.ACTION_WALLPAPER_CHANGED),
+            )
+            assertThat(latest).isTrue()
+            assertThat(underTest.sendLockscreenLayoutJob).isNotNull()
+            assertThat(underTest.sendLockscreenLayoutJob!!.isActive).isEqualTo(true)
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MAGIC_PORTRAIT_WALLPAPERS)
+    fun shouldSendNotificationLayout_setNotMagicPortraitWallpaper_cancelSendLayoutJob() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.shouldSendNotificationLayout)
+            val magicPortraitWallpaper = MAGIC_PORTRAIT_WP
+            whenever(wallpaperManager.getWallpaperInfoForUser(any()))
+                .thenReturn(magicPortraitWallpaper)
+            fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
+                context,
+                Intent(Intent.ACTION_WALLPAPER_CHANGED),
+            )
+            assertThat(latest).isTrue()
+            assertThat(underTest.sendLockscreenLayoutJob).isNotNull()
+            assertThat(underTest.sendLockscreenLayoutJob!!.isActive).isEqualTo(true)
+
+            val nonMagicPortraitWallpaper = UNSUPPORTED_WP
+            whenever(wallpaperManager.getWallpaperInfoForUser(any()))
+                .thenReturn(nonMagicPortraitWallpaper)
+            fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
+                context,
+                Intent(Intent.ACTION_WALLPAPER_CHANGED),
+            )
+            assertThat(latest).isFalse()
+            assertThat(underTest.sendLockscreenLayoutJob?.isCancelled).isEqualTo(true)
+        }
+
     private companion object {
         val USER_WITH_UNSUPPORTED_WP = UserInfo(/* id= */ 3, /* name= */ "user3", /* flags= */ 0)
         val UNSUPPORTED_WP =
@@ -433,5 +491,10 @@
         val USER_WITH_SUPPORTED_WP = UserInfo(/* id= */ 4, /* name= */ "user4", /* flags= */ 0)
         val SUPPORTED_WP =
             mock<WallpaperInfo>().apply { whenever(this.supportsAmbientMode()).thenReturn(true) }
+
+        val MAGIC_PORTRAIT_WP =
+            mock<WallpaperInfo>().apply {
+                whenever(this.component).thenReturn(ComponentName("", MAGIC_PORTRAIT_CLASSNAME))
+            }
     }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/app/WallpaperManagerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/app/WallpaperManagerKosmos.kt
new file mode 100644
index 0000000..2850ab7
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/app/WallpaperManagerKosmos.kt
@@ -0,0 +1,27 @@
+/*
+ * 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.app
+
+import android.app.WallpaperManager
+import android.content.applicationContext
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+
+@OptIn(ExperimentalCoroutinesApi::class)
+val Kosmos.wallpaperManager: WallpaperManager by Fixture {
+    WallpaperManager.getInstance(applicationContext)
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
index 6c3cf91..fcc83b3 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
@@ -24,16 +24,12 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import org.mockito.Mockito.`when` as whenever
 
 /** Creates a mock display. */
-fun display(
-    type: Int,
-    flags: Int = 0,
-    id: Int = 0,
-    state: Int? = null,
-): Display {
+fun display(type: Int, flags: Int = 0, id: Int = 0, state: Int? = null): Display {
     return mock {
         whenever(this.displayId).thenReturn(id)
         whenever(this.type).thenReturn(type)
@@ -51,10 +47,21 @@
 @SysUISingleton
 /** Fake [DisplayRepository] implementation for testing. */
 class FakeDisplayRepository @Inject constructor() : DisplayRepository {
-    private val flow = MutableSharedFlow<Set<Display>>(replay = 1)
+    private val flow = MutableStateFlow<Set<Display>>(emptySet())
     private val pendingDisplayFlow =
         MutableSharedFlow<DisplayRepository.PendingDisplay?>(replay = 1)
-    private val displayAdditionEventFlow = MutableSharedFlow<Display?>(replay = 1)
+    private val displayAdditionEventFlow = MutableSharedFlow<Display?>(replay = 0)
+    private val displayRemovalEventFlow = MutableSharedFlow<Int>(replay = 0)
+
+    suspend fun addDisplay(display: Display) {
+        flow.value += display
+        displayAdditionEventFlow.emit(display)
+    }
+
+    suspend fun removeDisplay(displayId: Int) {
+        flow.value = flow.value.filter { it.displayId != displayId }.toSet()
+        displayRemovalEventFlow.emit(displayId)
+    }
 
     /** Emits [value] as [displayAdditionEvent] flow value. */
     suspend fun emit(value: Display?) = displayAdditionEventFlow.emit(value)
@@ -65,7 +72,7 @@
     /** Emits [value] as [pendingDisplay] flow value. */
     suspend fun emit(value: DisplayRepository.PendingDisplay?) = pendingDisplayFlow.emit(value)
 
-    override val displays: Flow<Set<Display>>
+    override val displays: StateFlow<Set<Display>>
         get() = flow
 
     override val pendingDisplay: Flow<DisplayRepository.PendingDisplay?>
@@ -78,8 +85,11 @@
     override val displayAdditionEvent: Flow<Display?>
         get() = displayAdditionEventFlow
 
+    override val displayRemovalEvent: Flow<Int> = displayRemovalEventFlow
+
     private val _displayChangeEvent = MutableSharedFlow<Int>(replay = 1)
     override val displayChangeEvent: Flow<Int> = _displayChangeEvent
+
     suspend fun emitDisplayChangeEvent(displayId: Int) = _displayChangeEvent.emit(displayId)
 
     fun setDefaultDisplayOff(defaultDisplayOff: Boolean) {
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt
index 5e5f8cb..159dd34 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardClockRepository.kt
@@ -46,16 +46,27 @@
     private val _previewClock = MutableStateFlow(Mockito.mock(ClockController::class.java))
     override val previewClock: Flow<ClockController>
         get() = _previewClock
+
+    private val _notificationDefaultTop = MutableStateFlow(0F)
+    override val notificationDefaultTop: StateFlow<Float>
+        get() = _notificationDefaultTop
+
     override val clockEventController: ClockEventController
         get() = mock()
+
     override val shouldForceSmallClock: Boolean
         get() = _shouldForceSmallClock
+
     private var _shouldForceSmallClock: Boolean = false
 
     override fun setClockSize(size: ClockSize) {
         _clockSize.value = size
     }
 
+    override fun setNotificationDefaultTop(top: Float) {
+        _notificationDefaultTop.value = top
+    }
+
     fun setSelectedClockSize(size: ClockSizeSetting) {
         _selectedClockSize.value = size
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
index 54a6c0c..e513e8d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardRepository.kt
@@ -131,6 +131,10 @@
     private val _isEncryptedOrLockdown = MutableStateFlow(true)
     override val isEncryptedOrLockdown: Flow<Boolean> = _isEncryptedOrLockdown
 
+    private val _shortcutAbsoluteTop = MutableStateFlow(0F)
+    override val shortcutAbsoluteTop: StateFlow<Float>
+        get() = _shortcutAbsoluteTop.asStateFlow()
+
     private val _isKeyguardEnabled = MutableStateFlow(true)
     override val isKeyguardEnabled: StateFlow<Boolean> = _isKeyguardEnabled.asStateFlow()
 
@@ -241,7 +245,7 @@
 
     override fun setBiometricUnlockState(
         mode: BiometricUnlockMode,
-        source: BiometricUnlockSource?
+        source: BiometricUnlockSource?,
     ) {
         _biometricUnlockState.tryEmit(BiometricUnlockModel(mode, source))
     }
@@ -294,6 +298,10 @@
         return isShowKeyguardWhenReenabled
     }
 
+    override fun setShortcutAbsoluteTop(top: Float) {
+        _shortcutAbsoluteTop.value = top
+    }
+
     override fun setCanIgnoreAuthAndReturnToGone(canWake: Boolean) {
         _canIgnoreAuthAndReturnToGone.value = canWake
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryKosmos.kt
index 12d7c49..49a8c18 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryKosmos.kt
@@ -29,9 +29,10 @@
 import com.android.systemui.keyguard.ui.viewmodel.keyguardRootViewModel
 import com.android.systemui.keyguard.ui.viewmodel.keyguardSmartspaceViewModel
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.util.mockito.mock
+import com.android.systemui.shade.LargeScreenHeaderHelper
 import java.util.Optional
 import org.mockito.Mockito.spy
+import org.mockito.kotlin.mock
 
 val Kosmos.keyguardClockSection: ClockSection by
     Kosmos.Fixture {
@@ -43,6 +44,7 @@
             blueprintInteractor = { keyguardBlueprintInteractor },
             rootViewModel = keyguardRootViewModel,
             aodBurnInViewModel = aodBurnInViewModel,
+            largeScreenHeaderHelperLazy = { mock<LargeScreenHeaderHelper>() },
         )
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt
index d52883e..bdb9abb 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardClockInteractorKosmos.kt
@@ -27,13 +27,13 @@
 val Kosmos.keyguardClockInteractor by
     Kosmos.Fixture {
         KeyguardClockInteractor(
-            keyguardClockRepository = keyguardClockRepository,
-            applicationScope = applicationCoroutineScope,
             mediaCarouselInteractor = mediaCarouselInteractor,
             activeNotificationsInteractor = activeNotificationsInteractor,
             shadeInteractor = shadeInteractor,
             keyguardInteractor = keyguardInteractor,
             keyguardTransitionInteractor = keyguardTransitionInteractor,
             headsUpNotificationInteractor = headsUpNotificationInteractor,
+            applicationScope = applicationCoroutineScope,
+            keyguardClockRepository = keyguardClockRepository,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/wallpapers/data/repository/WallpaperRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/wallpapers/data/repository/WallpaperRepositoryKosmos.kt
new file mode 100644
index 0000000..1d8c891
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/wallpapers/data/repository/WallpaperRepositoryKosmos.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.wallpapers.data.repository
+
+import android.content.applicationContext
+import com.android.app.wallpaperManager
+import com.android.systemui.broadcast.broadcastDispatcher
+import com.android.systemui.keyguard.data.repository.keyguardClockRepository
+import com.android.systemui.keyguard.data.repository.keyguardRepository
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.user.data.repository.userRepository
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+
+@OptIn(ExperimentalCoroutinesApi::class)
+val Kosmos.wallpaperRepository by Fixture {
+    WallpaperRepositoryImpl(
+        context = applicationContext,
+        scope = testScope,
+        bgDispatcher = testDispatcher,
+        broadcastDispatcher = broadcastDispatcher,
+        userRepository = userRepository,
+        wallpaperManager = wallpaperManager,
+        keyguardClockRepository = keyguardClockRepository,
+        keyguardRepository = keyguardRepository,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/wallpapers/domain/interactor/WallpaperInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/wallpapers/domain/interactor/WallpaperInteractorKosmos.kt
new file mode 100644
index 0000000..5278351
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/wallpapers/domain/interactor/WallpaperInteractorKosmos.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.wallpapers.domain.interactor
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.wallpapers.data.repository.wallpaperRepository
+
+val Kosmos.wallpaperInteractor by
+    Kosmos.Fixture { WallpaperInteractor(wallpaperRepository = wallpaperRepository) }
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 42f69e9..95281c8 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -676,7 +676,7 @@
 
             final MacAddress macAddressObj = MacAddress.fromString(macAddress);
             mAssociationRequestsProcessor.createAssociation(userId, packageName, macAddressObj,
-                    null, null, null, false, null, null);
+                    null, null, null, false, null, null, null);
         }
 
         private void checkCanCallNotificationApi(String callingPackage, int userId) {
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
index 4fc9d55..2804945 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
@@ -106,8 +106,8 @@
                     boolean selfManaged = getNextBooleanArg();
                     final MacAddress macAddress = MacAddress.fromString(address);
                     mAssociationRequestsProcessor.createAssociation(userId, packageName, macAddress,
-                            deviceProfile, deviceProfile, /* associatedDevice */ null, selfManaged,
-                            /* callback */ null, /* resultReceiver */ null);
+                            deviceProfile, deviceProfile, /* associatedDevice */ null, false,
+                            /* callback */ null, /* resultReceiver */ null, /* deviceIcon */ null);
                 }
                 break;
 
diff --git a/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java b/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
index 0c54720..77b1780 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
@@ -17,10 +17,12 @@
 package com.android.server.companion.association;
 
 import static com.android.internal.util.XmlUtils.readBooleanAttribute;
+import static com.android.internal.util.XmlUtils.readByteArrayAttribute;
 import static com.android.internal.util.XmlUtils.readIntAttribute;
 import static com.android.internal.util.XmlUtils.readLongAttribute;
 import static com.android.internal.util.XmlUtils.readStringAttribute;
 import static com.android.internal.util.XmlUtils.writeBooleanAttribute;
+import static com.android.internal.util.XmlUtils.writeByteArrayAttribute;
 import static com.android.internal.util.XmlUtils.writeIntAttribute;
 import static com.android.internal.util.XmlUtils.writeLongAttribute;
 import static com.android.internal.util.XmlUtils.writeStringAttribute;
@@ -36,6 +38,7 @@
 import android.annotation.SuppressLint;
 import android.annotation.UserIdInt;
 import android.companion.AssociationInfo;
+import android.graphics.drawable.Icon;
 import android.net.MacAddress;
 import android.os.Environment;
 import android.util.AtomicFile;
@@ -51,6 +54,7 @@
 import org.xmlpull.v1.XmlSerializer;
 
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
@@ -172,6 +176,7 @@
     private static final String XML_ATTR_TIME_APPROVED = "time_approved";
     private static final String XML_ATTR_LAST_TIME_CONNECTED = "last_time_connected";
     private static final String XML_ATTR_SYSTEM_DATA_SYNC_FLAGS = "system_data_sync_flags";
+    private static final String XML_ATTR_DEVICE_ICON = "device_icon";
 
     private static final String LEGACY_XML_ATTR_DEVICE = "device";
 
@@ -393,7 +398,7 @@
         return new AssociationInfo(associationId, userId, appPackage, tag,
                 MacAddress.fromString(deviceAddress), null, profile, null,
                 /* managedByCompanionApp */ false, notify, /* revoked */ false, /* pending */ false,
-                timeApproved, Long.MAX_VALUE, /* systemDataSyncFlags */ 0);
+                timeApproved, Long.MAX_VALUE, /* systemDataSyncFlags */ 0, null);
     }
 
     private static Associations readAssociationsV1(@NonNull TypedXmlPullParser parser,
@@ -444,10 +449,12 @@
                 parser, XML_ATTR_LAST_TIME_CONNECTED, Long.MAX_VALUE);
         final int systemDataSyncFlags = readIntAttribute(parser,
                 XML_ATTR_SYSTEM_DATA_SYNC_FLAGS, 0);
+        final Icon deviceIcon = byteArrayToIcon(
+                readByteArrayAttribute(parser, XML_ATTR_DEVICE_ICON));
 
         return new AssociationInfo(associationId, userId, appPackage, tag, macAddress, displayName,
                 profile, null, selfManaged, notify, revoked, pending, timeApproved,
-                lastTimeConnected, systemDataSyncFlags);
+                lastTimeConnected, systemDataSyncFlags, deviceIcon);
     }
 
     private static void writeAssociations(@NonNull XmlSerializer parent,
@@ -480,6 +487,8 @@
         writeLongAttribute(
                 serializer, XML_ATTR_LAST_TIME_CONNECTED, a.getLastTimeConnectedMs());
         writeIntAttribute(serializer, XML_ATTR_SYSTEM_DATA_SYNC_FLAGS, a.getSystemDataSyncFlags());
+        writeByteArrayAttribute(
+                serializer, XML_ATTR_DEVICE_ICON, iconToByteArray(a.getDeviceIcon()));
 
         serializer.endTag(null, XML_TAG_ASSOCIATION);
     }
@@ -494,4 +503,24 @@
     private static @Nullable MacAddress stringToMacAddress(@Nullable String address) {
         return address != null ? MacAddress.fromString(address) : null;
     }
+
+    private static byte[] iconToByteArray(Icon deviceIcon)
+            throws IOException {
+        if (deviceIcon == null) {
+            return null;
+        }
+
+        ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
+        deviceIcon.writeToStream(byteStream);
+        return byteStream.toByteArray();
+    }
+
+    private static Icon byteArrayToIcon(byte[] bytes) throws IOException {
+        if (bytes == null) {
+            return null;
+        }
+
+        ByteArrayInputStream byteStream = new ByteArrayInputStream(bytes);
+        return Icon.createFromStream(byteStream);
+    }
 }
diff --git a/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java b/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
index d56f17b..aebd11a 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
@@ -48,6 +48,7 @@
 import android.content.Intent;
 import android.content.IntentSender;
 import android.content.pm.PackageManagerInternal;
+import android.graphics.drawable.Icon;
 import android.net.MacAddress;
 import android.os.Binder;
 import android.os.Bundle;
@@ -281,7 +282,7 @@
             createAssociation(userId, packageName, macAddress, request.getDisplayName(),
                     request.getDeviceProfile(), request.getAssociatedDevice(),
                     request.isSelfManaged(),
-                    callback, resultReceiver);
+                    callback, resultReceiver, request.getDeviceIcon());
         });
     }
 
@@ -292,15 +293,15 @@
             @Nullable MacAddress macAddress, @Nullable CharSequence displayName,
             @Nullable String deviceProfile, @Nullable AssociatedDevice associatedDevice,
             boolean selfManaged, @Nullable IAssociationRequestCallback callback,
-            @Nullable ResultReceiver resultReceiver) {
+            @Nullable ResultReceiver resultReceiver, @Nullable Icon deviceIcon) {
         final int id = mAssociationStore.getNextId();
         final long timestamp = System.currentTimeMillis();
 
         final AssociationInfo association = new AssociationInfo(id, userId, packageName,
                 /* tag */ null, macAddress, displayName, deviceProfile, associatedDevice,
                 selfManaged, /* notifyOnDeviceNearby */ false, /* revoked */ false,
-                /* pending */ false, timestamp, Long.MAX_VALUE, /* systemDataSyncFlags */ 0);
-
+                /* pending */ false, timestamp, Long.MAX_VALUE, /* systemDataSyncFlags */ 0,
+                deviceIcon);
         // Add role holder for association (if specified) and add new association to store.
         maybeGrantRoleAndStoreAssociation(association, callback, resultReceiver);
     }
diff --git a/services/companion/java/com/android/server/companion/utils/PermissionsUtils.java b/services/companion/java/com/android/server/companion/utils/PermissionsUtils.java
index c927cd0..f37e0c9 100644
--- a/services/companion/java/com/android/server/companion/utils/PermissionsUtils.java
+++ b/services/companion/java/com/android/server/companion/utils/PermissionsUtils.java
@@ -86,7 +86,7 @@
             @NonNull AssociationRequest request, int packageUid) {
         enforcePermissionForRequestingProfile(context, request.getDeviceProfile(), packageUid);
 
-        if (request.isSelfManaged()) {
+        if (request.isSelfManaged() || request.getDeviceIcon() != null) {
             enforcePermissionForRequestingSelfManaged(context, packageUid);
         }
     }
diff --git a/services/core/java/com/android/server/SystemConfig.java b/services/core/java/com/android/server/SystemConfig.java
index d80e40c..504137a 100644
--- a/services/core/java/com/android/server/SystemConfig.java
+++ b/services/core/java/com/android/server/SystemConfig.java
@@ -727,10 +727,8 @@
             return;
         }
         // Read configuration of features, libs and priv-app permissions from apex module.
-        int apexPermissionFlag = ALLOW_LIBS | ALLOW_FEATURES | ALLOW_PRIVAPP_PERMISSIONS;
-        if (android.permission.flags.Flags.apexSignaturePermissionAllowlistEnabled()) {
-            apexPermissionFlag |= ALLOW_SIGNATURE_PERMISSIONS;
-        }
+        int apexPermissionFlag = ALLOW_LIBS | ALLOW_FEATURES | ALLOW_PRIVAPP_PERMISSIONS
+                | ALLOW_SIGNATURE_PERMISSIONS;
         // TODO: Use a solid way to filter apex module folders?
         for (File f: FileUtils.listFilesOrEmpty(Environment.getApexDirectory())) {
             if (f.isFile() || f.getPath().contains("@")) {
diff --git a/services/core/java/com/android/server/am/ActiveServices.java b/services/core/java/com/android/server/am/ActiveServices.java
index 3666524..3bfbc55 100644
--- a/services/core/java/com/android/server/am/ActiveServices.java
+++ b/services/core/java/com/android/server/am/ActiveServices.java
@@ -1189,8 +1189,8 @@
             if (DEBUG_SERVICE) Slog.v(TAG_SERVICE, "START SERVICE WHILE RESTART PENDING: " + r);
         }
         final boolean wasStartRequested = r.startRequested;
-        r.lastActivity = SystemClock.uptimeMillis();
-        r.startRequested = true;
+        mAm.mProcessStateController.setServiceLastActivityTime(r, SystemClock.uptimeMillis());
+        mAm.mProcessStateController.setStartRequested(r, true);
         r.delayedStop = false;
         r.fgRequired = fgRequired;
         r.pendingStarts.add(new ServiceRecord.StartItem(r, false, r.makeNextStartId(),
@@ -1623,7 +1623,7 @@
             FrameworkStatsLog.write(FrameworkStatsLog.SERVICE_STATE_CHANGED, uid, packageName,
                     serviceName, FrameworkStatsLog.SERVICE_STATE_CHANGED__STATE__STOP);
             mAm.mBatteryStatsService.noteServiceStopRunning(uid, packageName, serviceName);
-            service.startRequested = false;
+            mAm.mProcessStateController.setStartRequested(service, false);
             if (service.tracker != null) {
                 synchronized (mAm.mProcessStats.mLock) {
                     service.tracker.setStarted(false, mAm.mProcessStats.getMemFactorLocked(),
@@ -1812,7 +1812,7 @@
             FrameworkStatsLog.write(FrameworkStatsLog.SERVICE_STATE_CHANGED, uid, packageName,
                     serviceName, FrameworkStatsLog.SERVICE_STATE_CHANGED__STATE__STOP);
             mAm.mBatteryStatsService.noteServiceStopRunning(uid, packageName, serviceName);
-            r.startRequested = false;
+            mAm.mProcessStateController.setStartRequested(r, false);
             if (r.tracker != null) {
                 synchronized (mAm.mProcessStats.mLock) {
                     r.tracker.setStarted(false, mAm.mProcessStats.getMemFactorLocked(),
@@ -2618,7 +2618,7 @@
                     }
                     notification.flags |= Notification.FLAG_FOREGROUND_SERVICE;
                     r.foregroundNoti = notification;
-                    r.foregroundServiceType = foregroundServiceType;
+                    mAm.mProcessStateController.setForegroundServiceType(r, foregroundServiceType);
                     if (!r.isForeground) {
                         final ServiceMap smap = getServiceMapLocked(r.userId);
                         if (smap != null) {
@@ -2643,7 +2643,7 @@
                             }
                             active.mNumActive++;
                         }
-                        r.isForeground = true;
+                        mAm.mProcessStateController.setIsForegroundService(r, true);
 
                         // The logging of FOREGROUND_SERVICE_STATE_CHANGED__STATE__ENTER event could
                         // be deferred, make a copy of mAllowStartForeground and
@@ -2772,7 +2772,7 @@
                     }
                 }
 
-                r.isForeground = false;
+                mAm.mProcessStateController.setIsForegroundService(r, false);
                 r.mFgsExitTime = SystemClock.uptimeMillis();
                 synchronized (mAm.mProcessStats.mLock) {
                     final ServiceState stracker = r.getTracker();
@@ -3565,7 +3565,7 @@
     private void maybeUpdateShortFgsTrackingLocked(ServiceRecord sr,
             boolean extendTimeout) {
         if (!sr.isShortFgs()) {
-            sr.clearShortFgsInfo(); // Just in case we have it.
+            mAm.mProcessStateController.clearShortFgsInfo(sr); // Just in case we have it.
             unscheduleShortFgsTimeoutLocked(sr);
             return;
         }
@@ -3581,7 +3581,7 @@
                 }
             }
             traceInstant("short FGS start/extend: ", sr);
-            sr.setShortFgsInfo(SystemClock.uptimeMillis());
+            mAm.mProcessStateController.setShortFgsInfo(sr, SystemClock.uptimeMillis());
 
             // We'll restart the timeout.
             unscheduleShortFgsTimeoutLocked(sr);
@@ -3605,7 +3605,7 @@
      * Stop the timeout for a ServiceRecord, if it's of a short-FGS.
      */
     private void maybeStopShortFgsTimeoutLocked(ServiceRecord sr) {
-        sr.clearShortFgsInfo(); // Always clear, just in case.
+        mAm.mProcessStateController.clearShortFgsInfo(sr); // Always clear, just in case.
         if (!sr.isShortFgs()) {
             return;
         }
@@ -3993,7 +3993,7 @@
     private void stopServiceAndUpdateAllowlistManagerLocked(ServiceRecord service) {
         maybeStopShortFgsTimeoutLocked(service);
         final ProcessServiceRecord psr = service.app.mServices;
-        psr.stopService(service);
+        mAm.mProcessStateController.stopService(psr, service);
         psr.updateBoundClientUids();
         if (service.allowlistManager) {
             updateAllowlistManagerLocked(psr);
@@ -4047,7 +4047,7 @@
             }
         }
         if (anyClientActivities != psr.hasClientActivities()) {
-            psr.setHasClientActivities(anyClientActivities);
+            mAm.mProcessStateController.setHasClientActivities(psr, anyClientActivities);
             if (updateLru) {
                 mAm.updateLruProcessLocked(psr.mApp, anyClientActivities, null);
             }
@@ -4216,7 +4216,8 @@
             }
 
             if ((flags&Context.BIND_AUTO_CREATE) != 0) {
-                s.lastActivity = SystemClock.uptimeMillis();
+                mAm.mProcessStateController.setServiceLastActivityTime(s,
+                        SystemClock.uptimeMillis());
                 if (!s.hasAutoCreateConnections()) {
                     // This is the first binding, let the tracker know.
                     synchronized (mAm.mProcessStats.mLock) {
@@ -4253,12 +4254,12 @@
             if (activity != null) {
                 activity.addConnection(c);
             }
-            clientPsr.addConnection(c);
+            mAm.mProcessStateController.addConnection(clientPsr, c);
             c.startAssociationIfNeeded();
             // Don't set hasAboveClient if binding to self to prevent modifyRawOomAdj() from
             // dropping the process' adjustment level.
             if (b.client != s.app && c.hasFlag(Context.BIND_ABOVE_CLIENT)) {
-                clientPsr.setHasAboveClient(true);
+                mAm.mProcessStateController.setHasAboveClient(clientPsr, true);
             }
             if (c.hasFlag(BIND_ALLOW_WHITELIST_MANAGEMENT)) {
                 s.allowlistManager = true;
@@ -4274,7 +4275,8 @@
             if (s.app != null && s.app.mState != null
                     && s.app.mState.getCurProcState() <= PROCESS_STATE_TOP
                     && c.hasFlag(Context.BIND_ALMOST_PERCEPTIBLE)) {
-                s.lastTopAlmostPerceptibleBindRequestUptimeMs = SystemClock.uptimeMillis();
+                mAm.mProcessStateController.setLastTopAlmostPerceptibleBindRequest(s,
+                        SystemClock.uptimeMillis());
             }
 
             if (s.app != null) {
@@ -4312,7 +4314,8 @@
 
             boolean needOomAdj = false;
             if (c.hasFlag(Context.BIND_AUTO_CREATE)) {
-                s.lastActivity = SystemClock.uptimeMillis();
+                mAm.mProcessStateController.setServiceLastActivityTime(s,
+                        SystemClock.uptimeMillis());
                 needOomAdj = (serviceBindingOomAdjPolicy
                         & SERVICE_BIND_OOMADJ_POLICY_SKIP_OOM_UPDATE_ON_CREATE) == 0;
                 if (bringUpServiceLocked(s, service.getFlags(), callerFg, false,
@@ -4328,7 +4331,7 @@
             if (s.app != null) {
                 ProcessServiceRecord servicePsr = s.app.mServices;
                 if (c.hasFlag(Context.BIND_TREAT_LIKE_ACTIVITY)) {
-                    servicePsr.setTreatLikeActivity(true);
+                    mAm.mProcessStateController.setTreatLikeActivity(servicePsr, true);
                 }
                 if (s.allowlistManager) {
                     servicePsr.mAllowlistManager = true;
@@ -4575,7 +4578,9 @@
                     }
                     // This could have made the service less important.
                     if (r.hasFlag(Context.BIND_TREAT_LIKE_ACTIVITY)) {
-                        psr.setTreatLikeActivity(true);
+                        // TODO(b/367545398): the following line is a bug. A service unbind
+                        //  should potentially lower a process's importance, not elevate it.
+                        mAm.mProcessStateController.setTreatLikeActivity(psr, true);
                         mAm.updateLruProcessLocked(app, true, null);
                     }
                     // If the bindee is more important than the binder, we may skip the OomAdjuster.
@@ -5077,7 +5082,10 @@
                         + " requires " + r.permission);
                 return new ServiceLookupResult(r.permission);
             } else if ((Manifest.permission.BIND_HOTWORD_DETECTION_SERVICE.equals(r.permission)
-                    || Manifest.permission.BIND_VISUAL_QUERY_DETECTION_SERVICE.equals(r.permission))
+                    || Manifest.permission.BIND_VISUAL_QUERY_DETECTION_SERVICE.equals(r.permission)
+                    || Manifest.permission.BIND_WEARABLE_SENSING_SERVICE.equals(r.permission)
+                    || Manifest.permission.BIND_ON_DEVICE_SANDBOXED_INFERENCE_SERVICE.equals(
+                    r.permission))
                     && callingUid != Process.SYSTEM_UID) {
                 // Hotword detection and visual query detection must run in its own sandbox, and we
                 // don't even trust its enclosing application to bind to it - only the system.
@@ -5162,8 +5170,9 @@
             }
             if (r.app != null) {
                 psr = r.app.mServices;
-                psr.startExecutingService(r);
-                psr.setExecServicesFg(psr.shouldExecServicesFg() || fg);
+                mAm.mProcessStateController.startExecutingService(psr, r);
+                mAm.mProcessStateController.setExecServicesFg(psr,
+                        psr.shouldExecServicesFg() || fg);
                 if (timeoutNeeded && psr.numberOfExecutingServices() == 1) {
                     if (!shouldSkipTimeout) {
                         scheduleServiceTimeoutLocked(r.app);
@@ -5175,7 +5184,7 @@
         } else if (r.app != null && fg) {
             psr = r.app.mServices;
             if (!psr.shouldExecServicesFg()) {
-                psr.setExecServicesFg(true);
+                mAm.mProcessStateController.setExecServicesFg(psr, true);
                 if (timeoutNeeded) {
                     if (!shouldSkipTimeout) {
                         scheduleServiceTimeoutLocked(r.app);
@@ -6020,11 +6029,13 @@
             Slog.v(TAG_MU, "realStartServiceLocked, ServiceRecord.uid = " + r.appInfo.uid
                     + ", ProcessRecord.uid = " + app.uid);
         r.setProcess(app, thread, pid, uidRecord);
-        r.restartTime = r.lastActivity = SystemClock.uptimeMillis();
+        final long now = SystemClock.uptimeMillis();
+        r.restartTime = now;
+        mAm.mProcessStateController.setServiceLastActivityTime(r, now);
         final boolean skipOomAdj = (serviceBindingOomAdjPolicy
                 & SERVICE_BIND_OOMADJ_POLICY_SKIP_OOM_UPDATE_ON_CREATE) != 0;
         final ProcessServiceRecord psr = app.mServices;
-        final boolean newService = psr.startService(r);
+        final boolean newService = mAm.mProcessStateController.startService(psr, r);
         bumpServiceExecutingLocked(r, execInFg, "create",
                 OOM_ADJ_REASON_NONE /* use "none" to avoid extra oom adj */,
                 skipOomAdj /* skipTimeoutIfPossible */);
@@ -6083,7 +6094,7 @@
 
                 // Cleanup.
                 if (newService) {
-                    psr.stopService(r);
+                    mAm.mProcessStateController.stopService(psr, r);
                     r.setProcess(null, null, 0, null);
                 }
 
@@ -6428,7 +6439,7 @@
             mAm.updateForegroundServiceUsageStats(r.name, r.userId, false);
         }
 
-        r.isForeground = false;
+        mAm.mProcessStateController.setIsForegroundService(r, false);
         r.mFgsNotificationWasDeferred = false;
         dropFgsNotificationStateLocked(r);
         r.foregroundId = 0;
@@ -6579,9 +6590,9 @@
         }
         if (b.client != skipApp) {
             final ProcessServiceRecord psr = b.client.mServices;
-            psr.removeConnection(c);
+            mAm.mProcessStateController.removeConnection(psr, c);
             if (c.hasFlag(Context.BIND_ABOVE_CLIENT)) {
-                psr.updateHasAboveClientLocked();
+                mAm.mProcessStateController.updateHasAboveClientLocked(psr);
             }
             // If this connection requested allowlist management, see if we should
             // now clear that state.
@@ -6597,7 +6608,7 @@
             }
             // And for almost perceptible exceptions.
             if (c.hasFlag(Context.BIND_ALMOST_PERCEPTIBLE)) {
-                psr.updateHasTopStartedAlmostPerceptibleServices();
+                mAm.mProcessStateController.updateHasTopStartedAlmostPerceptibleServices(psr);
             }
             if (s.app != null) {
                 updateServiceClientActivitiesLocked(s.app.mServices, c, true);
@@ -6796,8 +6807,8 @@
                 final ProcessServiceRecord psr = r.app.mServices;
                 if (DEBUG_SERVICE) Slog.v(TAG_SERVICE,
                         "Nesting at 0 of " + r.shortInstanceName);
-                psr.setExecServicesFg(false);
-                psr.stopExecutingService(r);
+                mAm.mProcessStateController.setExecServicesFg(psr, false);
+                mAm.mProcessStateController.stopExecutingService(psr, r);
                 if (psr.numberOfExecutingServices() == 0) {
                     if (DEBUG_SERVICE || DEBUG_SERVICE_EXECUTING) Slog.v(TAG_SERVICE_EXECUTING,
                             "No more executingServices of " + r.shortInstanceName);
@@ -6806,7 +6817,7 @@
                     // Need to re-evaluate whether the app still needs to be in the foreground.
                     for (int i = psr.numberOfExecutingServices() - 1; i >= 0; i--) {
                         if (psr.getExecutingServiceAt(i).executeFg) {
-                            psr.setExecServicesFg(true);
+                            mAm.mProcessStateController.setExecServicesFg(psr, true);
                             break;
                         }
                     }
@@ -6819,9 +6830,9 @@
                 }
                 if (oomAdjReason != OOM_ADJ_REASON_NONE) {
                     if (enqueueOomAdj) {
-                        mAm.enqueueOomAdjTargetLocked(r.app);
+                        mAm.mProcessStateController.enqueueUpdateTarget(r.app);
                     } else {
-                        mAm.updateOomAdjLocked(r.app, oomAdjReason);
+                        mAm.mProcessStateController.runUpdate(r.app, oomAdjReason);
                     }
                 } else {
                     // Skip oom adj if it wasn't bumped during the bumpServiceExecutingLocked()
@@ -7206,8 +7217,7 @@
             removeConnectionLocked(r, app, null, true);
         }
         updateServiceConnectionActivitiesLocked(psr);
-        psr.removeAllConnections();
-        psr.removeAllSdkSandboxConnections();
+        mAm.mProcessStateController.removeAllConnections(psr);
 
         psr.mAllowlistManager = false;
 
@@ -7217,7 +7227,7 @@
             mAm.mBatteryStatsService.noteServiceStopLaunch(sr.appInfo.uid, sr.name.getPackageName(),
                     sr.name.getClassName());
             if (sr.app != app && sr.app != null && !sr.app.isPersistent()) {
-                sr.app.mServices.stopService(sr);
+                mAm.mProcessStateController.stopService(psr, sr);
                 sr.app.mServices.updateBoundClientUids();
             }
             sr.setProcess(null, null, 0, null);
@@ -7287,7 +7297,7 @@
             // Unless the process is persistent, this process record is going away,
             // so make sure the service is cleaned out of it.
             if (!app.isPersistent()) {
-                psr.stopService(sr);
+                mAm.mProcessStateController.stopService(psr, sr);
                 psr.updateBoundClientUids();
             }
 
@@ -7328,7 +7338,7 @@
                     // Update to stopped state because the explicit start is gone. The service is
                     // scheduled to restart for other reason (e.g. connections) so we don't bring
                     // down it.
-                    sr.startRequested = false;
+                    mAm.mProcessStateController.setStartRequested(sr, false);
                     if (sr.tracker != null) {
                         synchronized (mAm.mProcessStats.mLock) {
                             sr.tracker.setStarted(false, mAm.mProcessStats.getMemFactorLocked(),
@@ -7342,7 +7352,7 @@
         mAm.updateOomAdjPendingTargetsLocked(OOM_ADJ_REASON_STOP_SERVICE);
 
         if (!allowRestart) {
-            psr.stopAllServices();
+            mAm.mProcessStateController.stopAllServices(psr);
             psr.clearBoundClientUids();
 
             // Make sure there are no more restarting services for this process.
@@ -7384,7 +7394,7 @@
             }
         }
 
-        psr.stopAllExecutingServices();
+        mAm.mProcessStateController.stopAllExecutingServices(psr);
         psr.noteScheduleServiceTimeoutPending(false);
     }
 
@@ -9210,14 +9220,14 @@
                 new ForegroundServiceDelegation(options, connection);
         r.mFgsDelegation = delegation;
         mFgsDelegations.put(delegation, r);
-        r.isForeground = true;
+        mAm.mProcessStateController.setIsForegroundService(r, true);
         r.mFgsEnterTime = SystemClock.uptimeMillis();
-        r.foregroundServiceType = options.mForegroundServiceTypes;
+        mAm.mProcessStateController.setForegroundServiceType(r, options.mForegroundServiceTypes);
         r.updateOomAdjSeq();
         setFgsRestrictionLocked(callingPackage, callingPid, callingUid, intent, r, userId,
                 BackgroundStartPrivileges.NONE,  false /* isBindService */);
         final ProcessServiceRecord psr = callerApp.mServices;
-        final boolean newService = psr.startService(r);
+        final boolean newService = mAm.mProcessStateController.startService(psr, r);
         // updateOomAdj.
         updateServiceForegroundLocked(psr, /* oomAdj= */ true);
 
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index 74437cd..bcca20b 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -626,6 +626,8 @@
             DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSSZ");
 
     OomAdjuster mOomAdjuster;
+    @GuardedBy("this")
+    ProcessStateController mProcessStateController;
 
     static final String EXTRA_TITLE = "android.intent.extra.TITLE";
     static final String EXTRA_DESCRIPTION = "android.intent.extra.DESCRIPTION";
@@ -1958,7 +1960,7 @@
                         new HostingRecord(HostingRecord.HOSTING_TYPE_SYSTEM));
                 app.setPersistent(true);
                 app.setPid(MY_PID);
-                app.mState.setMaxAdj(ProcessList.SYSTEM_ADJ);
+                mProcessStateController.setMaxAdj(app, ProcessList.SYSTEM_ADJ);
                 app.makeActive(new ApplicationThreadDeferred(mSystemThread.getApplicationThread()),
                         mProcessStats);
                 app.mProfile.addHostingComponentType(HOSTING_COMPONENT_TYPE_SYSTEM);
@@ -2394,9 +2396,11 @@
         mProcessList.init(this, activeUids, mPlatformCompat);
         mAppProfiler = new AppProfiler(this, BackgroundThread.getHandler().getLooper(), null);
         mPhantomProcessList = new PhantomProcessList(this);
-        mOomAdjuster = mConstants.ENABLE_NEW_OOMADJ
-                ? new OomAdjusterModernImpl(this, mProcessList, activeUids, handlerThread)
-                : new OomAdjuster(this, mProcessList, activeUids, handlerThread);
+        mProcessStateController = new ProcessStateController.Builder(this, mProcessList, activeUids)
+                .setHandlerThread(handlerThread)
+                .useModernOomAdjuster(mConstants.ENABLE_NEW_OOMADJ)
+                .build();
+        mOomAdjuster = mProcessStateController.getOomAdjuster();
 
         mIntentFirewall = injector.getIntentFirewall();
         mProcessStats = new ProcessStatsService(this, mContext.getCacheDir());
@@ -2459,9 +2463,10 @@
         mAppProfiler = new AppProfiler(this, BackgroundThread.getHandler().getLooper(),
                 new LowMemDetector(this));
         mPhantomProcessList = new PhantomProcessList(this);
-        mOomAdjuster = mConstants.ENABLE_NEW_OOMADJ
-                ? new OomAdjusterModernImpl(this, mProcessList, activeUids)
-                : new OomAdjuster(this, mProcessList, activeUids);
+        mProcessStateController = new ProcessStateController.Builder(this, mProcessList, activeUids)
+                .useModernOomAdjuster(mConstants.ENABLE_NEW_OOMADJ)
+                .build();
+        mOomAdjuster = mProcessStateController.getOomAdjuster();
 
         mBroadcastQueue = mInjector.getBroadcastQueue(this);
         mBroadcastController = new BroadcastController(mContext, this, mBroadcastQueue);
@@ -4574,7 +4579,7 @@
         EventLogTags.writeAmProcBound(app.userId, pid, app.processName);
 
         synchronized (mProcLock) {
-            mOomAdjuster.setAttachingProcessStatesLSP(app);
+            mProcessStateController.setAttachingProcessStatesLSP(app);
             clearProcessForegroundLocked(app);
             app.setDebugging(false);
             app.setKilledByAm(false);
@@ -4770,7 +4775,7 @@
                 app.makeActive(new ApplicationThreadDeferred(thread), mProcessStats);
                 checkTime(startTime, "attachApplicationLocked: immediately after bindApplication");
             }
-            app.setPendingFinishAttach(true);
+            mProcessStateController.setPendingFinishAttach(app, true);
 
             updateLruProcessLocked(app, false, null);
             checkTime(startTime, "attachApplicationLocked: after updateLruProcessLocked");
@@ -4854,7 +4859,7 @@
 
         synchronized (this) {
             // Mark the finish attach application phase as completed
-            app.setPendingFinishAttach(false);
+            mProcessStateController.setPendingFinishAttach(app, false);
 
             final boolean normalMode = mProcessesReady || isAllowedWhileBooting(app.info);
             final String processName = app.processName;
@@ -5009,7 +5014,7 @@
         // If another follow up update is needed, it will be scheduled by OomAdjuster.
         mHandler.removeMessages(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG);
         synchronized (this) {
-            mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+            mProcessStateController.runFollowUpUpdate();
         }
     }
 
@@ -5804,10 +5809,10 @@
                 if (pr == null) {
                     return;
                 }
-                pr.mState.setForcingToImportant(null);
+                mProcessStateController.setForcingToImportant(pr, null);
                 clearProcessForegroundLocked(pr);
             }
-            updateOomAdjLocked(pr, OOM_ADJ_REASON_UI_VISIBILITY);
+            mProcessStateController.runUpdate(pr, OOM_ADJ_REASON_UI_VISIBILITY);
         }
     }
 
@@ -5830,7 +5835,7 @@
                     oldToken.token.unlinkToDeath(oldToken, 0);
                     mImportantProcesses.remove(pid);
                     if (pr != null) {
-                        pr.mState.setForcingToImportant(null);
+                        mProcessStateController.setForcingToImportant(pr, null);
                     }
                     changed = true;
                 }
@@ -5844,7 +5849,7 @@
                     try {
                         token.linkToDeath(newToken, 0);
                         mImportantProcesses.put(pid, newToken);
-                        pr.mState.setForcingToImportant(newToken);
+                        mProcessStateController.setForcingToImportant(pr, newToken);
                         changed = true;
                     } catch (RemoteException e) {
                         // If the process died while doing this, we will later
@@ -5854,7 +5859,7 @@
             }
 
             if (changed) {
-                updateOomAdjLocked(pr, OOM_ADJ_REASON_UI_VISIBILITY);
+                mProcessStateController.runUpdate(pr, OOM_ADJ_REASON_UI_VISIBILITY);
             }
         }
     }
@@ -7274,7 +7279,7 @@
 
         if ((info.flags & PERSISTENT_MASK) == PERSISTENT_MASK) {
             app.setPersistent(true);
-            app.mState.setMaxAdj(ProcessList.PERSISTENT_PROC_ADJ);
+            mProcessStateController.setMaxAdj(app, ProcessList.PERSISTENT_PROC_ADJ);
         }
         if (app.getThread() == null && mPersistentStartingProcesses.indexOf(app) < 0) {
             mPersistentStartingProcesses.add(app);
@@ -7377,7 +7382,7 @@
                 mServices.updateScreenStateLocked(isAwake);
                 reportCurWakefulnessUsageEvent();
                 mActivityTaskManager.onScreenAwakeChanged(isAwake);
-                mOomAdjuster.onWakefulnessChanged(wakefulness);
+                mProcessStateController.setWakefulness(wakefulness);
 
                 updateOomAdjLocked(OOM_ADJ_REASON_UI_VISIBILITY);
             }
@@ -8346,16 +8351,10 @@
                         Slog.w(TAG, "setHasTopUi called on unknown pid: " + pid);
                         return;
                     }
-                    if (pr.mState.hasTopUi() != hasTopUi) {
-                        if (DEBUG_OOM_ADJ) {
-                            Slog.d(TAG, "Setting hasTopUi=" + hasTopUi + " for pid=" + pid);
-                        }
-                        pr.mState.setHasTopUi(hasTopUi);
-                        changed = true;
-                    }
+                    changed = mProcessStateController.setHasTopUi(pr, hasTopUi);
                 }
                 if (changed) {
-                    updateOomAdjLocked(pr, OOM_ADJ_REASON_UI_VISIBILITY);
+                    mProcessStateController.runUpdate(pr, OOM_ADJ_REASON_UI_VISIBILITY);
                 }
             }
         } finally {
@@ -14084,10 +14083,14 @@
                 proc.setInFullBackup(true);
             }
             r.app = proc;
+            // TODO(b/369300367): This code suggests there could be a previous backup being
+            //  replaced here, but an OomAdjsuter update is not triggered on the previous app
+            //  (whose state will change from being removed from mBackupTargets).
             final BackupRecord backupTarget = mBackupTargets.get(targetUserId);
             oldBackupUid = backupTarget != null ? backupTarget.appInfo.uid : -1;
             newBackupUid = proc.isInFullBackup() ? r.appInfo.uid : -1;
             mBackupTargets.put(targetUserId, r);
+            mProcessStateController.setBackupTarget(proc, targetUserId);
 
             proc.mProfile.addHostingComponentType(HOSTING_COMPONENT_TYPE_BACKUP);
 
@@ -14141,6 +14144,7 @@
                 }
                 mBackupTargets.removeAt(indexOfKey);
             }
+            mProcessStateController.stopBackupTarget(userId);
         }
 
         JobSchedulerInternal js = LocalServices.getService(JobSchedulerInternal.class);
@@ -14219,6 +14223,8 @@
 
                 // Not backing this app up any more; reset its OOM adjustment
                 final ProcessRecord proc = backupTarget.app;
+                // TODO(b/369300367): Triggering the update before the state is actually set
+                //  seems wrong.
                 updateOomAdjLocked(proc, OOM_ADJ_REASON_BACKUP);
                 proc.setInFullBackup(false);
                 proc.mProfile.clearHostingComponentType(HOSTING_COMPONENT_TYPE_BACKUP);
@@ -14237,6 +14243,7 @@
                 }
             } finally {
                 mBackupTargets.delete(userId);
+                mProcessStateController.stopBackupTarget(userId);
             }
         }
 
@@ -15309,7 +15316,8 @@
                             proc.info.packageName, proc.info.uid, proc.getPid(), isForeground);
                 }
             }
-            psr.setHasForegroundServices(isForeground, fgServiceTypes, hasTypeNoneFgs);
+            mProcessStateController.setHasForegroundServices(psr, isForeground, fgServiceTypes,
+                    hasTypeNoneFgs);
             ArrayList<ProcessRecord> curProcs = mForegroundPackages.get(proc.info.packageName,
                     proc.info.uid);
             if (isForeground) {
@@ -15340,7 +15348,7 @@
                     ProcessChangeItem.CHANGE_FOREGROUND_SERVICES, fgServiceTypes);
         }
         if (oomAdj) {
-            updateOomAdjLocked(proc, OOM_ADJ_REASON_UI_VISIBILITY);
+            mProcessStateController.runUpdate(proc, OOM_ADJ_REASON_UI_VISIBILITY);
         }
     }
 
@@ -15390,7 +15398,7 @@
      */
     @GuardedBy("this")
     void enqueueOomAdjTargetLocked(ProcessRecord app) {
-        mOomAdjuster.enqueueOomAdjTargetLocked(app);
+        mProcessStateController.enqueueUpdateTarget(app);
     }
 
     /**
@@ -15398,7 +15406,7 @@
      */
     @GuardedBy("this")
     void removeOomAdjTargetLocked(ProcessRecord app, boolean procDied) {
-        mOomAdjuster.removeOomAdjTargetLocked(app, procDied);
+        mProcessStateController.removeUpdateTarget(app, procDied);
     }
 
     /**
@@ -15407,7 +15415,7 @@
      */
     @GuardedBy("this")
     void updateOomAdjPendingTargetsLocked(@OomAdjReason int oomAdjReason) {
-        mOomAdjuster.updateOomAdjPendingTargetsLocked(oomAdjReason);
+        mProcessStateController.runPendingUpdate(oomAdjReason);
     }
 
     static final class ProcStatsRunnable implements Runnable {
@@ -15426,7 +15434,7 @@
 
     @GuardedBy("this")
     final void updateOomAdjLocked(@OomAdjReason int oomAdjReason) {
-        mOomAdjuster.updateOomAdjLocked(oomAdjReason);
+        mProcessStateController.runFullUpdate(oomAdjReason);
     }
 
     /**
@@ -15438,7 +15446,7 @@
      */
     @GuardedBy("this")
     final boolean updateOomAdjLocked(ProcessRecord app, @OomAdjReason int oomAdjReason) {
-        return mOomAdjuster.updateOomAdjLocked(app, oomAdjReason);
+        return mProcessStateController.runUpdate(app, oomAdjReason);
     }
 
     @Override
@@ -15703,7 +15711,7 @@
 
     @GuardedBy({"this", "mProcLock"})
     final void setUidTempAllowlistStateLSP(int uid, boolean onAllowlist) {
-        mOomAdjuster.setUidTempAllowlistStateLSP(uid, onAllowlist);
+        mProcessStateController.setUidTempAllowlistStateLSP(uid, onAllowlist);
     }
 
     private void trimApplications(boolean forceFullOomAdj, @OomAdjReason int oomAdjReason) {
@@ -16753,12 +16761,9 @@
                         return;
                     }
                 }
-                if (pr.mState.hasOverlayUi() == hasOverlayUi) {
-                    return;
+                if (mProcessStateController.setHasOverlayUi(pr, hasOverlayUi)) {
+                    mProcessStateController.runUpdate(pr, OOM_ADJ_REASON_UI_VISIBILITY);
                 }
-                pr.mState.setHasOverlayUi(hasOverlayUi);
-                //Slog.i(TAG, "Setting hasOverlayUi=" + pr.hasOverlayUi + " for pid=" + pid);
-                updateOomAdjLocked(pr, OOM_ADJ_REASON_UI_VISIBILITY);
             }
         }
 
diff --git a/services/core/java/com/android/server/am/AppProfiler.java b/services/core/java/com/android/server/am/AppProfiler.java
index dda48ad..4f2d69e 100644
--- a/services/core/java/com/android/server/am/AppProfiler.java
+++ b/services/core/java/com/android/server/am/AppProfiler.java
@@ -1356,6 +1356,7 @@
     @GuardedBy("mService")
     void setMemFactorOverrideLocked(@MemFactor int factor) {
         mMemFactorOverride = factor;
+        mService.mProcessStateController.setIsLastMemoryLevelNormal(isLastMemoryLevelNormal());
     }
 
     @GuardedBy({"mService", "mProcLock"})
@@ -1423,6 +1424,7 @@
         }
 
         mLastMemoryLevel = memFactor;
+        mService.mProcessStateController.setIsLastMemoryLevelNormal(isLastMemoryLevelNormal());
         mLastNumProcesses = mService.mProcessList.getLruSizeLOSP();
 
         // Dispatch UI_HIDDEN to processes that need it
diff --git a/services/core/java/com/android/server/am/ContentProviderHelper.java b/services/core/java/com/android/server/am/ContentProviderHelper.java
index da40826..221938a 100644
--- a/services/core/java/com/android/server/am/ContentProviderHelper.java
+++ b/services/core/java/com/android/server/am/ContentProviderHelper.java
@@ -325,7 +325,8 @@
                     final int verifiedAdj = cpr.proc.mState.getVerifiedAdj();
                     boolean success = !serviceBindingOomAdjPolicy()
                             || mService.mOomAdjuster.evaluateProviderConnectionAdd(r, cpr.proc)
-                            ? mService.updateOomAdjLocked(cpr.proc, OOM_ADJ_REASON_GET_PROVIDER)
+                            ? mService.mProcessStateController.runUpdate(cpr.proc,
+                            OOM_ADJ_REASON_GET_PROVIDER)
                             : true;
                     // XXX things have changed so updateOomAdjLocked doesn't actually tell us
                     // if the process has been successfully adjusted.  So to reduce races with
@@ -534,10 +535,9 @@
                             if (ActivityManagerDebugConfig.DEBUG_PROVIDER) {
                                 Slog.d(TAG, "Installing in existing process " + proc);
                             }
-                            final ProcessProviderRecord pr = proc.mProviders;
-                            if (!pr.hasProvider(cpi.name)) {
+                            if (mService.mProcessStateController.addPublishedProvider(proc,
+                                    cpi.name, cpr)) {
                                 checkTime(startTime, "getContentProviderImpl: scheduling install");
-                                pr.installProvider(cpi.name, cpr);
                                 mService.mOomAdjuster.unfreezeTemporarily(proc,
                                         CachedAppOptimizer.UNFREEZE_REASON_GET_PROVIDER);
                                 try {
@@ -881,7 +881,8 @@
             ComponentName comp = new ComponentName(cpr.info.packageName, cpr.info.name);
             ContentProviderRecord localCpr = mProviderMap.getProviderByClass(comp, userId);
             if (localCpr.hasExternalProcessHandles()) {
-                if (localCpr.removeExternalProcessHandleLocked(token)) {
+                if (mService.mProcessStateController.removeExternalProviderClient(localCpr,
+                        token)) {
                     mService.updateOomAdjLocked(localCpr.proc, OOM_ADJ_REASON_REMOVE_PROVIDER);
                 } else {
                     Slog.e(TAG, "Attempt to remove content provider " + localCpr
@@ -1447,7 +1448,8 @@
             String callingPackage, String callingTag, boolean stable, boolean updateLru,
             long startTime, ProcessList processList, @UserIdInt int expectedUserId) {
         if (r == null) {
-            cpr.addExternalProcessHandleLocked(externalProcessToken, callingUid, callingTag);
+            mService.mProcessStateController.addExternalProviderClient(cpr, externalProcessToken,
+                    callingUid, callingTag);
             return null;
         }
 
@@ -1470,7 +1472,7 @@
         if (cpr.proc != null) {
             cpr.proc.mProfile.addHostingComponentType(HOSTING_COMPONENT_TYPE_PROVIDER);
         }
-        pr.addProviderConnection(conn);
+        mService.mProcessStateController.addProviderConnection(r, conn);
         mService.startAssociationLocked(r.uid, r.processName, r.mState.getCurProcState(),
                 cpr.uid, cpr.appInfo.longVersionCode, cpr.name, cpr.info.processName);
         if (updateLru && cpr.proc != null
@@ -1493,7 +1495,8 @@
             ContentProviderRecord cpr, IBinder externalProcessToken, boolean stable,
             boolean enforceDelay, boolean updateOomAdj) {
         if (conn == null) {
-            cpr.removeExternalProcessHandleLocked(externalProcessToken);
+            mService.mProcessStateController.removeExternalProviderClient(cpr,
+                    externalProcessToken);
             return false;
         }
 
@@ -1537,14 +1540,15 @@
             if (cpr.proc != null && !hasProviderConnectionLocked(cpr.proc)) {
                 cpr.proc.mProfile.clearHostingComponentType(HOSTING_COMPONENT_TYPE_PROVIDER);
             }
-            conn.client.mProviders.removeProviderConnection(conn);
+            mService.mProcessStateController.removeProviderConnection(conn.client, conn);
             if (conn.client.mState.getSetProcState()
                     < ActivityManager.PROCESS_STATE_LAST_ACTIVITY) {
                 // The client is more important than last activity -- note the time this
                 // is happening, so we keep the old provider process around a bit as last
                 // activity to avoid thrashing it.
                 if (cpr.proc != null) {
-                    cpr.proc.mProviders.setLastProviderTime(SystemClock.uptimeMillis());
+                    mService.mProcessStateController.setLastProviderTime(cpr.proc,
+                            SystemClock.uptimeMillis());
                 }
             }
             mService.stopAssociationLocked(conn.client.uid, conn.client.processName, cpr.uid,
@@ -1821,7 +1825,7 @@
                 }
             }
             if (removed && cpr.proc != null) {
-                cpr.proc.mProviders.removeProvider(cpr.info.name);
+                mService.mProcessStateController.removePublishedProvider(cpr.proc, cpr.info.name);
             }
         }
 
diff --git a/services/core/java/com/android/server/am/OWNERS b/services/core/java/com/android/server/am/OWNERS
index 2a30ad0..61079fc 100644
--- a/services/core/java/com/android/server/am/OWNERS
+++ b/services/core/java/com/android/server/am/OWNERS
@@ -54,8 +54,9 @@
 per-file *Permission* = patb@google.com
 per-file *Package* = patb@google.com
 
-# OOM Adjuster
+# OOM Adjuster & ProcessStateController
 per-file *Oom* = file:/OOM_ADJUSTER_OWNERS
+per-file ProcessStateController.java = file:/OOM_ADJUSTER_OWNERS
 
 # Miscellaneous
 per-file SettingsToPropertiesMapper.java = omakoto@google.com, yamasani@google.com, dzshen@google.com, zhidou@google.com, tedbauer@google.com
diff --git a/services/core/java/com/android/server/am/OomAdjuster.java b/services/core/java/com/android/server/am/OomAdjuster.java
index 4073ab8..e8f7b5f 100644
--- a/services/core/java/com/android/server/am/OomAdjuster.java
+++ b/services/core/java/com/android/server/am/OomAdjuster.java
@@ -130,6 +130,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.UserIdInt;
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal.OomAdjReason;
 import android.app.ActivityThread;
@@ -376,6 +377,7 @@
 
     final ActivityManagerService mService;
     final Injector mInjector;
+    final GlobalState mGlobalState;
     final ProcessList mProcessList;
     final ActivityManagerGlobalLock mProcLock;
 
@@ -470,15 +472,23 @@
 
     }
 
+    // TODO(b/346822474): hook up global state usage.
+    interface GlobalState {
+        /** Is device's screen on. */
+        boolean isAwake();
+
+        /** What process is running a backup for a given userId. */
+        ProcessRecord getBackupTarget(@UserIdInt int userId);
+
+        /** Is memory level normal since last evaluation. */
+        boolean isLastMemoryLevelNormal();
+    }
+
     boolean isChangeEnabled(@CachedCompatChangeId int cachedCompatChangeId,
             ApplicationInfo app, boolean defaultValue) {
         return mInjector.isChangeEnabled(cachedCompatChangeId, app, defaultValue);
     }
 
-    OomAdjuster(ActivityManagerService service, ProcessList processList, ActiveUids activeUids) {
-        this(service, processList, activeUids, createAdjusterThread());
-    }
-
     static ServiceThread createAdjusterThread() {
         // The process group is usually critical to the response time of foreground app, so the
         // setter should apply it as soon as possible.
@@ -489,18 +499,9 @@
     }
 
     OomAdjuster(ActivityManagerService service, ProcessList processList, ActiveUids activeUids,
-            ServiceThread adjusterThread) {
-        this(service, processList, activeUids, adjusterThread, new Injector());
-    }
-
-    OomAdjuster(ActivityManagerService service, ProcessList processList, ActiveUids activeUids,
-            Injector injector) {
-        this(service, processList, activeUids, createAdjusterThread(), injector);
-    }
-
-    OomAdjuster(ActivityManagerService service, ProcessList processList, ActiveUids activeUids,
-            ServiceThread adjusterThread, Injector injector) {
+            ServiceThread adjusterThread, GlobalState globalState, Injector injector) {
         mService = service;
+        mGlobalState = globalState;
         mInjector = injector;
         mProcessList = processList;
         mProcLock = service.mProcLock;
@@ -1816,9 +1817,36 @@
         }
     }
 
+    private boolean isDeviceFullyAwake() {
+        if (Flags.pushGlobalStateToOomadjuster()) {
+            return mGlobalState.isAwake();
+        } else {
+            return mService.mWakefulness.get() == PowerManagerInternal.WAKEFULNESS_AWAKE;
+        }
+    }
+
     private boolean isScreenOnOrAnimatingLocked(ProcessStateRecord state) {
-        return mService.mWakefulness.get() == PowerManagerInternal.WAKEFULNESS_AWAKE
-                || state.isRunningRemoteAnimation();
+        return isDeviceFullyAwake() || state.isRunningRemoteAnimation();
+    }
+
+    private boolean isBackupProcess(ProcessRecord app) {
+        if (Flags.pushGlobalStateToOomadjuster()) {
+            return app == mGlobalState.getBackupTarget(app.userId);
+        } else {
+            final BackupRecord backupTarget = mService.mBackupTargets.get(app.userId);
+            if (backupTarget == null) {
+                return false;
+            }
+            return app == backupTarget.app;
+        }
+    }
+
+    private boolean isLastMemoryLevelNormal() {
+        if (Flags.pushGlobalStateToOomadjuster()) {
+            return mGlobalState.isLastMemoryLevelNormal();
+        } else {
+            return mService.mAppProfiler.isLastMemoryLevelNormal();
+        }
     }
 
     @GuardedBy({"mService", "mProcLock"})
@@ -2259,8 +2287,7 @@
         state.setHasStartedServices(false);
         state.setAdjSeq(mAdjSeq);
 
-        final BackupRecord backupTarget = mService.mBackupTargets.get(app.userId);
-        if (backupTarget != null && app == backupTarget.app) {
+        if (isBackupProcess(app)) {
             // If possible we want to avoid killing apps while they're being backed up
             if (adj > BACKUP_APP_ADJ) {
                 if (DEBUG_BACKUP) Slog.v(TAG_BACKUP, "oom BACKUP_APP_ADJ for " + app);
@@ -2526,8 +2553,7 @@
                     double cachedRestoreThreshold =
                             mProcessList.getCachedRestoreThresholdKb() * thresholdModifier;
 
-                    if (!mService.mAppProfiler.isLastMemoryLevelNormal()
-                            && lastPssOrRss >= cachedRestoreThreshold) {
+                    if (isLastMemoryLevelNormal() && lastPssOrRss >= cachedRestoreThreshold) {
                         state.setServiceHighRam(true);
                         state.setServiceB(true);
                         //Slog.i(TAG, "ADJ " + app + " high ram!");
@@ -2621,7 +2647,7 @@
         // Put bound foreground services in a special sched group for additional
         // restrictions on screen off
         if (state.getCurProcState() >= PROCESS_STATE_BOUND_FOREGROUND_SERVICE
-                && mService.mWakefulness.get() != PowerManagerInternal.WAKEFULNESS_AWAKE
+                && !isDeviceFullyAwake()
                 && !state.shouldScheduleLikeTopApp()) {
             if (schedGroup > SCHED_GROUP_RESTRICTED) {
                 schedGroup = SCHED_GROUP_RESTRICTED;
@@ -2910,8 +2936,7 @@
                         clientProcState = PROCESS_STATE_FOREGROUND_SERVICE;
                     } else if (cr.hasFlag(Context.BIND_FOREGROUND_SERVICE)) {
                         clientProcState = PROCESS_STATE_BOUND_FOREGROUND_SERVICE;
-                    } else if (mService.mWakefulness.get()
-                            == PowerManagerInternal.WAKEFULNESS_AWAKE
+                    } else if (isDeviceFullyAwake()
                             && cr.hasFlag(Context.BIND_FOREGROUND_SERVICE_WHILE_AWAKE)) {
                         clientProcState = PROCESS_STATE_BOUND_FOREGROUND_SERVICE;
                     } else {
diff --git a/services/core/java/com/android/server/am/OomAdjusterModernImpl.java b/services/core/java/com/android/server/am/OomAdjusterModernImpl.java
index fb1c2e9..e452c45 100644
--- a/services/core/java/com/android/server/am/OomAdjusterModernImpl.java
+++ b/services/core/java/com/android/server/am/OomAdjusterModernImpl.java
@@ -756,18 +756,9 @@
             new ComputeConnectionsConsumer();
 
     OomAdjusterModernImpl(ActivityManagerService service, ProcessList processList,
-            ActiveUids activeUids) {
-        this(service, processList, activeUids, createAdjusterThread());
-    }
-
-    OomAdjusterModernImpl(ActivityManagerService service, ProcessList processList,
-            ActiveUids activeUids, ServiceThread adjusterThread) {
-        super(service, processList, activeUids, adjusterThread);
-    }
-
-    OomAdjusterModernImpl(ActivityManagerService service, ProcessList processList,
-            ActiveUids activeUids, Injector injector) {
-        super(service, processList, activeUids, injector);
+            ActiveUids activeUids, ServiceThread adjusterThread, GlobalState globalState,
+            Injector injector) {
+        super(service, processList, activeUids, adjusterThread, globalState, injector);
     }
 
     private final ProcessRecordNodes mProcessRecordProcStateNodes = new ProcessRecordNodes(
diff --git a/services/core/java/com/android/server/am/ProcessList.java b/services/core/java/com/android/server/am/ProcessList.java
index 57922d5..2485626 100644
--- a/services/core/java/com/android/server/am/ProcessList.java
+++ b/services/core/java/com/android/server/am/ProcessList.java
@@ -3439,12 +3439,12 @@
             state.setCurrentSchedulingGroup(ProcessList.SCHED_GROUP_DEFAULT);
             state.setSetSchedGroup(ProcessList.SCHED_GROUP_DEFAULT);
             r.setPersistent(true);
-            state.setMaxAdj(ProcessList.PERSISTENT_PROC_ADJ);
+            mService.mProcessStateController.setMaxAdj(r, ProcessList.PERSISTENT_PROC_ADJ);
         }
         if (isolated && isolatedUid != 0) {
             // Special case for startIsolatedProcess (internal only) - assume the process
             // is required by the system server to prevent it being killed.
-            state.setMaxAdj(ProcessList.PERSISTENT_SERVICE_ADJ);
+            mService.mProcessStateController.setMaxAdj(r, ProcessList.PERSISTENT_SERVICE_ADJ);
         }
         addProcessNameLocked(r);
         return r;
diff --git a/services/core/java/com/android/server/am/ProcessRecord.java b/services/core/java/com/android/server/am/ProcessRecord.java
index 3e71d00..b51db13 100644
--- a/services/core/java/com/android/server/am/ProcessRecord.java
+++ b/services/core/java/com/android/server/am/ProcessRecord.java
@@ -17,6 +17,7 @@
 package com.android.server.am;
 
 import static android.app.ActivityManagerInternal.OOM_ADJ_REASON_ACTIVITY;
+import static android.app.ActivityManagerInternal.OOM_ADJ_REASON_UI_VISIBILITY;
 
 import static com.android.internal.util.Preconditions.checkArgument;
 import static com.android.server.am.ActivityManagerDebugConfig.TAG_AM;
@@ -1192,7 +1193,7 @@
         setWaitingToKill(null);
 
         mState.onCleanupApplicationRecordLSP();
-        mServices.onCleanupApplicationRecordLocked();
+        mService.mProcessStateController.onCleanupApplicationRecord(mServices);
         mReceivers.onCleanupApplicationRecordLocked();
         mService.mOomAdjuster.onProcessEndLocked(this);
 
@@ -1638,7 +1639,7 @@
             updateProcessInfo(false /* updateServiceConnectionActivities */,
                     true /* activityChange */, true /* updateOomAdj */);
             setPendingUiClean(true);
-            mState.setHasShownUi(true);
+            mService.mProcessStateController.setHasShownUi(this, true);
             mState.forceProcessStateUpTo(topProcessState);
         }
     }
@@ -1657,7 +1658,10 @@
             return;
         }
         synchronized (mService) {
-            mState.setRunningRemoteAnimation(runningRemoteAnimation);
+            if (mService.mProcessStateController.setRunningRemoteAnimation(this,
+                    runningRemoteAnimation)) {
+                mService.mProcessStateController.runUpdate(this, OOM_ADJ_REASON_UI_VISIBILITY);
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/am/ProcessStateController.java b/services/core/java/com/android/server/am/ProcessStateController.java
new file mode 100644
index 0000000..428df23
--- /dev/null
+++ b/services/core/java/com/android/server/am/ProcessStateController.java
@@ -0,0 +1,645 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.am;
+
+import static com.android.server.am.ActivityManagerDebugConfig.DEBUG_OOM_ADJ;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.UserIdInt;
+import android.app.ActivityManager;
+import android.app.ActivityManagerInternal;
+import android.content.pm.ServiceInfo;
+import android.os.IBinder;
+import android.os.PowerManagerInternal;
+import android.util.Slog;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.server.ServiceThread;
+
+/**
+ * ProcessStateController is responsible for maintaining state that can affect the OomAdjuster
+ * computations of a process. Any state that can affect a process's importance must be set by
+ * only ProcessStateController.
+ */
+public class ProcessStateController {
+    public static String TAG = "ProcessStateController";
+
+    private final OomAdjuster mOomAdjuster;
+
+    private final GlobalState mGlobalState = new GlobalState();
+
+    private ProcessStateController(ActivityManagerService ams, ProcessList processList,
+            ActiveUids activeUids, ServiceThread handlerThread, OomAdjuster.Injector oomAdjInjector,
+            boolean useOomAdjusterModernImpl) {
+        mOomAdjuster = useOomAdjusterModernImpl
+                ? new OomAdjusterModernImpl(ams, processList, activeUids, handlerThread,
+                mGlobalState, oomAdjInjector)
+                : new OomAdjuster(ams, processList, activeUids, handlerThread, mGlobalState,
+                        oomAdjInjector);
+    }
+
+    /**
+     * Get the instance of OomAdjuster that ProcessStateController is using.
+     * Must only be interacted with while holding the ActivityManagerService lock.
+     */
+    public OomAdjuster getOomAdjuster() {
+        return mOomAdjuster;
+    }
+
+    /**
+     * Add a process to evaluated the next time an update is run.
+     */
+    public void enqueueUpdateTarget(@NonNull ProcessRecord proc) {
+        mOomAdjuster.enqueueOomAdjTargetLocked(proc);
+    }
+
+    /**
+     * Remove a process that was added by {@link #enqueueUpdateTarget}.
+     */
+    public void removeUpdateTarget(@NonNull ProcessRecord proc, boolean procDied) {
+        mOomAdjuster.removeOomAdjTargetLocked(proc, procDied);
+    }
+
+    /**
+     * Trigger an update on a single process (and any processes that have been enqueued with
+     * {@link #enqueueUpdateTarget}).
+     */
+    public boolean runUpdate(@NonNull ProcessRecord proc,
+            @ActivityManagerInternal.OomAdjReason int oomAdjReason) {
+        return mOomAdjuster.updateOomAdjLocked(proc, oomAdjReason);
+    }
+
+    /**
+     * Trigger an update on all processes that have been enqueued with {@link #enqueueUpdateTarget}.
+     */
+    public void runPendingUpdate(@ActivityManagerInternal.OomAdjReason int oomAdjReason) {
+        mOomAdjuster.updateOomAdjPendingTargetsLocked(oomAdjReason);
+    }
+
+    /**
+     * Trigger an update on all processes.
+     */
+    public void runFullUpdate(@ActivityManagerInternal.OomAdjReason int oomAdjReason) {
+        mOomAdjuster.updateOomAdjLocked(oomAdjReason);
+    }
+
+    /**
+     * Trigger an update on any processes that have been marked for follow up during a previous
+     * update.
+     */
+    public void runFollowUpUpdate() {
+        mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+    }
+
+    private static class GlobalState implements OomAdjuster.GlobalState {
+        public boolean isAwake = true;
+        // TODO(b/369300367): Maintaining global state for backup processes is a bit convoluted.
+        //  ideally the state gets migrated to ProcessStateRecord.
+        public final SparseArray<ProcessRecord> backupTargets = new SparseArray<>();
+        public boolean isLastMemoryLevelNormal = true;
+
+        public boolean isAwake() {
+            return isAwake;
+        }
+
+        public ProcessRecord getBackupTarget(@UserIdInt int userId) {
+            return backupTargets.get(userId);
+        }
+
+        public boolean isLastMemoryLevelNormal() {
+            return isLastMemoryLevelNormal;
+        }
+    }
+
+    /*************************** Global State Events ***************************/
+    /**
+     * Set which process state Top processes should get.
+     */
+    public void setTopProcessState(@ActivityManager.ProcessState int procState) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set whether to give Top processes the Top sched group.
+     */
+    public void setUseTopSchedGroupForTopProcess(boolean useTopSchedGroup) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set the Top process.
+     */
+    public void setTopApp(@Nullable ProcessRecord proc) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set which process is considered the Home process, if any.
+     */
+    public void setHomeProcess(@Nullable ProcessRecord proc) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set which process is considered the Heavy Weight process, if any.
+     */
+    public void setHeavyWeightProcess(@Nullable ProcessRecord proc) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set which process is showing UI while the screen is off, if any.
+     */
+    public void setVisibleDozeUiProcess(@Nullable ProcessRecord proc) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set which process is considered the Previous process, if any.
+     */
+    public void setPreviousProcess(@Nullable ProcessRecord proc) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set what wakefulness state the screen is in.
+     */
+    public void setWakefulness(int wakefulness) {
+        mGlobalState.isAwake = (wakefulness == PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mOomAdjuster.onWakefulnessChanged(wakefulness);
+    }
+
+    /**
+     * Set for a given user what process is currently running a backup, if any.
+     */
+    public void setBackupTarget(@NonNull ProcessRecord proc, @UserIdInt int userId) {
+        mGlobalState.backupTargets.put(userId, proc);
+    }
+
+    /**
+     * No longer consider any process running a backup for a given user.
+     */
+    public void stopBackupTarget(@UserIdInt int userId) {
+        mGlobalState.backupTargets.delete(userId);
+    }
+
+    /**
+     * Set whether the last known memory level is normal.
+     */
+    public void setIsLastMemoryLevelNormal(boolean isMemoryNormal) {
+        mGlobalState.isLastMemoryLevelNormal = isMemoryNormal;
+    }
+
+    /***************************** UID State Events ****************************/
+    /**
+     * Set a UID as temp allowlisted.
+     */
+    public void setUidTempAllowlistStateLSP(int uid, boolean allowList) {
+        mOomAdjuster.setUidTempAllowlistStateLSP(uid, allowList);
+    }
+
+    /*********************** Process Miscellaneous Events **********************/
+    /**
+     * Set the maximum adj score a process can be assigned.
+     */
+    public void setMaxAdj(@NonNull ProcessRecord proc, int adj) {
+        proc.mState.setMaxAdj(adj);
+    }
+
+    /**
+     * Initialize a process that is being attached.
+     */
+    @GuardedBy({"mService", "mProcLock"})
+    public void setAttachingProcessStatesLSP(@NonNull ProcessRecord proc) {
+        mOomAdjuster.setAttachingProcessStatesLSP(proc);
+    }
+
+    /**
+     * Note whether a process is pending attach or not.
+     */
+    public void setPendingFinishAttach(@NonNull ProcessRecord proc, boolean pendingFinishAttach) {
+        proc.setPendingFinishAttach(pendingFinishAttach);
+    }
+
+    /**
+     * Set what sched group to grant a process due to running a broadcast.
+     * {@link ProcessList.SCHED_GROUP_UNDEFINED} means the process is not running a broadcast.
+     */
+    public void setBroadcastSchedGroup(@NonNull ProcessRecord proc, int schedGroup) {
+        // TODO(b/302575389): Migrate state pulled from BroadcastQueue to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /********************* Process Visibility State Events *********************/
+    /**
+     * Note whether a process has Top UI or not.
+     *
+     * @return true if the state changed, otherwise returns false.
+     */
+    public boolean setHasTopUi(@NonNull ProcessRecord proc, boolean hasTopUi) {
+        if (proc.mState.hasTopUi() == hasTopUi) return false;
+        if (DEBUG_OOM_ADJ) {
+            Slog.d(TAG, "Setting hasTopUi=" + hasTopUi + " for pid=" + proc.getPid());
+        }
+        proc.mState.setHasTopUi(hasTopUi);
+        return true;
+    }
+
+    /**
+     * Note whether a process is displaying Overlay UI or not.
+     *
+     * @return true if the state changed, otherwise returns false.
+     */
+    public boolean setHasOverlayUi(@NonNull ProcessRecord proc, boolean hasOverlayUi) {
+        if (proc.mState.hasOverlayUi() == hasOverlayUi) return false;
+        proc.mState.setHasOverlayUi(hasOverlayUi);
+        return true;
+    }
+
+
+    /**
+     * Note whether a process is running a remote animation.
+     *
+     * @return true if the state changed, otherwise returns false.
+     */
+    public boolean setRunningRemoteAnimation(@NonNull ProcessRecord proc,
+            boolean runningRemoteAnimation) {
+        if (proc.mState.isRunningRemoteAnimation() == runningRemoteAnimation) return false;
+        if (DEBUG_OOM_ADJ) {
+            Slog.i(TAG, "Setting runningRemoteAnimation=" + runningRemoteAnimation
+                    + " for pid=" + proc.getPid());
+        }
+        proc.mState.setRunningRemoteAnimation(runningRemoteAnimation);
+        return true;
+    }
+
+    /**
+     * Note that the process is showing a toast.
+     */
+    public void setForcingToImportant(@NonNull ProcessRecord proc,
+            @Nullable Object forcingToImportant) {
+        if (proc.mState.getForcingToImportant() == forcingToImportant) return;
+        proc.mState.setForcingToImportant(forcingToImportant);
+    }
+
+    /**
+     * Note that the process has shown UI at some point in its life.
+     */
+    public void setHasShownUi(@NonNull ProcessRecord proc, boolean hasShownUi) {
+        // This arguably should be turned into an internal state of OomAdjuster.
+        if (proc.mState.hasShownUi() == hasShownUi) return;
+        proc.mState.setHasShownUi(hasShownUi);
+    }
+
+    /**
+     * Note whether the process has an activity or not.
+     */
+    public void setHasActivity(@NonNull ProcessRecord proc, boolean hasActivity) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        // Possibly not needed, maybe can use ActivityStateFlags.
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Note whether the process has a visibly activity or not.
+     */
+    public void setHasVisibleActivity(@NonNull ProcessRecord proc, boolean hasVisibleActivity) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        // maybe used ActivityStateFlags instead.
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /**
+     * Set the Activity State Flags for a process.
+     */
+    public void setActivityStateFlags(@NonNull ProcessRecord proc, int flags) {
+        // TODO(b/302575389): Migrate state pulled from ATMS to a pushed model
+        throw new UnsupportedOperationException("Not implemented yet");
+    }
+
+    /********************** Content Provider State Events **********************/
+    /**
+     * Note that a process is hosting a content provider.
+     */
+    public boolean addPublishedProvider(@NonNull ProcessRecord proc, String name,
+            ContentProviderRecord cpr) {
+        final ProcessProviderRecord providers = proc.mProviders;
+        if (providers.hasProvider(name)) return false;
+        providers.installProvider(name, cpr);
+        return true;
+    }
+
+    /**
+     * Remove a published content provider from a process.
+     */
+    public void removePublishedProvider(@NonNull ProcessRecord proc, String name) {
+        final ProcessProviderRecord providers = proc.mProviders;
+        providers.removeProvider(name);
+    }
+
+    /**
+     * Note that a content provider has an external client.
+     */
+    public void addExternalProviderClient(@NonNull ContentProviderRecord cpr,
+            IBinder externalProcessToken, int callingUid, String callingTag) {
+        cpr.addExternalProcessHandleLocked(externalProcessToken, callingUid, callingTag);
+    }
+
+    /**
+     * Remove an external client from a conetnt provider.
+     */
+    public boolean removeExternalProviderClient(@NonNull ContentProviderRecord cpr,
+            IBinder externalProcessToken) {
+        return cpr.removeExternalProcessHandleLocked(externalProcessToken);
+    }
+
+    /**
+     * Note the time a process is no longer hosting any content providers.
+     */
+    public void setLastProviderTime(@NonNull ProcessRecord proc, long uptimeMs) {
+        proc.mProviders.setLastProviderTime(uptimeMs);
+    }
+
+    /**
+     * Note that a process has connected to a content provider.
+     */
+    public void addProviderConnection(@NonNull ProcessRecord client,
+            ContentProviderConnection cpc) {
+        client.mProviders.addProviderConnection(cpc);
+    }
+
+    /**
+     * Note that a process is no longer connected to a content provider.
+     */
+    public void removeProviderConnection(@NonNull ProcessRecord client,
+            ContentProviderConnection cpc) {
+        client.mProviders.addProviderConnection(cpc);
+    }
+
+    /********************** Content Provider State Events **********************/
+    /*************************** Service State Events **************************/
+    /**
+     * Note that a process has started hosting a service.
+     */
+    public boolean startService(@NonNull ProcessServiceRecord psr, ServiceRecord sr) {
+        return psr.startService(sr);
+    }
+
+    /**
+     * Note that a process has stopped hosting a service.
+     */
+    public boolean stopService(@NonNull ProcessServiceRecord psr, ServiceRecord sr) {
+        return psr.stopService(sr);
+    }
+
+    /**
+     * Remove all services that the process is hosting.
+     */
+    public void stopAllServices(@NonNull ProcessServiceRecord psr) {
+        psr.stopAllServices();
+    }
+
+    /**
+     * Note that a process's service has started executing.
+     */
+    public void startExecutingService(@NonNull ProcessServiceRecord psr, ServiceRecord sr) {
+        psr.startExecutingService(sr);
+    }
+
+    /**
+     * Note that a process's service has stopped executing.
+     */
+    public void stopExecutingService(@NonNull ProcessServiceRecord psr, ServiceRecord sr) {
+        psr.stopExecutingService(sr);
+    }
+
+    /**
+     * Note all executing services a process has has stopped.
+     */
+    public void stopAllExecutingServices(@NonNull ProcessServiceRecord psr) {
+        psr.stopAllExecutingServices();
+    }
+
+    /**
+     * Note that process has bound to a service.
+     */
+    public void addConnection(@NonNull ProcessServiceRecord psr, ConnectionRecord cr) {
+        psr.addConnection(cr);
+    }
+
+    /**
+     * Note that process has unbound from a service.
+     */
+    public void removeConnection(@NonNull ProcessServiceRecord psr, ConnectionRecord cr) {
+        psr.removeConnection(cr);
+    }
+
+    /**
+     * Remove all bindings a process has to services.
+     */
+    public void removeAllConnections(@NonNull ProcessServiceRecord psr) {
+        psr.removeAllConnections();
+        psr.removeAllSdkSandboxConnections();
+    }
+
+    /**
+     * Note whether an executing service should be considered in the foreground or not.
+     */
+    public void setExecServicesFg(@NonNull ProcessServiceRecord psr, boolean execServicesFg) {
+        psr.setExecServicesFg(execServicesFg);
+    }
+
+    /**
+     * Note whether a service is in the foreground or not and what type of FGS, if so.
+     */
+    public void setHasForegroundServices(@NonNull ProcessServiceRecord psr,
+            boolean hasForegroundServices,
+            int fgServiceTypes, boolean hasTypeNoneFgs) {
+        psr.setHasForegroundServices(hasForegroundServices, fgServiceTypes, hasTypeNoneFgs);
+    }
+
+    /**
+     * Note whether a service has a client activity or not.
+     */
+    public void setHasClientActivities(@NonNull ProcessServiceRecord psr,
+            boolean hasClientActivities) {
+        psr.setHasClientActivities(hasClientActivities);
+    }
+
+    /**
+     * Note whether a service should be treated like an activity or not.
+     */
+    public void setTreatLikeActivity(@NonNull ProcessServiceRecord psr, boolean treatLikeActivity) {
+        psr.setTreatLikeActivity(treatLikeActivity);
+    }
+
+    /**
+     * Note whether a process has bound to a service with
+     * {@link android.content.Context.BIND_ABOVE_CLIENT} or not.
+     */
+    public void setHasAboveClient(@NonNull ProcessServiceRecord psr, boolean hasAboveClient) {
+        psr.setHasAboveClient(hasAboveClient);
+    }
+
+    /**
+     * Recompute whether a process has bound to a service with
+     * {@link android.content.Context.BIND_ABOVE_CLIENT} or not.
+     */
+    public void updateHasAboveClientLocked(@NonNull ProcessServiceRecord psr) {
+        psr.updateHasAboveClientLocked();
+    }
+
+    /**
+     * Cleanup a process's state.
+     */
+    public void onCleanupApplicationRecord(@NonNull ProcessServiceRecord psr) {
+        psr.onCleanupApplicationRecordLocked();
+    }
+
+    /**
+     * Set which process is hosting a service.
+     */
+    public void setHostProcess(@NonNull ServiceRecord sr, @Nullable ProcessRecord host) {
+        sr.app = host;
+    }
+
+    /**
+     * Note whether a service is a Foreground Service or not
+     */
+    public void setIsForegroundService(@NonNull ServiceRecord sr, boolean isFgs) {
+        sr.isForeground = isFgs;
+    }
+
+    /**
+     * Note the Foreground Service type of a service.
+     */
+    public void setForegroundServiceType(@NonNull ServiceRecord sr,
+            @ServiceInfo.ForegroundServiceType int fgsType) {
+        sr.foregroundServiceType = fgsType;
+    }
+
+    /**
+     * Note the start time of a short foreground service.
+     */
+    public void setShortFgsInfo(@NonNull ServiceRecord sr, long uptimeNow) {
+        sr.setShortFgsInfo(uptimeNow);
+    }
+
+    /**
+     * Note that a short foreground service has stopped.
+     */
+    public void clearShortFgsInfo(@NonNull ServiceRecord sr) {
+        sr.clearShortFgsInfo();
+    }
+
+    /**
+     * Note the last time a service was active.
+     */
+    public void setServiceLastActivityTime(@NonNull ServiceRecord sr, long lastActivityUpdateMs) {
+        sr.lastActivity = lastActivityUpdateMs;
+    }
+
+    /**
+     * Note that a service start was requested.
+     */
+    public void setStartRequested(@NonNull ServiceRecord sr, boolean startRequested) {
+        sr.startRequested = startRequested;
+    }
+
+    /**
+     * Note the last time the service was bound by a Top process with
+     * {@link android.content.Context.BIND_ALMOST_PERCEPTIBLE}
+     */
+    public void setLastTopAlmostPerceptibleBindRequest(@NonNull ServiceRecord sr,
+            long lastTopAlmostPerceptibleBindRequestUptimeMs) {
+        sr.lastTopAlmostPerceptibleBindRequestUptimeMs =
+                lastTopAlmostPerceptibleBindRequestUptimeMs;
+    }
+
+    /**
+     * Recompute whether a process has bound to a service with
+     * {@link android.content.Context.BIND_ALMOST_PERCEPTIBLE} or not.
+     */
+    public void updateHasTopStartedAlmostPerceptibleServices(@NonNull ProcessServiceRecord psr) {
+        psr.updateHasTopStartedAlmostPerceptibleServices();
+    }
+
+    /**
+     * Builder for ProcessStateController.
+     */
+    public static class Builder {
+        private final ActivityManagerService mAms;
+        private final ProcessList mProcessList;
+        private final ActiveUids mActiveUids;
+
+        private ServiceThread mHandlerThread = null;
+        private OomAdjuster.Injector mOomAdjInjector = null;
+        private boolean mUseOomAdjusterModernImpl = false;
+
+        public Builder(ActivityManagerService ams, ProcessList processList, ActiveUids activeUids) {
+            mAms = ams;
+            mProcessList = processList;
+            mActiveUids = activeUids;
+        }
+
+        /**
+         * Build the ProcessStateController object.
+         */
+        public ProcessStateController build() {
+            if (mHandlerThread == null) {
+                mHandlerThread = OomAdjuster.createAdjusterThread();
+            }
+            if (mOomAdjInjector == null) {
+                mOomAdjInjector = new OomAdjuster.Injector();
+            }
+            return new ProcessStateController(mAms, mProcessList, mActiveUids, mHandlerThread,
+                    mOomAdjInjector, mUseOomAdjusterModernImpl);
+        }
+
+        /**
+         * For Testing Purposes. Set what thread OomAdjuster will offload tasks on to.
+         */
+        public Builder setHandlerThread(ServiceThread handlerThread) {
+            mHandlerThread = handlerThread;
+            return this;
+        }
+
+        /**
+         * For Testing Purposes. Set an injector for OomAdjuster.
+         */
+        public Builder setOomAdjusterInjector(OomAdjuster.Injector injector) {
+            mOomAdjInjector = injector;
+            return this;
+        }
+
+        /**
+         * Set which implementation of OomAdjuster to use.
+         */
+        public Builder useModernOomAdjuster(boolean use) {
+            mUseOomAdjusterModernImpl = use;
+            return this;
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/am/ProcessStateRecord.java b/services/core/java/com/android/server/am/ProcessStateRecord.java
index bc990d9..b0f808b 100644
--- a/services/core/java/com/android/server/am/ProcessStateRecord.java
+++ b/services/core/java/com/android/server/am/ProcessStateRecord.java
@@ -19,14 +19,11 @@
 import static android.app.ActivityManager.PROCESS_CAPABILITY_NONE;
 import static android.app.ActivityManager.PROCESS_STATE_CACHED_EMPTY;
 import static android.app.ActivityManager.PROCESS_STATE_NONEXISTENT;
-import static android.app.ActivityManagerInternal.OOM_ADJ_REASON_UI_VISIBILITY;
 import static android.app.ProcessMemoryState.HOSTING_COMPONENT_TYPE_ACTIVITY;
 import static android.app.ProcessMemoryState.HOSTING_COMPONENT_TYPE_BROADCAST_RECEIVER;
 import static android.app.ProcessMemoryState.HOSTING_COMPONENT_TYPE_STARTED_SERVICE;
 
-import static com.android.server.am.ActivityManagerDebugConfig.DEBUG_OOM_ADJ;
 import static com.android.server.am.ProcessList.CACHED_APP_MIN_ADJ;
-import static com.android.server.am.ProcessRecord.TAG;
 import static com.android.server.wm.WindowProcessController.ACTIVITY_STATE_FLAG_IS_PAUSING_OR_PAUSED;
 import static com.android.server.wm.WindowProcessController.ACTIVITY_STATE_FLAG_IS_STOPPING;
 import static com.android.server.wm.WindowProcessController.ACTIVITY_STATE_FLAG_IS_STOPPING_FINISHING;
@@ -38,7 +35,6 @@
 import android.content.ComponentName;
 import android.os.SystemClock;
 import android.os.Trace;
-import android.util.Slog;
 import android.util.TimeUtils;
 
 import com.android.internal.annotations.CompositeRWLock;
@@ -790,15 +786,7 @@
 
     @GuardedBy("mService")
     void setRunningRemoteAnimation(boolean runningRemoteAnimation) {
-        if (mRunningRemoteAnimation == runningRemoteAnimation) {
-            return;
-        }
         mRunningRemoteAnimation = runningRemoteAnimation;
-        if (DEBUG_OOM_ADJ) {
-            Slog.i(TAG, "Setting runningRemoteAnimation=" + runningRemoteAnimation
-                    + " for pid=" + mApp.getPid());
-        }
-        mService.updateOomAdjLocked(mApp, OOM_ADJ_REASON_UI_VISIBILITY);
     }
 
     @GuardedBy({"mService", "mProcLock"})
diff --git a/services/core/java/com/android/server/am/ServiceRecord.java b/services/core/java/com/android/server/am/ServiceRecord.java
index b9cdf27..92d33c9 100644
--- a/services/core/java/com/android/server/am/ServiceRecord.java
+++ b/services/core/java/com/android/server/am/ServiceRecord.java
@@ -1248,7 +1248,7 @@
             app.mServices.updateBoundClientUids();
             app.mServices.updateHostingComonentTypeForBindingsLocked();
         }
-        app = proc;
+        ams.mProcessStateController.setHostProcess(this, proc);
         updateProcessStateOnRequest();
         if (pendingConnectionGroup > 0 && proc != null) {
             final ProcessServiceRecord psr = proc.mServices;
diff --git a/services/core/java/com/android/server/am/flags.aconfig b/services/core/java/com/android/server/am/flags.aconfig
index 7873d34..adf0e64 100644
--- a/services/core/java/com/android/server/am/flags.aconfig
+++ b/services/core/java/com/android/server/am/flags.aconfig
@@ -218,3 +218,10 @@
     description: "Set reset_on_fork flag."
     bug: "370988407"
 }
+
+flag {
+    name: "push_global_state_to_oomadjuster"
+    namespace: "backstage_power"
+    description: "Migrate OomAdjuster pulled device state to a push model"
+    bug: "302575389"
+}
diff --git a/services/core/java/com/android/server/appop/AppOpsService.java b/services/core/java/com/android/server/appop/AppOpsService.java
index 596e375..e0cf96f 100644
--- a/services/core/java/com/android/server/appop/AppOpsService.java
+++ b/services/core/java/com/android/server/appop/AppOpsService.java
@@ -72,9 +72,6 @@
 import static android.content.pm.PermissionInfo.PROTECTION_FLAG_APPOP;
 import static android.permission.flags.Flags.deviceAwareAppOpNewSchemaEnabled;
 
-import static com.android.internal.util.FrameworkStatsLog.APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED__BINDER_API__CHECK_OPERATION;
-import static com.android.internal.util.FrameworkStatsLog.APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED__BINDER_API__NOTE_OPERATION;
-import static com.android.internal.util.FrameworkStatsLog.APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED__BINDER_API__NOTE_PROXY_OPERATION;
 import static com.android.server.appop.AppOpsService.ModeCallback.ALL_OPS;
 
 import android.Manifest;
@@ -163,7 +160,6 @@
 import com.android.internal.pm.pkg.component.ParsedAttribution;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.DumpUtils;
-import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.Preconditions;
 import com.android.internal.util.XmlUtils;
 import com.android.internal.util.function.pooled.PooledLambda;
@@ -2833,26 +2829,12 @@
 
     @Override
     public int checkOperation(int code, int uid, String packageName) {
-        if (Flags.appopAccessTrackingLoggingEnabled()) {
-            FrameworkStatsLog.write(
-                    FrameworkStatsLog.APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED,
-                    uid, code,
-                    APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED__BINDER_API__CHECK_OPERATION,
-                    false);
-        }
         return mCheckOpsDelegateDispatcher.checkOperation(code, uid, packageName, null,
                 Context.DEVICE_ID_DEFAULT, false /*raw*/);
     }
 
     @Override
     public int checkOperationForDevice(int code, int uid, String packageName, int virtualDeviceId) {
-        if (Flags.appopAccessTrackingLoggingEnabled()) {
-            FrameworkStatsLog.write(
-                    FrameworkStatsLog.APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED,
-                    uid, code,
-                    APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED__BINDER_API__CHECK_OPERATION,
-                    false);
-        }
         return mCheckOpsDelegateDispatcher.checkOperation(code, uid, packageName, null,
                 virtualDeviceId, false /*raw*/);
     }
@@ -3033,14 +3015,6 @@
     public SyncNotedAppOp noteProxyOperationWithState(int code,
             AttributionSourceState attributionSourceState, boolean shouldCollectAsyncNotedOp,
             String message, boolean shouldCollectMessage, boolean skipProxyOperation) {
-        if (Flags.appopAccessTrackingLoggingEnabled()) {
-            FrameworkStatsLog.write(
-                    FrameworkStatsLog.APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED,
-                    attributionSourceState.uid, code,
-                    APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED__BINDER_API__NOTE_PROXY_OPERATION,
-                    attributionSourceState.attributionTag != null);
-        }
-
         AttributionSource attributionSource = new AttributionSource(attributionSourceState);
         return mCheckOpsDelegateDispatcher.noteProxyOperation(code, attributionSource,
                 shouldCollectAsyncNotedOp, message, shouldCollectMessage, skipProxyOperation);
@@ -3122,14 +3096,6 @@
     public SyncNotedAppOp noteOperation(int code, int uid, String packageName,
             String attributionTag, boolean shouldCollectAsyncNotedOp, String message,
             boolean shouldCollectMessage) {
-        if (Flags.appopAccessTrackingLoggingEnabled()) {
-            FrameworkStatsLog.write(
-                    FrameworkStatsLog.APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED,
-                    uid, code,
-                    APP_OP_NOTE_OP_OR_CHECK_OP_BINDER_API_CALLED__BINDER_API__NOTE_OPERATION,
-                    attributionTag != null);
-        }
-
         return mCheckOpsDelegateDispatcher.noteOperation(code, uid, packageName,
                 attributionTag, Context.DEVICE_ID_DEFAULT, shouldCollectAsyncNotedOp, message,
                 shouldCollectMessage);
diff --git a/services/core/java/com/android/server/devicestate/DeviceStateManagerService.java b/services/core/java/com/android/server/devicestate/DeviceStateManagerService.java
index 1094bee..6216a58 100644
--- a/services/core/java/com/android/server/devicestate/DeviceStateManagerService.java
+++ b/services/core/java/com/android/server/devicestate/DeviceStateManagerService.java
@@ -785,8 +785,8 @@
                         + " IDeviceStateManagerCallback.");
             }
 
-            ProcessRecord record = new ProcessRecord(callback, pid, this::handleProcessDied,
-                    mHandler);
+            final ProcessRecord record =
+                    new ProcessRecord(callback, pid, this::handleProcessDied, mHandler);
             try {
                 callback.asBinder().linkToDeath(record, 0);
             } catch (RemoteException ex) {
@@ -797,8 +797,8 @@
             // Callback clients should not be notified of invalid device states, so calls to
             // #getDeviceStateInfoLocked should be gated on checks if a committed state is present
             // before getting the device state info.
-            DeviceStateInfo currentInfo = mCommittedState.isPresent()
-                    ? getDeviceStateInfoLocked() : null;
+            final DeviceStateInfo currentInfo =
+                    mCommittedState.isPresent() ? getDeviceStateInfoLocked() : null;
             if (currentInfo != null) {
                 // If there is not a committed state we'll wait to notify the process of the initial
                 // value.
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index d71826f..179ec63 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -1978,7 +1978,7 @@
                         // handles stopping the projection.
                         Slog.w(TAG, "Content Recording: failed to start mirroring - "
                                 + "releasing virtual display " + displayId);
-                        releaseVirtualDisplayInternal(callback.asBinder());
+                        releaseVirtualDisplayInternal(callback.asBinder(), callingUid);
                         return Display.INVALID_DISPLAY;
                     } else if (projection != null) {
                         // Indicate that this projection has been used to record, and can't be used
@@ -2067,7 +2067,7 @@
         // Something weird happened and the logical display was not created.
         Slog.w(TAG, "Rejecting request to create virtual display "
                 + "because the logical display was not created.");
-        mVirtualDisplayAdapter.releaseVirtualDisplayLocked(callback.asBinder());
+        mVirtualDisplayAdapter.releaseVirtualDisplayLocked(callback.asBinder(), callingUid);
         mDisplayDeviceRepo.onDisplayDeviceEvent(device,
                 DisplayAdapter.DISPLAY_DEVICE_EVENT_REMOVED);
         return -1;
@@ -2094,14 +2094,14 @@
         }
     }
 
-    private void releaseVirtualDisplayInternal(IBinder appToken) {
+    private void releaseVirtualDisplayInternal(IBinder appToken, int callingUid) {
         synchronized (mSyncRoot) {
             if (mVirtualDisplayAdapter == null) {
                 return;
             }
 
             DisplayDevice device =
-                    mVirtualDisplayAdapter.releaseVirtualDisplayLocked(appToken);
+                    mVirtualDisplayAdapter.releaseVirtualDisplayLocked(appToken, callingUid);
             Slog.d(TAG, "Virtual Display: Display Device released");
             if (device != null) {
                 // TODO: multi-display - handle virtual displays the same as other display adapters.
@@ -4620,9 +4620,10 @@
 
         @Override // Binder call
         public void releaseVirtualDisplay(IVirtualDisplayCallback callback) {
+            final int callingUid = Binder.getCallingUid();
             final long token = Binder.clearCallingIdentity();
             try {
-                releaseVirtualDisplayInternal(callback.asBinder());
+                releaseVirtualDisplayInternal(callback.asBinder(), callingUid);
             } finally {
                 Binder.restoreCallingIdentity(token);
             }
diff --git a/services/core/java/com/android/server/display/VirtualDisplayAdapter.java b/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
index e77c5ec..4211453 100644
--- a/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
+++ b/services/core/java/com/android/server/display/VirtualDisplayAdapter.java
@@ -55,12 +55,14 @@
 import android.os.SystemProperties;
 import android.util.ArrayMap;
 import android.util.Slog;
+import android.util.SparseIntArray;
 import android.view.Display;
 import android.view.DisplayCutout;
 import android.view.DisplayShape;
 import android.view.Surface;
 import android.view.SurfaceControl;
 
+import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.display.feature.DisplayManagerFlags;
 
@@ -85,6 +87,11 @@
     private static final AtomicInteger sNextUniqueIndex = new AtomicInteger(0);
 
     private final ArrayMap<IBinder, VirtualDisplayDevice> mVirtualDisplayDevices = new ArrayMap<>();
+
+    private final int mMaxDevices;
+    private final int mMaxDevicesPerPackage;
+    private final SparseIntArray mNoOfDevicesPerPackage = new SparseIntArray();
+
     private final Handler mHandler;
     private final SurfaceControlDisplayFactory mSurfaceControlDisplayFactory;
 
@@ -114,8 +121,31 @@
         super(syncRoot, context, handler, listener, TAG, featureFlags);
         mHandler = handler;
         mSurfaceControlDisplayFactory = surfaceControlDisplayFactory;
+
+        mMaxDevices = context.getResources().getInteger(R.integer.config_virtualDisplayLimit);
+        if (mMaxDevices < 1) {
+            throw new IllegalArgumentException("The limit of virtual displays must be >= 1");
+        }
+        mMaxDevicesPerPackage =
+                context.getResources().getInteger(R.integer.config_virtualDisplayLimitPerPackage);
+        if (mMaxDevicesPerPackage < 1) {
+            throw new IllegalArgumentException(
+                    "The limit of virtual displays per package must be >= 1");
+        }
     }
 
+    /**
+     * Create a virtual display
+     * @param callback The callback
+     * @param projection The media projection
+     * @param ownerUid The UID of the package creating a display
+     * @param ownerPackageName The name of the package creating a display
+     * @param uniqueId The unique ID of the display device
+     * @param surface The surface
+     * @param flags The flags
+     * @param virtualDisplayConfig The config
+     * @return The display device created
+     */
     public DisplayDevice createVirtualDisplayLocked(IVirtualDisplayCallback callback,
             IMediaProjection projection, int ownerUid, String ownerPackageName, String uniqueId,
             Surface surface, int flags, VirtualDisplayConfig virtualDisplayConfig) {
@@ -126,6 +156,22 @@
             return null;
         }
 
+        if (getFeatureFlags().isVirtualDisplayLimitEnabled()
+                && mVirtualDisplayDevices.size() >= mMaxDevices) {
+            Slog.w(TAG, "Rejecting request to create private virtual display because "
+                    + mMaxDevices + " devices already exist.");
+            return null;
+        }
+
+        int noOfDevices = mNoOfDevicesPerPackage.get(ownerUid, /* valueIfKeyNotFound= */ 0);
+        if (getFeatureFlags().isVirtualDisplayLimitEnabled()
+                && noOfDevices >= mMaxDevicesPerPackage) {
+            Slog.w(TAG, "Rejecting request to create private virtual display because "
+                    + mMaxDevicesPerPackage + " devices already exist for package "
+                    + ownerPackageName + ".");
+            return null;
+        }
+
         String name = virtualDisplayConfig.getName();
         boolean secure = (flags & VIRTUAL_DISPLAY_FLAG_SECURE) != 0;
 
@@ -140,6 +186,9 @@
                 projection, mediaProjectionCallback, uniqueId, virtualDisplayConfig);
 
         mVirtualDisplayDevices.put(appToken, device);
+        if (getFeatureFlags().isVirtualDisplayLimitEnabled()) {
+            mNoOfDevicesPerPackage.put(ownerUid, noOfDevices + 1);
+        }
 
         try {
             if (projection != null) {
@@ -150,7 +199,7 @@
             appToken.linkToDeath(device, 0);
         } catch (RemoteException ex) {
             Slog.e(TAG, "Virtual Display: error while setting up VirtualDisplayDevice", ex);
-            mVirtualDisplayDevices.remove(appToken);
+            removeVirtualDisplayDeviceLocked(appToken, ownerUid);
             device.destroyLocked(false);
             return null;
         }
@@ -194,8 +243,15 @@
         }
     }
 
-    public DisplayDevice releaseVirtualDisplayLocked(IBinder appToken) {
-        VirtualDisplayDevice device = mVirtualDisplayDevices.remove(appToken);
+    /**
+     * Release a virtual display that was previously created
+     * @param appToken The token to identify the display
+     * @param ownerUid The UID of the package, used to keep track of and limit the number of
+     *                 displays created per package
+     * @return The display device that has been removed
+     */
+    public DisplayDevice releaseVirtualDisplayLocked(IBinder appToken, int ownerUid) {
+        VirtualDisplayDevice device = removeVirtualDisplayDeviceLocked(appToken, ownerUid);
         if (device != null) {
             Slog.v(TAG, "Release VirtualDisplay " + device.mName);
             device.destroyLocked(true);
@@ -228,10 +284,6 @@
                 : ("," + uid + "," + config.getName() + "," + sNextUniqueIndex.getAndIncrement()));
     }
 
-    private void handleBinderDiedLocked(IBinder appToken) {
-        mVirtualDisplayDevices.remove(appToken);
-    }
-
     private void handleMediaProjectionStoppedLocked(IBinder appToken) {
         VirtualDisplayDevice device = mVirtualDisplayDevices.get(appToken);
         if (device != null) {
@@ -241,6 +293,18 @@
         }
     }
 
+    private VirtualDisplayDevice removeVirtualDisplayDeviceLocked(IBinder appToken, int ownerUid) {
+        int noOfDevices = mNoOfDevicesPerPackage.get(ownerUid, /* valueIfKeyNotFound= */ 0);
+        if (getFeatureFlags().isVirtualDisplayLimitEnabled()) {
+            if (noOfDevices <= 1) {
+                mNoOfDevicesPerPackage.delete(ownerUid);
+            } else {
+                mNoOfDevicesPerPackage.put(ownerUid, noOfDevices - 1);
+            }
+        }
+        return mVirtualDisplayDevices.remove(appToken);
+    }
+
     private final class VirtualDisplayDevice extends DisplayDevice implements DeathRecipient {
         private static final int PENDING_SURFACE_CHANGE = 0x01;
         private static final int PENDING_RESIZE = 0x02;
@@ -300,7 +364,7 @@
         @Override
         public void binderDied() {
             synchronized (getSyncRoot()) {
-                handleBinderDiedLocked(mAppToken);
+                removeVirtualDisplayDeviceLocked(mAppToken, mOwnerUid);
                 Slog.i(TAG, "Virtual display device released because application token died: "
                     + mOwnerPackageName);
                 destroyLocked(false);
diff --git a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
index 5284d1c..99ced7f 100644
--- a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
+++ b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
@@ -199,6 +199,11 @@
             Flags.FLAG_IDLE_SCREEN_CONFIG_IN_SUBSCRIBING_LIGHT_SENSOR,
             Flags::idleScreenConfigInSubscribingLightSensor);
 
+    private final FlagState mVirtualDisplayLimit =
+            new FlagState(
+                    Flags.FLAG_VIRTUAL_DISPLAY_LIMIT,
+                    Flags::virtualDisplayLimit);
+
     private final FlagState mNormalBrightnessForDozeParameter = new FlagState(
             Flags.FLAG_NORMAL_BRIGHTNESS_FOR_DOZE_PARAMETER,
             Flags::normalBrightnessForDozeParameter
@@ -207,6 +212,10 @@
             Flags.FLAG_BLOCK_AUTOBRIGHTNESS_CHANGES_ON_STYLUS_USAGE,
             Flags::blockAutobrightnessChangesOnStylusUsage
     );
+    private final FlagState mIsUserRefreshRateForExternalDisplayEnabled = new FlagState(
+            Flags.FLAG_ENABLE_USER_REFRESH_RATE_FOR_EXTERNAL_DISPLAY,
+            Flags::enableUserRefreshRateForExternalDisplay
+    );
 
     private final FlagState mEnableBatteryStatsForAllDisplays = new FlagState(
             Flags.FLAG_ENABLE_BATTERY_STATS_FOR_ALL_DISPLAYS,
@@ -416,6 +425,10 @@
         return mNewHdrBrightnessModifier.isEnabled();
     }
 
+    public boolean isVirtualDisplayLimitEnabled() {
+        return mVirtualDisplayLimit.isEnabled();
+    }
+
     /**
      * @return Whether the useDozeBrightness parameter should be used
      */
@@ -447,6 +460,14 @@
     }
 
     /**
+     * @return {@code true} if need to use user refresh rate settings for
+     * external displays.
+     */
+    public boolean isUserRefreshRateForExternalDisplayEnabled() {
+        return mIsUserRefreshRateForExternalDisplayEnabled.isEnabled();
+    }
+
+    /**
      * dumps all flagstates
      * @param pw printWriter
      */
@@ -487,10 +508,12 @@
         pw.println(" " + mOffloadDozeOverrideHoldsWakelock);
         pw.println(" " + mOffloadSessionCancelBlockScreenOn);
         pw.println(" " + mNewHdrBrightnessModifier);
+        pw.println(" " + mVirtualDisplayLimit);
         pw.println(" " + mNormalBrightnessForDozeParameter);
         pw.println(" " + mIdleScreenConfigInSubscribingLightSensor);
         pw.println(" " + mEnableBatteryStatsForAllDisplays);
         pw.println(" " + mBlockAutobrightnessChangesOnStylusUsage);
+        pw.println(" " + mIsUserRefreshRateForExternalDisplayEnabled);
     }
 
     private static class FlagState {
diff --git a/services/core/java/com/android/server/display/feature/display_flags.aconfig b/services/core/java/com/android/server/display/feature/display_flags.aconfig
index 252ed09..2f04d9e 100644
--- a/services/core/java/com/android/server/display/feature/display_flags.aconfig
+++ b/services/core/java/com/android/server/display/feature/display_flags.aconfig
@@ -348,6 +348,14 @@
 }
 
 flag {
+    name: "virtual_display_limit"
+    namespace: "display_manager"
+    description: "Limit the number of virtual displays that can be created."
+    bug: "261791612"
+    is_fixed_read_only: true
+}
+
+flag {
     name: "idle_screen_config_in_subscribing_light_sensor"
     namespace: "display_manager"
     description: "Account for Idle screen refresh rate configs while subscribing to light sensor"
@@ -373,3 +381,14 @@
     bug: "352411468"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "enable_user_refresh_rate_for_external_display"
+    namespace: "display_manager"
+    description: "Apply refresh rate from user preferred display mode to external displays"
+    bug: "370657357"
+    is_fixed_read_only: true
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
index 18e0d6e..ffa64bf 100644
--- a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
+++ b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
@@ -1194,6 +1194,13 @@
         @GuardedBy("mLock")
         private void updateRefreshRateSettingLocked(float minRefreshRate, float peakRefreshRate,
                 float defaultRefreshRate, int displayId) {
+            if (mDisplayObserver.isExternalDisplayLocked(displayId)) {
+                if (mLoggingEnabled) {
+                    Slog.d(TAG, "skip updateRefreshRateSettingLocked for external display "
+                            + displayId);
+                }
+                return;
+            }
             // TODO(b/156304339): The logic in here, aside from updating the refresh rate votes, is
             // used to predict if we're going to be doing frequent refresh rate switching, and if
             // so, enable the brightness observer. The logic here is more complicated and fragile
@@ -1243,6 +1250,8 @@
         }
 
         private void removeRefreshRateSetting(int displayId) {
+            mVotesStorage.updateVote(displayId, Vote.PRIORITY_USER_SETTING_PEAK_REFRESH_RATE,
+                    null);
             mVotesStorage.updateVote(displayId, Vote.PRIORITY_USER_SETTING_PEAK_RENDER_FRAME_RATE,
                     null);
             mVotesStorage.updateVote(displayId, Vote.PRIORITY_USER_SETTING_MIN_RENDER_FRAME_RATE,
@@ -1458,11 +1467,11 @@
         public void onDisplayAdded(int displayId) {
             updateDisplayDeviceConfig(displayId);
             DisplayInfo displayInfo = getDisplayInfo(displayId);
+            registerExternalDisplay(displayInfo);
             updateDisplayModes(displayId, displayInfo);
             updateLayoutLimitedFrameRate(displayId, displayInfo);
             updateUserSettingDisplayPreferredSize(displayInfo);
             updateDisplaysPeakRefreshRateAndResolution(displayInfo);
-            addDisplaysSynchronizedPeakRefreshRate(displayInfo);
         }
 
         @Override
@@ -1477,7 +1486,7 @@
             updateLayoutLimitedFrameRate(displayId, null);
             removeUserSettingDisplayPreferredSize(displayId);
             removeDisplaysPeakRefreshRateAndResolution(displayId);
-            removeDisplaysSynchronizedPeakRefreshRate(displayId);
+            unregisterExternalDisplay(displayId);
         }
 
         @Override
@@ -1489,6 +1498,30 @@
             updateUserSettingDisplayPreferredSize(displayInfo);
         }
 
+        private void registerExternalDisplay(DisplayInfo displayInfo) {
+            if (displayInfo == null || displayInfo.type != Display.TYPE_EXTERNAL) {
+                return;
+            }
+            synchronized (mLock) {
+                mExternalDisplaysConnected.add(displayInfo.displayId);
+                if (mExternalDisplaysConnected.size() == 1) {
+                    addDisplaysSynchronizedPeakRefreshRate();
+                }
+            }
+        }
+
+        private void unregisterExternalDisplay(int displayId) {
+            synchronized (mLock) {
+                if (!isExternalDisplayLocked(displayId)) {
+                    return;
+                }
+                mExternalDisplaysConnected.remove(displayId);
+                if (mExternalDisplaysConnected.isEmpty()) {
+                    removeDisplaysSynchronizedPeakRefreshRate();
+                }
+            }
+        }
+
         boolean isExternalDisplayLocked(int displayId) {
             return mExternalDisplaysConnected.contains(displayId);
         }
@@ -1534,10 +1567,24 @@
                 return;
             }
 
-            mVotesStorage.updateVote(info.displayId,
-                    Vote.PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE,
-                    Vote.forSize(/* width */ preferredMode.getPhysicalWidth(),
-                            /* height */ preferredMode.getPhysicalHeight()));
+            if (info.type == Display.TYPE_EXTERNAL
+                    && mDisplayManagerFlags.isUserRefreshRateForExternalDisplayEnabled()
+                    && !isRefreshRateSynchronizationEnabled()) {
+                mVotesStorage.updateVote(info.displayId,
+                        Vote.PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE,
+                        Vote.forSizeAndPhysicalRefreshRatesRange(
+                                /* minWidth */ preferredMode.getPhysicalWidth(),
+                                /* minHeight */ preferredMode.getPhysicalHeight(),
+                                /* width */ preferredMode.getPhysicalWidth(),
+                                /* height */ preferredMode.getPhysicalHeight(),
+                                /* minRefreshRate */ preferredMode.getRefreshRate(),
+                                /* maxRefreshRate */ preferredMode.getRefreshRate()));
+            } else {
+                mVotesStorage.updateVote(info.displayId,
+                        Vote.PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE,
+                        Vote.forSize(/* width */ preferredMode.getPhysicalWidth(),
+                                /* height */ preferredMode.getPhysicalHeight()));
+            }
         }
 
         @Nullable
@@ -1584,17 +1631,10 @@
          * Sets 60Hz target refresh rate as the vote with
          * {@link Vote#PRIORITY_SYNCHRONIZED_REFRESH_RATE} priority.
          */
-        private void addDisplaysSynchronizedPeakRefreshRate(@Nullable final DisplayInfo info) {
-            if (info == null || info.type != Display.TYPE_EXTERNAL
-                    || !isRefreshRateSynchronizationEnabled()) {
+        private void addDisplaysSynchronizedPeakRefreshRate() {
+            if (!isRefreshRateSynchronizationEnabled()) {
                 return;
             }
-            synchronized (mLock) {
-                mExternalDisplaysConnected.add(info.displayId);
-                if (mExternalDisplaysConnected.size() != 1) {
-                    return;
-                }
-            }
             // set minRefreshRate as the max refresh rate.
             mVotesStorage.updateGlobalVote(Vote.PRIORITY_SYNCHRONIZED_REFRESH_RATE,
                     Vote.forPhysicalRefreshRates(
@@ -1610,19 +1650,10 @@
                                     + SYNCHRONIZED_REFRESH_RATE_TOLERANCE));
         }
 
-        private void removeDisplaysSynchronizedPeakRefreshRate(final int displayId) {
+        private void removeDisplaysSynchronizedPeakRefreshRate() {
             if (!isRefreshRateSynchronizationEnabled()) {
                 return;
             }
-            synchronized (mLock) {
-                if (!isExternalDisplayLocked(displayId)) {
-                    return;
-                }
-                mExternalDisplaysConnected.remove(displayId);
-                if (!mExternalDisplaysConnected.isEmpty()) {
-                    return;
-                }
-            }
             mVotesStorage.updateGlobalVote(Vote.PRIORITY_SYNCHRONIZED_REFRESH_RATE, null);
             mVotesStorage.updateGlobalVote(Vote.PRIORITY_SYNCHRONIZED_RENDER_FRAME_RATE, null);
         }
diff --git a/services/core/java/com/android/server/display/mode/Vote.java b/services/core/java/com/android/server/display/mode/Vote.java
index 459f9a6..f5abb05 100644
--- a/services/core/java/com/android/server/display/mode/Vote.java
+++ b/services/core/java/com/android/server/display/mode/Vote.java
@@ -44,7 +44,7 @@
     // It votes [minRefreshRate, Float.POSITIVE_INFINITY]
     int PRIORITY_USER_SETTING_MIN_RENDER_FRAME_RATE = 3;
 
-    // User setting preferred display resolution.
+    // User setting preferred display resolution, for external displays also includes refresh rate.
     int PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE = 4;
 
     // APP_REQUEST_RENDER_FRAME_RATE_RANGE is used to for internal apps to limit the render
diff --git a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
index 8b06dad..d1d5d48 100644
--- a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
+++ b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
@@ -1112,19 +1112,12 @@
                 continue;
             }
 
-            Log.w(
-                    TAG,
-                    TextUtils.formatSimple(
-                            "Revoking access to manager record id: %d, package: %s, userId:"
-                                    + " %d",
-                            manager.mManagerId, manager.mOwnerPackageName, userRecord.mUserId));
-
+            Slog.w(TAG, "Revoking access for " + manager.getDebugString());
             unregisterManagerLocked(manager.mManager, /* died */ false);
-
             try {
                 manager.mManager.invalidateInstance();
             } catch (RemoteException ex) {
-                Slog.w(TAG, "Failed to notify manager= " + manager + " of permission revocation.");
+                manager.logRemoteException("invalidateInstance", ex);
             }
         }
     }
@@ -2428,10 +2421,7 @@
             try {
                 mManager.notifyRequestFailed(requestId, reason);
             } catch (RemoteException ex) {
-                Slog.w(
-                        TAG,
-                        "Failed to notify manager of the request failure. Manager probably died.",
-                        ex);
+                logRemoteException("notifyRequestFailed", ex);
             }
         }
 
@@ -2444,7 +2434,7 @@
             try {
                 mManager.notifyRoutesUpdated(routes);
             } catch (RemoteException ex) {
-                Slog.w(TAG, "Failed to notify routes. Manager probably died.", ex);
+                logRemoteException("notifyRoutesUpdated", ex);
             }
         }
 
@@ -2457,10 +2447,7 @@
             try {
                 mManager.notifySessionUpdated(sessionInfo);
             } catch (RemoteException ex) {
-                Slog.w(
-                        TAG,
-                        "notifySessionUpdatedToManagers: Failed to notify. Manager probably died.",
-                        ex);
+                logRemoteException("notifySessionUpdated", ex);
             }
         }
 
@@ -2473,13 +2460,18 @@
             try {
                 mManager.notifySessionReleased(sessionInfo);
             } catch (RemoteException ex) {
-                Slog.w(
-                        TAG,
-                        "notifySessionReleasedToManagers: Failed to notify. Manager probably died.",
-                        ex);
+                logRemoteException("notifySessionReleased", ex);
             }
         }
 
+        private void logRemoteException(String operation, RemoteException exception) {
+            String message =
+                    TextUtils.formatSimple(
+                            "%s failed for %s due to %s",
+                            operation, getDebugString(), exception.toString());
+            Slog.w(TAG, message);
+        }
+
         private void updateScanningState(@ScanningState int scanningState) {
             if (mScanningState == scanningState) {
                 return;
@@ -2492,9 +2484,16 @@
                             UserHandler::updateDiscoveryPreferenceOnHandler, mUserRecord.mHandler));
         }
 
-        @Override
-        public String toString() {
-            return "Manager " + mOwnerPackageName + " (pid " + mOwnerPid + ")";
+        /** Returns a human readable representation of this manager record for logging purposes. */
+        public String getDebugString() {
+            return TextUtils.formatSimple(
+                    "Manager %s (id=%d,pid=%d,userId=%d,uid=%d,targetPkg=%s)",
+                    mOwnerPackageName,
+                    mManagerId,
+                    mOwnerPid,
+                    mUserRecord.mUserId,
+                    mOwnerUid,
+                    mTargetPackageName);
         }
     }
 
diff --git a/services/core/java/com/android/server/notification/ManagedServices.java b/services/core/java/com/android/server/notification/ManagedServices.java
index cd0a2a7..03fc60c 100644
--- a/services/core/java/com/android/server/notification/ManagedServices.java
+++ b/services/core/java/com/android/server/notification/ManagedServices.java
@@ -106,8 +106,7 @@
     protected final String TAG = getClass().getSimpleName().replace('$', '.');
     protected final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
 
-    protected static final int ON_BINDING_DIED_REBIND_DELAY_MS = 10000;
-    protected static final int ON_BINDING_DIED_REBIND_MSG = 1234;
+    private static final int ON_BINDING_DIED_REBIND_DELAY_MS = 10000;
     protected static final String ENABLED_SERVICES_SEPARATOR = ":";
     private static final String DB_VERSION_1 = "1";
     private static final String DB_VERSION_2 = "2";
@@ -857,13 +856,7 @@
             String approvedItem = getApprovedValue(pkgOrComponent);
 
             if (approvedItem != null) {
-                final ComponentName component = ComponentName.unflattenFromString(approvedItem);
                 if (enabled) {
-                    if (component != null && !isValidService(component, userId)) {
-                        Log.e(TAG, "Skip allowing " + mConfig.caption + " " + pkgOrComponent
-                                + " (userSet: " + userSet + ") for invalid service");
-                        return;
-                    }
                     approved.add(approvedItem);
                 } else {
                     approved.remove(approvedItem);
@@ -961,7 +954,7 @@
                 || isPackageOrComponentAllowed(component.getPackageName(), userId))) {
             return false;
         }
-        return isValidService(component, userId);
+        return componentHasBindPermission(component, userId);
     }
 
     private boolean componentHasBindPermission(ComponentName component, int userId) {
@@ -1313,12 +1306,11 @@
                     if (TextUtils.equals(getPackageName(approvedPackageOrComponent), packageName)) {
                         final ComponentName component = ComponentName.unflattenFromString(
                                 approvedPackageOrComponent);
-                        if (component != null && !isValidService(component, userId)) {
+                        if (component != null && !componentHasBindPermission(component, userId)) {
                             approved.removeAt(j);
                             if (DEBUG) {
                                 Slog.v(TAG, "Removing " + approvedPackageOrComponent
-                                        + " from approved list; no bind permission or "
-                                        + "service interface filter found "
+                                        + " from approved list; no bind permission found "
                                         + mConfig.bindPermission);
                             }
                         }
@@ -1337,11 +1329,6 @@
         }
     }
 
-    protected boolean isValidService(ComponentName component, int userId) {
-        return componentHasBindPermission(component, userId) && queryPackageForServices(
-                component.getPackageName(), userId).contains(component);
-    }
-
     protected boolean isValidEntry(String packageOrComponent, int userId) {
         return hasMatchingServices(packageOrComponent, userId);
     }
@@ -1499,25 +1486,23 @@
      * Called when user switched to unbind all services from other users.
      */
     @VisibleForTesting
-    void unbindOtherUserServices(int switchedToUser) {
+    void unbindOtherUserServices(int currentUser) {
         TimingsTraceAndSlog t = new TimingsTraceAndSlog();
-        t.traceBegin("ManagedServices.unbindOtherUserServices_current" + switchedToUser);
-        unbindServicesImpl(switchedToUser, true /* allExceptUser */);
+        t.traceBegin("ManagedServices.unbindOtherUserServices_current" + currentUser);
+        unbindServicesImpl(currentUser, true /* allExceptUser */);
         t.traceEnd();
     }
 
-    void unbindUserServices(int removedUser) {
+    void unbindUserServices(int user) {
         TimingsTraceAndSlog t = new TimingsTraceAndSlog();
-        t.traceBegin("ManagedServices.unbindUserServices" + removedUser);
-        unbindServicesImpl(removedUser, false /* allExceptUser */);
+        t.traceBegin("ManagedServices.unbindUserServices" + user);
+        unbindServicesImpl(user, false /* allExceptUser */);
         t.traceEnd();
     }
 
     void unbindServicesImpl(int user, boolean allExceptUser) {
         final SparseArray<Set<ComponentName>> componentsToUnbind = new SparseArray<>();
         synchronized (mMutex) {
-            // Remove enqueued rebinds to avoid rebinding services for a switched user
-            mHandler.removeMessages(ON_BINDING_DIED_REBIND_MSG);
             final Set<ManagedServiceInfo> removableBoundServices = getRemovableConnectedServices();
             for (ManagedServiceInfo info : removableBoundServices) {
                 if ((allExceptUser && (info.userid != user))
@@ -1712,7 +1697,6 @@
                             mServicesRebinding.add(servicesBindingTag);
                             mHandler.postDelayed(() ->
                                     reregisterService(name, userid),
-                                    ON_BINDING_DIED_REBIND_MSG,
                                     ON_BINDING_DIED_REBIND_DELAY_MS);
                         } else {
                             Slog.v(TAG, getCaption() + " not rebinding in user " + userid
diff --git a/services/core/java/com/android/server/power/feature/PowerManagerFlags.java b/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
index fd60e06..db57d11 100644
--- a/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
+++ b/services/core/java/com/android/server/power/feature/PowerManagerFlags.java
@@ -47,6 +47,9 @@
             Flags::perDisplayWakeByTouch
     );
 
+    private final FlagState mFrameworkWakelockInfo =
+            new FlagState(Flags.FLAG_FRAMEWORK_WAKELOCK_INFO, Flags::frameworkWakelockInfo);
+
     /** Returns whether early-screen-timeout-detector is enabled on not. */
     public boolean isEarlyScreenTimeoutDetectorEnabled() {
         return mEarlyScreenTimeoutDetectorFlagState.isEnabled();
@@ -67,6 +70,13 @@
     }
 
     /**
+     * @return Whether FrameworkWakelockInfo atom logging is enabled or not.
+     */
+    public boolean isFrameworkWakelockInfoEnabled() {
+        return mFrameworkWakelockInfo.isEnabled();
+    }
+
+    /**
      * dumps all flagstates
      * @param pw printWriter
      */
@@ -75,6 +85,7 @@
         pw.println(" " + mEarlyScreenTimeoutDetectorFlagState);
         pw.println(" " + mImproveWakelockLatency);
         pw.println(" " + mPerDisplayWakeByTouch);
+        pw.println(" " + mFrameworkWakelockInfo);
     }
 
     private static class FlagState {
diff --git a/services/core/java/com/android/server/power/feature/power_flags.aconfig b/services/core/java/com/android/server/power/feature/power_flags.aconfig
index 9cf3bb6..8bb69ba 100644
--- a/services/core/java/com/android/server/power/feature/power_flags.aconfig
+++ b/services/core/java/com/android/server/power/feature/power_flags.aconfig
@@ -26,3 +26,10 @@
     bug: "343295183"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "framework_wakelock_info"
+    namespace: "power"
+    description: "Feature flag to enable statsd pulling of FrameworkWakelockInfo atoms"
+    bug: "352602149"
+}
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index 936fadf..391071f 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -145,6 +145,7 @@
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
 import com.android.server.LocalServices;
+import com.android.server.power.feature.PowerManagerFlags;
 import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
@@ -5142,6 +5143,10 @@
             mFrameworkStatsLogger.wakelockStateChanged(mapIsolatedUid(uid), wc, name,
                     uidStats.mProcessState, true /* acquired */,
                     getPowerManagerWakeLockLevel(type));
+            if (mPowerManagerFlags.isFrameworkWakelockInfoEnabled()) {
+                mFrameworkEvents.noteStartWakeLock(
+                        mapIsolatedUid(uid), name, getPowerManagerWakeLockLevel(type), uptimeMs);
+            }
         }
     }
 
@@ -5187,6 +5192,10 @@
             mFrameworkStatsLogger.wakelockStateChanged(mapIsolatedUid(uid), wc, name,
                     uidStats.mProcessState, false/* acquired */,
                     getPowerManagerWakeLockLevel(type));
+            if (mPowerManagerFlags.isFrameworkWakelockInfoEnabled()) {
+                mFrameworkEvents.noteStopWakeLock(
+                        mapIsolatedUid(uid), name, getPowerManagerWakeLockLevel(type), uptimeMs);
+            }
 
             if (mappedUid != uid) {
                 // Decrement the ref count for the isolated uid and delete the mapping if uneeded.
@@ -11343,6 +11352,9 @@
         return mTmpCpuTimeInFreq;
     }
 
+    WakelockStatsFrameworkEvents mFrameworkEvents = new WakelockStatsFrameworkEvents();
+    PowerManagerFlags mPowerManagerFlags = new PowerManagerFlags();
+
     public BatteryStatsImpl(@NonNull BatteryStatsConfig config, @NonNull Clock clock,
             @NonNull MonotonicClock monotonicClock, @Nullable File systemDir,
             @NonNull Handler handler, @Nullable PlatformIdleStateCallback platformIdleStateCallback,
@@ -15964,6 +15976,10 @@
                 // Already plugged in. Schedule the long plug in alarm.
                 scheduleNextResetWhilePluggedInCheck();
             }
+
+            if (mPowerManagerFlags.isFrameworkWakelockInfoEnabled()) {
+                mFrameworkEvents.initialize(context);
+            }
         }
     }
 
@@ -16791,7 +16807,8 @@
         }
         int NSORPMS = in.readInt();
         if (NSORPMS > 10000) {
-            throw new ParcelFormatException("File corrupt: too many screen-off rpm stats " + NSORPMS);
+            throw new ParcelFormatException(
+                    "File corrupt: too many screen-off rpm stats " + NSORPMS);
         }
         for (int irpm = 0; irpm < NSORPMS; irpm++) {
             if (in.readInt() != 0) {
diff --git a/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java b/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java
new file mode 100644
index 0000000..c9693bd
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WakelockStatsFrameworkEvents.java
@@ -0,0 +1,328 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.app.StatsManager;
+import android.content.Context;
+import android.os.SystemClock;
+import android.util.Log;
+import android.util.StatsEvent;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.util.ConcurrentUtils;
+import com.android.internal.util.FrameworkStatsLog;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+/** A class to initialise and log metrics pulled by statsd. */
+public class WakelockStatsFrameworkEvents {
+    // statsd has a dimensional limit on the number of different keys it can handle.
+    // Beyond that limit, statsd will drop data.
+    //
+    // When we have seem SUMMARY_THRESHOLD distinct (uid, tag, wakeLockLevel) keys,
+    // we start summarizing new keys as (uid, OVERFLOW_TAG, OVERFLOW_LEVEL) to
+    // reduce the number of keys we pass to statsd.
+    //
+    // When we reach MAX_WAKELOCK_DIMENSIONS distinct keys, we summarize all new keys
+    // as (OVERFLOW_UID, HARD_CAP_TAG, OVERFLOW_LEVEL) to hard cap the number of
+    // distinct keys we pass to statsd.
+    @VisibleForTesting public static final int SUMMARY_THRESHOLD = 500;
+    @VisibleForTesting public static final int MAX_WAKELOCK_DIMENSIONS = 1000;
+
+    @VisibleForTesting public static final int HARD_CAP_UID = -1;
+    @VisibleForTesting public static final String OVERFLOW_TAG = "*overflow*";
+    @VisibleForTesting public static final String HARD_CAP_TAG = "*overflow hard cap*";
+    @VisibleForTesting public static final int OVERFLOW_LEVEL = 1;
+
+    private static class WakeLockKey {
+        private int uid;
+        private String tag;
+        private int powerManagerWakeLockLevel;
+        private int hashCode;
+
+        WakeLockKey(int uid, String tag, int powerManagerWakeLockLevel) {
+            this.uid = uid;
+            this.tag = new String(tag);
+            this.powerManagerWakeLockLevel = powerManagerWakeLockLevel;
+
+            this.hashCode = Objects.hash(uid, tag, powerManagerWakeLockLevel);
+        }
+
+        int getUid() {
+            return uid;
+        }
+
+        String getTag() {
+            return tag;
+        }
+
+        int getPowerManagerWakeLockLevel() {
+            return powerManagerWakeLockLevel;
+        }
+
+        void setOverflow() {
+            tag = OVERFLOW_TAG;
+            powerManagerWakeLockLevel = OVERFLOW_LEVEL;
+            this.hashCode = Objects.hash(uid, tag, powerManagerWakeLockLevel);
+        }
+
+        void setHardCap() {
+            uid = HARD_CAP_UID;
+            tag = HARD_CAP_TAG;
+            powerManagerWakeLockLevel = OVERFLOW_LEVEL;
+            this.hashCode = Objects.hash(uid, tag, powerManagerWakeLockLevel);
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || !(o instanceof WakeLockKey)) return false;
+
+            WakeLockKey that = (WakeLockKey) o;
+            return uid == that.uid
+                    && tag.equals(that.tag)
+                    && powerManagerWakeLockLevel == that.powerManagerWakeLockLevel;
+        }
+
+        @Override
+        public int hashCode() {
+            return this.hashCode;
+        }
+    }
+
+    private static class WakeLockStats {
+        // accumulated uptime attributed to this WakeLock since boot, where overlap
+        // (including nesting) is ignored
+        public long uptimeMillis = 0;
+
+        // count of WakeLocks that have been acquired and then released
+        public long completedCount = 0;
+    }
+
+    private final Object mLock = new Object();
+
+    @GuardedBy("mLock")
+    private final Map<WakeLockKey, WakeLockStats> mWakeLockStats = new HashMap<>();
+
+    private static class WakeLockData {
+        // uptime millis when first acquired
+        public long acquireUptimeMillis = 0;
+        public int refCount = 0;
+
+        WakeLockData(long uptimeMillis) {
+            acquireUptimeMillis = uptimeMillis;
+        }
+    }
+
+    @GuardedBy("mLock")
+    private final Map<WakeLockKey, WakeLockData> mOpenWakeLocks = new HashMap<>();
+
+    public void noteStartWakeLock(
+            int uid, String tag, int powerManagerWakeLockLevel, long eventUptimeMillis) {
+        final WakeLockKey key = new WakeLockKey(uid, tag, powerManagerWakeLockLevel);
+
+        synchronized (mLock) {
+            WakeLockData data =
+                    mOpenWakeLocks.computeIfAbsent(key, k -> new WakeLockData(eventUptimeMillis));
+            data.refCount++;
+            mOpenWakeLocks.put(key, data);
+        }
+    }
+
+    @VisibleForTesting
+    public boolean inOverflow() {
+        synchronized (mLock) {
+            return inOverflowLocked();
+        }
+    }
+
+    @GuardedBy("mLock")
+    private boolean inOverflowLocked() {
+        return mWakeLockStats.size() >= SUMMARY_THRESHOLD;
+    }
+
+    @VisibleForTesting
+    public boolean inHardCap() {
+        synchronized (mLock) {
+            return inHardCapLocked();
+        }
+    }
+
+    @GuardedBy("mLock")
+    private boolean inHardCapLocked() {
+        return mWakeLockStats.size() >= MAX_WAKELOCK_DIMENSIONS;
+    }
+
+    public void noteStopWakeLock(
+            int uid, String tag, int powerManagerWakeLockLevel, long eventUptimeMillis) {
+        WakeLockKey key = new WakeLockKey(uid, tag, powerManagerWakeLockLevel);
+
+        synchronized (mLock) {
+            WakeLockData data = mOpenWakeLocks.get(key);
+            if (data == null) {
+                Log.e(TAG, "WakeLock not found when stopping: " + uid + " " + tag);
+                return;
+            }
+
+            if (data.refCount == 1) {
+                mOpenWakeLocks.remove(key);
+                long wakeLockDur = eventUptimeMillis - data.acquireUptimeMillis;
+
+                // Rewrite key if in an overflow state.
+                if (inOverflowLocked() && !mWakeLockStats.containsKey(key)) {
+                    key.setOverflow();
+                    if (inHardCapLocked() && !mWakeLockStats.containsKey(key)) {
+                        key.setHardCap();
+                    }
+                }
+
+                WakeLockStats stats = mWakeLockStats.computeIfAbsent(key, k -> new WakeLockStats());
+                stats.uptimeMillis += wakeLockDur;
+                stats.completedCount++;
+                mWakeLockStats.put(key, stats);
+            } else {
+                data.refCount--;
+                mOpenWakeLocks.put(key, data);
+            }
+        }
+    }
+
+    public List<StatsEvent> pullFrameworkWakelockInfoAtoms() {
+        return pullFrameworkWakelockInfoAtoms(SystemClock.uptimeMillis());
+    }
+
+    public List<StatsEvent> pullFrameworkWakelockInfoAtoms(long nowMillis) {
+        List<StatsEvent> result = new ArrayList<>();
+        HashSet<WakeLockKey> keys = new HashSet<>();
+
+        // Used to collect open WakeLocks when in an overflow state.
+        HashMap<WakeLockKey, WakeLockStats> openOverflowStats = new HashMap<>();
+
+        synchronized (mLock) {
+            keys.addAll(mWakeLockStats.keySet());
+
+            // If we are in an overflow state, an open wakelock may have a new key
+            // that needs to be summarized.
+            if (inOverflowLocked()) {
+                for (WakeLockKey key : mOpenWakeLocks.keySet()) {
+                    if (!mWakeLockStats.containsKey(key)) {
+                        WakeLockData data = mOpenWakeLocks.get(key);
+
+                        key.setOverflow();
+                        if (inHardCapLocked() && !mWakeLockStats.containsKey(key)) {
+                            key.setHardCap();
+                        }
+                        keys.add(key);
+
+                        WakeLockStats stats =
+                                openOverflowStats.computeIfAbsent(key, k -> new WakeLockStats());
+                        stats.uptimeMillis += nowMillis - data.acquireUptimeMillis;
+                        openOverflowStats.put(key, stats);
+                    }
+                }
+            } else {
+                keys.addAll(mOpenWakeLocks.keySet());
+            }
+
+            for (WakeLockKey key : keys) {
+                long openWakeLockUptime = 0;
+                WakeLockData data = mOpenWakeLocks.get(key);
+                if (data != null) {
+                    openWakeLockUptime = nowMillis - data.acquireUptimeMillis;
+                }
+
+                WakeLockStats stats = mWakeLockStats.computeIfAbsent(key, k -> new WakeLockStats());
+                WakeLockStats extraTime =
+                        openOverflowStats.computeIfAbsent(key, k -> new WakeLockStats());
+
+                stats.uptimeMillis += openWakeLockUptime + extraTime.uptimeMillis;
+
+                StatsEvent event =
+                        StatsEvent.newBuilder()
+                                .setAtomId(FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO)
+                                .writeInt(key.getUid())
+                                .writeString(key.getTag())
+                                .writeInt(key.getPowerManagerWakeLockLevel())
+                                .writeLong(stats.uptimeMillis)
+                                .writeLong(stats.completedCount)
+                                .build();
+                result.add(event);
+            }
+        }
+
+        return result;
+    }
+
+    private static final String TAG = "BatteryStatsPulledMetrics";
+
+    private final StatsPullCallbackHandler mStatsPullCallbackHandler =
+            new StatsPullCallbackHandler();
+
+    private boolean mIsInitialized = false;
+
+    public void initialize(Context context) {
+        if (mIsInitialized) {
+            return;
+        }
+
+        final StatsManager statsManager = context.getSystemService(StatsManager.class);
+        if (statsManager == null) {
+            Log.e(
+                    TAG,
+                    "Error retrieving StatsManager. Cannot initialize BatteryStatsPulledMetrics.");
+        } else {
+            Log.d(TAG, "Registering callback with StatsManager");
+
+            // DIRECT_EXECUTOR means that callback will run on binder thread.
+            statsManager.setPullAtomCallback(
+                    FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO,
+                    null /* metadata */,
+                    ConcurrentUtils.DIRECT_EXECUTOR,
+                    mStatsPullCallbackHandler);
+            mIsInitialized = true;
+        }
+    }
+
+    private class StatsPullCallbackHandler implements StatsManager.StatsPullAtomCallback {
+        @Override
+        public int onPullAtom(int atomTag, List<StatsEvent> data) {
+            // handle the tags appropriately.
+            List<StatsEvent> events = pullEvents(atomTag);
+            if (events == null) {
+                return StatsManager.PULL_SKIP;
+            }
+
+            data.addAll(events);
+            return StatsManager.PULL_SUCCESS;
+        }
+
+        private List<StatsEvent> pullEvents(int atomTag) {
+            switch (atomTag) {
+                case FrameworkStatsLog.FRAMEWORK_WAKELOCK_INFO:
+                    return pullFrameworkWakelockInfoAtoms();
+                default:
+                    return null;
+            }
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
index 74c1124e..e7735d8 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -63,6 +63,7 @@
 import static com.android.internal.util.FrameworkStatsLog.TIME_ZONE_DETECTOR_STATE__DETECTION_MODE__TELEPHONY;
 import static com.android.internal.util.FrameworkStatsLog.TIME_ZONE_DETECTOR_STATE__DETECTION_MODE__UNKNOWN;
 import static com.android.server.am.MemoryStatUtil.readMemoryStatFromFilesystem;
+import static com.android.server.stats.Flags.accumulateNetworkStatsSinceBoot;
 import static com.android.server.stats.Flags.addMobileBytesTransferByProcStatePuller;
 import static com.android.server.stats.Flags.applyNetworkStatsPollRateLimit;
 import static com.android.server.stats.pull.IonMemoryUtil.readProcessSystemIonHeapSizesFromDebugfs;
@@ -229,6 +230,7 @@
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 import com.android.server.stats.pull.IonMemoryUtil.IonAllocations;
 import com.android.server.stats.pull.ProcfsMemoryUtil.MemorySnapshot;
+import com.android.server.stats.pull.netstats.NetworkStatsAccumulator;
 import com.android.server.stats.pull.netstats.NetworkStatsExt;
 import com.android.server.stats.pull.netstats.SubInfo;
 import com.android.server.storage.DiskStatsFileLogger;
@@ -424,6 +426,14 @@
     @GuardedBy("mDataBytesTransferLock")
     private final ArrayList<NetworkStatsExt> mNetworkStatsBaselines = new ArrayList<>();
 
+    // Accumulates NetworkStats from initialization till the present moment.
+    // It is necessary to accumulate stats internally, because NetworkStats persists data for a
+    // limited amount of time, after which diff becomes incorrect without accumulation.
+    @NonNull
+    @GuardedBy("mDataBytesTransferLock")
+    private final ArrayList<NetworkStatsAccumulator> mNetworkStatsAccumulators =
+            new ArrayList<>();
+
     @GuardedBy("mDataBytesTransferLock")
     private long mLastNetworkStatsPollTime = -NETSTATS_POLL_RATE_LIMIT_MS;
 
@@ -1265,49 +1275,39 @@
             case FrameworkStatsLog.WIFI_BYTES_TRANSFER: {
                 final NetworkStats stats = getUidNetworkStatsSnapshotForTransportLocked(
                         TRANSPORT_WIFI);
-                if (stats != null) {
-                    ret.add(new NetworkStatsExt(sliceNetworkStatsByUid(stats),
-                            new int[]{TRANSPORT_WIFI}, /*slicedByFgbg=*/false));
-                }
+                ret.add(new NetworkStatsExt(sliceNetworkStatsByUid(stats),
+                        new int[]{TRANSPORT_WIFI}, /*slicedByFgbg=*/false));
                 break;
             }
             case FrameworkStatsLog.WIFI_BYTES_TRANSFER_BY_FG_BG: {
                 final NetworkStats stats = getUidNetworkStatsSnapshotForTransportLocked(
                         TRANSPORT_WIFI);
-                if (stats != null) {
-                    ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
-                            new int[]{TRANSPORT_WIFI}, /*slicedByFgbg=*/true));
-                }
+                ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
+                        new int[]{TRANSPORT_WIFI}, /*slicedByFgbg=*/true));
                 break;
             }
             case FrameworkStatsLog.MOBILE_BYTES_TRANSFER: {
                 final NetworkStats stats =
                         getUidNetworkStatsSnapshotForTransportLocked(TRANSPORT_CELLULAR);
-                if (stats != null) {
-                    ret.add(new NetworkStatsExt(sliceNetworkStatsByUid(stats),
-                            new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/false));
-                }
+                ret.add(new NetworkStatsExt(sliceNetworkStatsByUid(stats),
+                        new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/false));
                 break;
             }
             case FrameworkStatsLog.MOBILE_BYTES_TRANSFER_BY_FG_BG: {
                 final NetworkStats stats =
                         getUidNetworkStatsSnapshotForTransportLocked(TRANSPORT_CELLULAR);
-                if (stats != null) {
-                    ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
-                            new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/true));
-                }
+                ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
+                        new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/true));
                 break;
             }
             case FrameworkStatsLog.PROXY_BYTES_TRANSFER_BY_FG_BG: {
                 final NetworkStats stats = getUidNetworkStatsSnapshotForTemplateLocked(
                         new NetworkTemplate.Builder(MATCH_PROXY).build(),  /*includeTags=*/false);
-                if (stats != null) {
-                    ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
-                            new int[]{TRANSPORT_BLUETOOTH},
-                            /*slicedByFgbg=*/true, /*slicedByTag=*/false,
-                            /*slicedByMetered=*/false, TelephonyManager.NETWORK_TYPE_UNKNOWN,
-                            /*subInfo=*/null, OEM_MANAGED_ALL, /*isTypeProxy=*/true));
-                }
+                ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
+                        new int[]{TRANSPORT_BLUETOOTH},
+                        /*slicedByFgbg=*/true, /*slicedByTag=*/false,
+                        /*slicedByMetered=*/false, TelephonyManager.NETWORK_TYPE_UNKNOWN,
+                        /*subInfo=*/null, OEM_MANAGED_ALL, /*isTypeProxy=*/true));
                 break;
             }
             case FrameworkStatsLog.BYTES_TRANSFER_BY_TAG_AND_METERED: {
@@ -1316,14 +1316,12 @@
                 final NetworkStats cellularStats = getUidNetworkStatsSnapshotForTemplateLocked(
                         new NetworkTemplate.Builder(MATCH_MOBILE)
                                 .setMeteredness(METERED_YES).build(), /*includeTags=*/true);
-                if (wifiStats != null && cellularStats != null) {
-                    final NetworkStats stats = wifiStats.add(cellularStats);
-                    ret.add(new NetworkStatsExt(sliceNetworkStatsByUidTagAndMetered(stats),
-                            new int[]{TRANSPORT_WIFI, TRANSPORT_CELLULAR},
-                            /*slicedByFgbg=*/false, /*slicedByTag=*/true,
-                            /*slicedByMetered=*/true, TelephonyManager.NETWORK_TYPE_UNKNOWN,
-                            /*subInfo=*/null, OEM_MANAGED_ALL, /*isTypeProxy=*/false));
-                }
+                final NetworkStats stats = wifiStats.add(cellularStats);
+                ret.add(new NetworkStatsExt(sliceNetworkStatsByUidTagAndMetered(stats),
+                        new int[]{TRANSPORT_WIFI, TRANSPORT_CELLULAR},
+                        /*slicedByFgbg=*/false, /*slicedByTag=*/true,
+                        /*slicedByMetered=*/true, TelephonyManager.NETWORK_TYPE_UNKNOWN,
+                        /*subInfo=*/null, OEM_MANAGED_ALL, /*isTypeProxy=*/false));
                 break;
             }
             case FrameworkStatsLog.DATA_USAGE_BYTES_TRANSFER: {
@@ -1509,12 +1507,10 @@
                 final NetworkStats stats = getUidNetworkStatsSnapshotForTemplateLocked(
                         template, false);
                 final Integer transport = ruleAndTransport.second;
-                if (stats != null) {
-                    ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
-                            new int[]{transport}, /*slicedByFgbg=*/true, /*slicedByTag=*/false,
-                            /*slicedByMetered=*/false, TelephonyManager.NETWORK_TYPE_UNKNOWN,
-                            /*subInfo=*/null, oemManaged, /*isTypeProxy=*/false));
-                }
+                ret.add(new NetworkStatsExt(sliceNetworkStatsByUidAndFgbg(stats),
+                        new int[]{transport}, /*slicedByFgbg=*/true, /*slicedByTag=*/false,
+                        /*slicedByMetered=*/false, TelephonyManager.NETWORK_TYPE_UNKNOWN,
+                        /*subInfo=*/null, oemManaged, /*isTypeProxy=*/false));
             }
         }
 
@@ -1525,7 +1521,7 @@
      * Create a snapshot of NetworkStats for a given transport.
      */
     @GuardedBy("mDataBytesTransferLock")
-    @Nullable
+    @NonNull
     private NetworkStats getUidNetworkStatsSnapshotForTransportLocked(int transport) {
         NetworkTemplate template = null;
         switch (transport) {
@@ -1564,20 +1560,49 @@
      * some traffic before boot.
      */
     @GuardedBy("mDataBytesTransferLock")
-    @Nullable
+    @NonNull
     private NetworkStats getUidNetworkStatsSnapshotForTemplateLocked(
             @NonNull NetworkTemplate template, boolean includeTags) {
         final long elapsedMillisSinceBoot = SystemClock.elapsedRealtime();
-        final long currentTimeInMillis = MICROSECONDS.toMillis(SystemClock.currentTimeMicro());
-        final long bucketDuration = Settings.Global.getLong(mContext.getContentResolver(),
+        final long currentTimeMillis = MICROSECONDS.toMillis(SystemClock.currentTimeMicro());
+        final long bootTimeMillis = currentTimeMillis - elapsedMillisSinceBoot;
+        final long bucketDurationMillis = Settings.Global.getLong(mContext.getContentResolver(),
                 NETSTATS_UID_BUCKET_DURATION, NETSTATS_UID_DEFAULT_BUCKET_DURATION_MS);
 
-        // Set startTime before boot so that NetworkStats includes at least one full bucket.
-        // Set endTime in the future so that NetworkStats includes everything in the active bucket.
-        final long startTime = currentTimeInMillis - elapsedMillisSinceBoot - bucketDuration;
-        final long endTime = currentTimeInMillis + bucketDuration;
+        if (accumulateNetworkStatsSinceBoot()) {
+            NetworkStatsAccumulator accumulator = CollectionUtils.find(
+                    mNetworkStatsAccumulators, it -> it.hasEqualParameters(template, includeTags));
+            if (accumulator == null) {
+                accumulator = new NetworkStatsAccumulator(
+                        template,
+                        includeTags,
+                        bucketDurationMillis,
+                        bootTimeMillis - bucketDurationMillis);
+                mNetworkStatsAccumulators.add(accumulator);
+            }
 
-        // NetworkStatsManager#forceUpdate updates stats for all networks
+            return accumulator.queryStats(currentTimeMillis,
+                    (aTemplate, aIncludeTags, aStartTime, aEndTime) -> {
+                        synchronized (mDataBytesTransferLock) {
+                            return getUidNetworkStatsSnapshotForTemplateLocked(aTemplate,
+                                    aIncludeTags, aStartTime, aEndTime);
+                        }
+                    });
+
+        } else {
+            // Set end time in the future to include all stats in the active bucket.
+            return getUidNetworkStatsSnapshotForTemplateLocked(
+                    template, includeTags,
+                    bootTimeMillis - bucketDurationMillis,
+                    currentTimeMillis + bucketDurationMillis);
+        }
+    }
+
+    @GuardedBy("mDataBytesTransferLock")
+    @NonNull
+    private NetworkStats getUidNetworkStatsSnapshotForTemplateLocked(
+            @NonNull NetworkTemplate template, boolean includeTags, long startTime, long endTime) {
+        final long elapsedMillisSinceBoot = SystemClock.elapsedRealtime();
         if (applyNetworkStatsPollRateLimit()) {
             // The new way: rate-limit force-polling for all NetworkStats queries
             if (elapsedMillisSinceBoot - mLastNetworkStatsPollTime >= NETSTATS_POLL_RATE_LIMIT_MS) {
@@ -1621,12 +1646,10 @@
                             .setMeteredness(METERED_YES).build();
             final NetworkStats stats =
                     getUidNetworkStatsSnapshotForTemplateLocked(template, /*includeTags=*/false);
-            if (stats != null) {
-                ret.add(new NetworkStatsExt(sliceNetworkStatsByFgbg(stats),
-                        new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/true,
-                        /*slicedByTag=*/false, /*slicedByMetered=*/false, ratType, subInfo,
-                        OEM_MANAGED_ALL, /*isTypeProxy=*/false));
-            }
+            ret.add(new NetworkStatsExt(sliceNetworkStatsByFgbg(stats),
+                    new int[]{TRANSPORT_CELLULAR}, /*slicedByFgbg=*/true,
+                    /*slicedByTag=*/false, /*slicedByMetered=*/false, ratType, subInfo,
+                    OEM_MANAGED_ALL, /*isTypeProxy=*/false));
         }
         return ret;
     }
diff --git a/services/core/java/com/android/server/stats/pull/netstats/NetworkStatsAccumulator.java b/services/core/java/com/android/server/stats/pull/netstats/NetworkStatsAccumulator.java
new file mode 100644
index 0000000..e798bc4
--- /dev/null
+++ b/services/core/java/com/android/server/stats/pull/netstats/NetworkStatsAccumulator.java
@@ -0,0 +1,108 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.stats.pull.netstats;
+
+import android.annotation.NonNull;
+import android.net.NetworkStats;
+import android.net.NetworkTemplate;
+
+import java.util.Objects;
+
+/**
+ * A class that queries NetworkStats, accumulates query results, and exposes cumulative stats across
+ * the time range covered by the queries. This class is not thread-safe.
+ * <p>
+ * This is class should be used when querying NetworkStats since boot, as NetworkStats persists data
+ * only for a limited period of time.
+ *
+ * @hide
+ */
+public class NetworkStatsAccumulator {
+
+    private final NetworkTemplate mTemplate;
+    private final boolean mWithTags;
+    private final long mBucketDurationMillis;
+    private NetworkStats mSnapshot;
+    private long mSnapshotEndTimeMillis;
+
+    public NetworkStatsAccumulator(@NonNull NetworkTemplate template, boolean withTags,
+            long bucketDurationMillis, long snapshotEndTimeMillis) {
+        mTemplate = template;
+        mWithTags = withTags;
+        mBucketDurationMillis = bucketDurationMillis;
+        mSnapshot = new NetworkStats(0, 1);
+        mSnapshotEndTimeMillis = snapshotEndTimeMillis;
+    }
+
+    /**
+     * Provides cumulative NetworkStats until given timestamp.
+     * <p>
+     * This method method may call {@code queryFunction} more than once, which includes maintaining
+     * an internal cumulative stats snapshot and querying stats after the snapshot.
+     */
+    @NonNull
+    public NetworkStats queryStats(long currentTimeMillis,
+            @NonNull StatsQueryFunction queryFunction) {
+        maybeExpandSnapshot(currentTimeMillis, queryFunction);
+        return snapshotPlusFollowingStats(currentTimeMillis, queryFunction);
+    }
+
+    /**
+     * Returns true if the accumulator is using given query parameters.
+     */
+    public boolean hasEqualParameters(@NonNull NetworkTemplate template, boolean withTags) {
+        return Objects.equals(mTemplate, template) && mWithTags == withTags;
+    }
+
+    /**
+     * Expands the internal cumulative stats snapshot, if possible, by querying NetworkStats.
+     */
+    private void maybeExpandSnapshot(long currentTimeMillis,
+            @NonNull StatsQueryFunction queryFunction) {
+        // Update snapshot only if it is possible to expand it by at least one full bucket, and only
+        // if the new snapshot's end is not in the active bucket.
+        long newEndTimeMillis = currentTimeMillis - mBucketDurationMillis;
+        if (newEndTimeMillis - mSnapshotEndTimeMillis > mBucketDurationMillis) {
+            NetworkStats extraStats = queryFunction.queryNetworkStats(mTemplate, mWithTags,
+                    mSnapshotEndTimeMillis, newEndTimeMillis);
+            mSnapshot = mSnapshot.add(extraStats);
+            mSnapshotEndTimeMillis = newEndTimeMillis;
+        }
+    }
+
+    /**
+     * Adds up stats in the internal cumulative snapshot and the stats that follow after it.
+     */
+    @NonNull
+    private NetworkStats snapshotPlusFollowingStats(long currentTimeMillis,
+            @NonNull StatsQueryFunction queryFunction) {
+        // Set end time in the future to include all stats in the active bucket.
+        NetworkStats extraStats = queryFunction.queryNetworkStats(mTemplate, mWithTags,
+                mSnapshotEndTimeMillis, currentTimeMillis + mBucketDurationMillis);
+        return mSnapshot.add(extraStats);
+    }
+
+    @FunctionalInterface
+    public interface StatsQueryFunction {
+        /**
+         * Returns network stats during the given time period.
+         */
+        @NonNull
+        NetworkStats queryNetworkStats(@NonNull NetworkTemplate template, boolean includeTags,
+                long startTime, long endTime);
+    }
+}
diff --git a/services/core/java/com/android/server/stats/stats_flags.aconfig b/services/core/java/com/android/server/stats/stats_flags.aconfig
index afea303..8686458f 100644
--- a/services/core/java/com/android/server/stats/stats_flags.aconfig
+++ b/services/core/java/com/android/server/stats/stats_flags.aconfig
@@ -30,3 +30,11 @@
     bug: "352495181"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "accumulate_network_stats_since_boot"
+    namespace: "statsd"
+    description: "Accumulate results of NetworkStats queries to avoid hitting NetworkStats persistence limit"
+    bug: "352537247"
+    is_fixed_read_only: true
+}
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 31f4d08..c7fafe9 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -3198,6 +3198,9 @@
         if (!compatEnabled && !mWmService.mConstants.mIgnoreActivityOrientationRequest) {
             return false;
         }
+        if (mWmService.mConstants.isPackageOptOutIgnoreActivityOrientationRequest(packageName)) {
+            return false;
+        }
         // If the user preference respects aspect ratio, then it becomes non-resizable.
         return !mAppCompatController.getAppCompatOverrides().getAppCompatAspectRatioOverrides()
                 .shouldApplyUserMinAspectRatioOverride();
@@ -5877,6 +5880,7 @@
             return;
         }
 
+        final State prevState = mState;
         mState = state;
 
         if (getTaskFragment() != null) {
@@ -5917,6 +5921,14 @@
                 mAtmService.updateBatteryStats(this, false);
                 mAtmService.updateActivityUsageStats(this, Event.ACTIVITY_PAUSED);
                 break;
+            case STOPPING:
+                // It is possible that an Activity is scheduled to be STOPPED directly from RESUMED
+                // state. Updating the PAUSED usage state in that case, since the Activity will be
+                // STOPPED while cycled through the PAUSED state.
+                if (prevState == RESUMED) {
+                    mAtmService.updateActivityUsageStats(this, Event.ACTIVITY_PAUSED);
+                }
+                break;
             case STOPPED:
                 mAtmService.updateActivityUsageStats(this, Event.ACTIVITY_STOPPED);
                 if (mDisplayContent != null) {
diff --git a/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java b/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java
index 1d00136..8c5689c1 100644
--- a/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java
+++ b/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java
@@ -22,6 +22,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.app.CameraCompatTaskInfo;
 import android.content.pm.ActivityInfo.ScreenOrientation;
 import android.content.res.Configuration;
 import android.widget.Toast;
@@ -250,6 +251,14 @@
                 cameraCompatFreeformPolicyAspectRatio);
     }
 
+    @CameraCompatTaskInfo.FreeformCameraCompatMode
+    static int getCameraCompatFreeformMode(@NonNull ActivityRecord activity) {
+        final AppCompatCameraPolicy cameraPolicy = getAppCompatCameraPolicy(activity);
+        return cameraPolicy != null && cameraPolicy.mCameraCompatFreeformPolicy != null
+                ? cameraPolicy.mCameraCompatFreeformPolicy.getCameraCompatMode(activity)
+                : CameraCompatTaskInfo.CAMERA_COMPAT_FREEFORM_NONE;
+    }
+
     /**
      * Whether we should apply the min aspect ratio per-app override only when an app is connected
      * to the camera.
diff --git a/services/core/java/com/android/server/wm/AppCompatUtils.java b/services/core/java/com/android/server/wm/AppCompatUtils.java
index 69421d0..8d84248 100644
--- a/services/core/java/com/android/server/wm/AppCompatUtils.java
+++ b/services/core/java/com/android/server/wm/AppCompatUtils.java
@@ -185,8 +185,8 @@
                         && aspectRatioOverrides.shouldEnableUserAspectRatioSettings();
         appCompatTaskInfo.setEligibleForUserAspectRatioButton(eligibleForAspectRatioButton);
         appCompatTaskInfo.setTopActivityLetterboxed(top.areBoundsLetterboxed());
-        appCompatTaskInfo.cameraCompatTaskInfo.freeformCameraCompatMode = top.mAppCompatController
-                .getAppCompatCameraOverrides().getFreeformCameraCompatMode();
+        appCompatTaskInfo.cameraCompatTaskInfo.freeformCameraCompatMode =
+                AppCompatCameraPolicy.getCameraCompatFreeformMode(top);
         appCompatTaskInfo.setHasMinAspectRatioOverride(top.mAppCompatController
                 .getDesktopAppCompatAspectRatioPolicy().hasMinAspectRatioOverride(task));
     }
diff --git a/services/core/java/com/android/server/wm/CameraCompatFreeformPolicy.java b/services/core/java/com/android/server/wm/CameraCompatFreeformPolicy.java
index 4f0cbf9..2a0252a 100644
--- a/services/core/java/com/android/server/wm/CameraCompatFreeformPolicy.java
+++ b/services/core/java/com/android/server/wm/CameraCompatFreeformPolicy.java
@@ -21,6 +21,8 @@
 import static android.app.CameraCompatTaskInfo.CAMERA_COMPAT_FREEFORM_NONE;
 import static android.app.CameraCompatTaskInfo.CAMERA_COMPAT_FREEFORM_PORTRAIT_DEVICE_IN_LANDSCAPE;
 import static android.app.CameraCompatTaskInfo.CAMERA_COMPAT_FREEFORM_PORTRAIT_DEVICE_IN_PORTRAIT;
+import static android.app.WindowConfiguration.WINDOW_CONFIG_APP_BOUNDS;
+import static android.app.WindowConfiguration.WINDOW_CONFIG_DISPLAY_ROTATION;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_LOCKED;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_NOSENSOR;
 import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
@@ -35,6 +37,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.app.CameraCompatTaskInfo;
 import android.content.pm.ActivityInfo;
 import android.content.res.Configuration;
 import android.view.DisplayInfo;
@@ -64,8 +67,6 @@
     @NonNull
     private final CameraStateMonitor mCameraStateMonitor;
 
-    private boolean mIsCameraCompatTreatmentPending = false;
-
     @Nullable
     private Task mCameraTask;
 
@@ -100,13 +101,27 @@
         return mIsRunning;
     }
 
-    // Refreshing only when configuration changes after rotation or camera split screen aspect ratio
-    // treatment is enabled.
+    // Refreshing only when configuration changes after applying camera compat treatment.
     @Override
     public boolean shouldRefreshActivity(@NonNull ActivityRecord activity,
             @NonNull Configuration newConfig,
             @NonNull Configuration lastReportedConfig) {
-        return isTreatmentEnabledForActivity(activity) && mIsCameraCompatTreatmentPending;
+        return isTreatmentEnabledForActivity(activity, /* shouldCheckOrientation= */ true)
+                && haveCameraCompatAttributesChanged(newConfig, lastReportedConfig);
+    }
+
+    private boolean haveCameraCompatAttributesChanged(@NonNull Configuration newConfig,
+            @NonNull Configuration lastReportedConfig) {
+        // Camera compat treatment changes the following:
+        // - Letterboxes app bounds to camera compat aspect ratio in app's requested orientation,
+        // - Changes display rotation so it matches what the app expects in its chosen orientation,
+        // - Rotate-and-crop camera feed to match that orientation (this changes iff the display
+        //     rotation changes, so no need to check).
+        final long diff = newConfig.windowConfiguration.diff(lastReportedConfig.windowConfiguration,
+                /* compareUndefined= */ true);
+        final boolean appBoundsChanged = (diff & WINDOW_CONFIG_APP_BOUNDS) != 0;
+        final boolean displayRotationChanged = (diff & WINDOW_CONFIG_DISPLAY_ROTATION) != 0;
+        return appBoundsChanged || displayRotationChanged;
     }
 
     /**
@@ -132,22 +147,26 @@
     @Override
     public void onCameraOpened(@NonNull ActivityRecord cameraActivity,
             @NonNull String cameraId) {
-        if (!isTreatmentEnabledForActivity(cameraActivity)) {
+        // Do not check orientation outside of the config recompute, as the app's orientation intent
+        // might be obscured by a fullscreen override. Especially for apps which have a camera
+        // functionality which is not the main focus of the app: while most of the app might work
+        // well in fullscreen, often the camera setup still assumes it will run on a portrait device
+        // in its natural orientation and comes out stretched or sideways.
+        // Config recalculation will later check the original orientation to avoid applying
+        // treatment to apps optimized for large screens.
+        if (!isTreatmentEnabledForActivity(cameraActivity, /* shouldCheckOrientation= */ false)) {
             return;
         }
-        final int existingCameraCompatMode = cameraActivity.mAppCompatController
-                .getAppCompatCameraOverrides()
-                        .getFreeformCameraCompatMode();
-        final int newCameraCompatMode = getCameraCompatMode(cameraActivity);
-        if (newCameraCompatMode != existingCameraCompatMode) {
-            mIsCameraCompatTreatmentPending = true;
-            mCameraTask = cameraActivity.getTask();
-            cameraActivity.mAppCompatController.getAppCompatCameraOverrides()
-                    .setFreeformCameraCompatMode(newCameraCompatMode);
-            forceUpdateActivityAndTask(cameraActivity);
-        } else {
-            mIsCameraCompatTreatmentPending = false;
-        }
+
+        cameraActivity.recomputeConfiguration();
+        updateCameraCompatMode(cameraActivity);
+        cameraActivity.getTask().dispatchTaskInfoChangedIfNeeded(/* force= */ true);
+        cameraActivity.ensureActivityConfiguration(/* ignoreVisibility= */ false);
+    }
+
+    private void updateCameraCompatMode(@NonNull ActivityRecord cameraActivity) {
+        cameraActivity.mAppCompatController.getAppCompatCameraOverrides()
+                .setFreeformCameraCompatMode(getCameraCompatMode(cameraActivity));
     }
 
     @Override
@@ -167,7 +186,6 @@
             }
         }
         mCameraTask = null;
-        mIsCameraCompatTreatmentPending = false;
         return true;
     }
 
@@ -184,13 +202,12 @@
         // Camera compat should direct aspect ratio when in camera compat mode, unless an app has a
         // different camera compat aspect ratio set: this allows per-app camera compat override
         // aspect ratio to be smaller than the default.
-        return isInCameraCompatMode(activity) && !activity.mAppCompatController
+        return isInFreeformCameraCompatMode(activity) && !activity.mAppCompatController
                 .getAppCompatCameraOverrides().isOverrideMinAspectRatioForCameraEnabled();
     }
 
-    private boolean isInCameraCompatMode(@NonNull ActivityRecord activity) {
-        return activity.mAppCompatController.getAppCompatCameraOverrides()
-                .getFreeformCameraCompatMode() != CAMERA_COMPAT_FREEFORM_NONE;
+    boolean isInFreeformCameraCompatMode(@NonNull ActivityRecord activity) {
+        return getCameraCompatMode(activity) != CAMERA_COMPAT_FREEFORM_NONE;
     }
 
     float getCameraCompatAspectRatio(@NonNull ActivityRecord activityRecord) {
@@ -201,16 +218,11 @@
         return MIN_FIXED_ORIENTATION_LETTERBOX_ASPECT_RATIO;
     }
 
-    private void forceUpdateActivityAndTask(ActivityRecord cameraActivity) {
-        cameraActivity.recomputeConfiguration();
-        cameraActivity.updateReportedConfigurationAndSend();
-        Task cameraTask = cameraActivity.getTask();
-        if (cameraTask != null) {
-            cameraTask.dispatchTaskInfoChangedIfNeeded(/* force= */ true);
+    @CameraCompatTaskInfo.FreeformCameraCompatMode
+    int getCameraCompatMode(@NonNull ActivityRecord topActivity) {
+        if (!isTreatmentEnabledForActivity(topActivity, /* shouldCheckOrientation= */ true)) {
+            return CAMERA_COMPAT_FREEFORM_NONE;
         }
-    }
-
-    private static int getCameraCompatMode(@NonNull ActivityRecord topActivity) {
         final int appOrientation = topActivity.getRequestedConfigurationOrientation();
         // It is very important to check the original (actual) display rotation, and not the
         // sandboxed rotation that camera compat treatment sets.
@@ -250,15 +262,24 @@
      * <ul>
      *     <li>Treatment is enabled.
      *     <li>Camera is active for the package.
-     *     <li>The app has a fixed orientation.
+     *     <li>The app has a fixed orientation if {@code checkOrientation} is true.
      *     <li>The app is in freeform windowing mode.
      * </ul>
+     *
+     * @param checkOrientation Whether to take apps orientation into account for this check. Only
+     *                         fixed-orientation apps should be targeted, but this might be
+     *                         obscured by OEMs via fullscreen override and the app's original
+     *                         intent inaccessible when the camera opens. Thus, policy would pass
+     *                         {@code false} here when considering whether to trigger config
+     *                         recalculation, and later pass {@code true} during recalculation.
      */
-    private boolean isTreatmentEnabledForActivity(@NonNull ActivityRecord activity) {
+    @VisibleForTesting
+    boolean isTreatmentEnabledForActivity(@NonNull ActivityRecord activity,
+            boolean checkOrientation) {
         int orientation = activity.getRequestedConfigurationOrientation();
         return isCameraCompatForFreeformEnabledForActivity(activity)
                 && mCameraStateMonitor.isCameraRunningForActivity(activity)
-                && orientation != ORIENTATION_UNDEFINED
+                && (!checkOrientation || orientation != ORIENTATION_UNDEFINED)
                 && activity.inFreeformWindowingMode()
                 // "locked" and "nosensor" values are often used by camera apps that can't
                 // handle dynamic changes so we shouldn't force-letterbox them.
@@ -270,7 +291,7 @@
 
     private boolean isActivityForCameraIdRefreshing(@NonNull ActivityRecord topActivity,
             @NonNull String cameraId) {
-        if (!isTreatmentEnabledForActivity(topActivity)
+        if (!isTreatmentEnabledForActivity(topActivity, /* checkOrientation= */ true)
                 || mCameraStateMonitor.isCameraWithIdRunningForActivity(topActivity, cameraId)) {
             return false;
         }
diff --git a/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java b/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java
index 1a8f5fc..fcf88d3 100644
--- a/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java
+++ b/services/core/java/com/android/server/wm/DesktopModeBoundsCalculator.java
@@ -36,7 +36,7 @@
 import android.os.SystemProperties;
 import android.util.Size;
 import android.view.Gravity;
-import android.window.flags.DesktopModeFlags;
+import android.window.DesktopModeFlags;
 
 import java.util.function.Consumer;
 
diff --git a/services/core/java/com/android/server/wm/DesktopModeHelper.java b/services/core/java/com/android/server/wm/DesktopModeHelper.java
index b5ea0bd..6bf1c46 100644
--- a/services/core/java/com/android/server/wm/DesktopModeHelper.java
+++ b/services/core/java/com/android/server/wm/DesktopModeHelper.java
@@ -19,7 +19,7 @@
 import android.annotation.NonNull;
 import android.content.Context;
 import android.os.SystemProperties;
-import android.window.flags.DesktopModeFlags;
+import android.window.DesktopModeFlags;
 
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
diff --git a/services/core/java/com/android/server/wm/InputManagerCallback.java b/services/core/java/com/android/server/wm/InputManagerCallback.java
index dcf0319..a4fe064 100644
--- a/services/core/java/com/android/server/wm/InputManagerCallback.java
+++ b/services/core/java/com/android/server/wm/InputManagerCallback.java
@@ -16,6 +16,7 @@
 package com.android.server.wm;
 
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
+import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_INPUT;
@@ -28,9 +29,12 @@
 import android.gui.StalledTransactionInfo;
 import android.os.Debug;
 import android.os.IBinder;
+import android.os.Trace;
 import android.util.Slog;
+import android.util.SparseIntArray;
 import android.view.Display;
 import android.view.InputApplicationHandle;
+import android.view.InputDevice;
 import android.view.KeyEvent;
 import android.view.SurfaceControl;
 import android.view.WindowManager;
@@ -64,6 +68,12 @@
     // which point the ActivityManager will enable dispatching.
     private boolean mInputDispatchEnabled;
 
+    /**
+     * The last input devices info which may affect display configuration. This is a quick lookup
+     * to detect interested changes without entering WM lock.
+     */
+    private SparseIntArray mLastInputConfigurationSources;
+
     public InputManagerCallback(WindowManagerService service) {
         mService = service;
     }
@@ -117,8 +127,16 @@
     /** Notifies that the input device configuration has changed. */
     @Override
     public void notifyConfigurationChanged() {
-        synchronized (mService.mGlobalLock) {
-            mService.mRoot.forAllDisplays(DisplayContent::sendNewConfiguration);
+        Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "notifyConfigurationChanged");
+        final boolean changed = !com.android.window.flags.Flags.filterIrrelevantInputDeviceChange()
+                || updateLastInputConfigurationSources();
+
+        if (changed) {
+            synchronized (mService.mGlobalLock) {
+                Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "inputDeviceConfigChanged");
+                mService.mRoot.forAllDisplays(DisplayContent::sendNewConfiguration);
+                Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+            }
         }
 
         synchronized (mInputDevicesReadyMonitor) {
@@ -127,6 +145,40 @@
                 mInputDevicesReadyMonitor.notifyAll();
             }
         }
+        Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+    }
+
+    /** Returns {@code true} if the change of input devices may affect display configuration. */
+    private boolean updateLastInputConfigurationSources() {
+        final InputDevice[] devices = mService.mInputManager.getInputDevices();
+        final SparseIntArray newSources = new SparseIntArray(8);
+        final SparseIntArray lastSources = mLastInputConfigurationSources;
+        boolean changed = lastSources == null;
+        for (InputDevice device : devices) {
+            final String descriptor = device.getDescriptor();
+            if (descriptor == null || device.isVirtual()) {
+                continue;
+            }
+            final int key = descriptor.hashCode();
+            // The interested attributes from DisplayContent#computeScreenConfiguration.
+            int newSourceHash = device.getSources();
+            newSourceHash = newSourceHash * 31 + device.getKeyboardType();
+            newSourceHash = newSourceHash * 31 + device.getAssociatedDisplayId();
+            newSourceHash = newSourceHash * 31 + (device.isExternal() ? 1 : 0);
+            newSourceHash = newSourceHash * 31 + (device.isEnabled() ? 1 : 0);
+            newSources.put(key, newSourceHash);
+            if (lastSources != null && !changed) {
+                final int lastSource = lastSources.get(key, 0 /* valueIfKeyNotFound */);
+                if (lastSource != newSourceHash) {
+                    changed = true;
+                }
+            }
+        }
+        if (lastSources != null && lastSources.size() != newSources.size()) {
+            changed = true;
+        }
+        mLastInputConfigurationSources = newSources;
+        return changed;
     }
 
     /** Notifies that the pointer location configuration has changed. */
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index 0c489d6..2fde5aa 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -62,6 +62,7 @@
 import com.android.internal.policy.IKeyguardDismissCallback;
 import com.android.server.inputmethod.InputMethodManagerInternal;
 import com.android.server.policy.WindowManagerPolicy;
+import com.android.window.flags.Flags;
 
 import java.io.PrintWriter;
 
@@ -73,6 +74,9 @@
  */
 class KeyguardController {
 
+    private static final boolean ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS =
+            Flags.ensureKeyguardDoesTransitionStarting();
+
     private static final String TAG = TAG_WITH_CLASS_NAME ? "KeyguardController" : TAG_ATM;
 
     static final String KEYGUARD_SLEEP_TOKEN_TAG = "keyguard";
@@ -201,6 +205,19 @@
             setWakeTransitionReady();
             return;
         }
+
+        if (ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS) {
+            final TransitionController transitionController =
+                    mWindowManager.mAtmService.getTransitionController();
+            final Transition transition = transitionController.getCollectingTransition();
+            if (transition != null && displayId == DEFAULT_DISPLAY) {
+                if (!keyguardShowing && state.mKeyguardShowing) {
+                    transition.addFlag(TRANSIT_FLAG_KEYGUARD_GOING_AWAY);
+                } else if (keyguardShowing && !state.mKeyguardShowing) {
+                    transition.addFlag(TRANSIT_FLAG_KEYGUARD_APPEARING);
+                }
+            }
+        }
         // Update the task snapshot if the screen will not be turned off. To make sure that the
         // unlocking animation can animate consistent content. The conditions are:
         // - Either AOD or keyguard changes to be showing. So if the states change individually,
@@ -231,8 +248,10 @@
                     || (keyguardShowing && !Display.isOffState(dc.getDisplayInfo().state))) {
                 // Keyguard decided to show or stopped going away. Send a transition to animate back
                 // to the locked state before holding the sleep token again
-                dc.requestTransitionAndLegacyPrepare(
-                        TRANSIT_TO_FRONT, TRANSIT_FLAG_KEYGUARD_APPEARING);
+                if (!ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS) {
+                    dc.requestTransitionAndLegacyPrepare(
+                            TRANSIT_TO_FRONT, TRANSIT_FLAG_KEYGUARD_APPEARING);
+                }
                 dc.mWallpaperController.adjustWallpaperWindows();
                 dc.executeAppTransition();
             }
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index af57c84..95cf6bc 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -36,7 +36,7 @@
 import static android.view.WindowManager.LayoutParams.INVALID_WINDOW_TYPE;
 import static android.view.WindowManager.TRANSIT_CHANGE;
 import static android.window.TaskFragmentAnimationParams.DEFAULT_ANIMATION_BACKGROUND_COLOR;
-import static android.window.flags.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
+import static android.window.DesktopModeFlags.ENABLE_CAPTION_COMPAT_INSET_FORCE_CONSUMPTION;
 
 import static com.android.internal.protolog.WmProtoLogGroups.WM_DEBUG_ANIM;
 import static com.android.internal.protolog.WmProtoLogGroups.WM_DEBUG_APP_TRANSITIONS;
diff --git a/services/core/java/com/android/server/wm/WindowManagerConstants.java b/services/core/java/com/android/server/wm/WindowManagerConstants.java
index e0f24d8..31ca24c 100644
--- a/services/core/java/com/android/server/wm/WindowManagerConstants.java
+++ b/services/core/java/com/android/server/wm/WindowManagerConstants.java
@@ -22,10 +22,12 @@
 import android.provider.AndroidDeviceConfig;
 import android.provider.DeviceConfig;
 import android.provider.DeviceConfigInterface;
+import android.util.ArraySet;
 
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.PrintWriter;
+import java.util.Arrays;
 import java.util.Objects;
 import java.util.concurrent.Executor;
 
@@ -38,6 +40,10 @@
     private static final String KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST =
             "ignore_activity_orientation_request";
 
+    /** The packages that ignore {@link #KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST}. */
+    private static final String KEY_OPT_OUT_IGNORE_ACTIVITY_ORIENTATION_REQUEST_LIST =
+            "opt_out_ignore_activity_orientation_request_list";
+
     /**
      * The minimum duration between gesture exclusion logging for a given window in
      * milliseconds.
@@ -65,6 +71,9 @@
     /** @see #KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST */
     boolean mIgnoreActivityOrientationRequest;
 
+    /** @see #KEY_OPT_OUT_IGNORE_ACTIVITY_ORIENTATION_REQUEST_LIST */
+    private ArraySet<String> mOptOutIgnoreActivityOrientationRequestPackages;
+
     private final WindowManagerGlobalLock mGlobalLock;
     private final Runnable mUpdateSystemGestureExclusionCallback;
     private final DeviceConfigInterface mDeviceConfig;
@@ -97,6 +106,7 @@
         updateSystemGestureExclusionLimitDp();
         updateSystemGestureExcludedByPreQStickyImmersive();
         updateIgnoreActivityOrientationRequest();
+        updateOptOutIgnoreActivityOrientationRequestList();
     }
 
     private void onAndroidPropertiesChanged(DeviceConfig.Properties properties) {
@@ -138,6 +148,9 @@
                     case KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST:
                         updateIgnoreActivityOrientationRequest();
                         break;
+                    case KEY_OPT_OUT_IGNORE_ACTIVITY_ORIENTATION_REQUEST_LIST:
+                        updateOptOutIgnoreActivityOrientationRequestList();
+                        break;
                     default:
                         break;
                 }
@@ -169,6 +182,25 @@
                 KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST, false);
     }
 
+    private void updateOptOutIgnoreActivityOrientationRequestList() {
+        final String packageList = mDeviceConfig.getString(
+                DeviceConfig.NAMESPACE_WINDOW_MANAGER,
+                KEY_OPT_OUT_IGNORE_ACTIVITY_ORIENTATION_REQUEST_LIST, "");
+        if (packageList.isEmpty()) {
+            mOptOutIgnoreActivityOrientationRequestPackages = null;
+            return;
+        }
+        mOptOutIgnoreActivityOrientationRequestPackages = new ArraySet<>();
+        mOptOutIgnoreActivityOrientationRequestPackages.addAll(
+                Arrays.asList(packageList.split(",")));
+    }
+
+    boolean isPackageOptOutIgnoreActivityOrientationRequest(String packageName) {
+        return mIgnoreActivityOrientationRequest
+                && mOptOutIgnoreActivityOrientationRequestPackages != null
+                && mOptOutIgnoreActivityOrientationRequestPackages.contains(packageName);
+    }
+
     void dump(PrintWriter pw) {
         pw.println("WINDOW MANAGER CONSTANTS (dumpsys window constants):");
 
@@ -180,6 +212,10 @@
         pw.print("="); pw.println(mSystemGestureExcludedByPreQStickyImmersive);
         pw.print("  "); pw.print(KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST);
         pw.print("="); pw.println(mIgnoreActivityOrientationRequest);
+        if (mOptOutIgnoreActivityOrientationRequestPackages != null) {
+            pw.print("  "); pw.print(KEY_OPT_OUT_IGNORE_ACTIVITY_ORIENTATION_REQUEST_LIST);
+            pw.print("="); pw.println(mOptOutIgnoreActivityOrientationRequestPackages);
+        }
         pw.println();
     }
 }
diff --git a/services/core/java/com/android/server/wm/WindowOrganizerController.java b/services/core/java/com/android/server/wm/WindowOrganizerController.java
index f8d0bc2..2229807 100644
--- a/services/core/java/com/android/server/wm/WindowOrganizerController.java
+++ b/services/core/java/com/android/server/wm/WindowOrganizerController.java
@@ -51,6 +51,7 @@
 import static android.window.WindowContainerTransaction.Change.CHANGE_FORCE_TRANSLUCENT;
 import static android.window.WindowContainerTransaction.Change.CHANGE_HIDDEN;
 import static android.window.WindowContainerTransaction.Change.CHANGE_RELATIVE_BOUNDS;
+import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_SET_KEYGUARD_STATE;
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_ADD_INSETS_FRAME_PROVIDER;
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_ADD_TASK_FRAGMENT_OPERATION;
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_CHILDREN_TASKS_REPARENT;
@@ -121,6 +122,7 @@
 import android.window.ITransitionPlayer;
 import android.window.IWindowContainerTransactionCallback;
 import android.window.IWindowOrganizerController;
+import android.window.KeyguardState;
 import android.window.RemoteTransition;
 import android.window.TaskFragmentAnimationParams;
 import android.window.TaskFragmentCreationParams;
@@ -1253,6 +1255,10 @@
                         caller, errorCallbackToken, organizer);
                 break;
             }
+            case HIERARCHY_OP_TYPE_SET_KEYGUARD_STATE: {
+                effects |= applyKeyguardState(hop);
+                break;
+            }
             case HIERARCHY_OP_TYPE_PENDING_INTENT: {
                 final Bundle launchOpts = hop.getLaunchOptions();
                 ActivityOptions activityOptions = launchOpts != null
@@ -1788,6 +1794,19 @@
         return effects;
     }
 
+    private int applyKeyguardState(@NonNull WindowContainerTransaction.HierarchyOp hop) {
+        int effects = TRANSACT_EFFECTS_NONE;
+
+        final KeyguardState keyguardState = hop.getKeyguardState();
+        if (keyguardState != null) {
+            int displayId = keyguardState.getDisplayId();
+            boolean keyguardShowing = keyguardState.getKeyguardShowing();
+            boolean aodShowing = keyguardState.getAodShowing();
+            mService.mKeyguardController.setKeyguardShown(displayId, keyguardShowing, aodShowing);
+        }
+        return effects;
+    }
+
     /**
      * Executes the provided {@code runnable} after the {@code transition}. If the
      * {@code transition} is {@code null}, the {@code runnable} is executed immediately.
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayManagerServiceTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayManagerServiceTest.java
index 6093a67..1a1c8e5 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayManagerServiceTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayManagerServiceTest.java
@@ -50,6 +50,7 @@
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.atLeastOnce;
@@ -3445,6 +3446,31 @@
         verify(dpc).onDisplayChanged(/* hbmMetadata= */ null, Layout.NO_LEAD_DISPLAY);
     }
 
+    @Test
+    public void testCreateAndReleaseVirtualDisplay_CalledWithTheSameUid() {
+        DisplayManagerService displayManager =
+                new DisplayManagerService(mContext, mShortMockedInjector);
+        registerDefaultDisplays(displayManager);
+        DisplayManagerService.BinderService bs = displayManager.new BinderService();
+        VirtualDisplayConfig config = mock(VirtualDisplayConfig.class);
+        Surface surface = mock(Surface.class);
+        when(config.getSurface()).thenReturn(surface);
+        int callingUid = Binder.getCallingUid();
+        IBinder binder = mock(IBinder.class);
+        when(mMockAppToken.asBinder()).thenReturn(binder);
+        String uniqueId = "123";
+        when(config.getUniqueId()).thenReturn(uniqueId);
+
+        bs.createVirtualDisplay(config, mMockAppToken, /* projection= */ null, PACKAGE_NAME);
+        verify(mMockVirtualDisplayAdapter).createVirtualDisplayLocked(eq(mMockAppToken),
+                /* projection= */ isNull(), eq(callingUid), eq(PACKAGE_NAME),
+                eq("virtual:" + PACKAGE_NAME + ":" + uniqueId), eq(surface), /* flags= */ anyInt(),
+                eq(config));
+
+        bs.releaseVirtualDisplay(mMockAppToken);
+        verify(mMockVirtualDisplayAdapter).releaseVirtualDisplayLocked(binder, callingUid);
+    }
+
     private void initDisplayPowerController(DisplayManagerInternal localService) {
         localService.initPowerManagement(new DisplayManagerInternal.DisplayPowerCallbacks() {
             @Override
diff --git a/services/tests/displayservicetests/src/com/android/server/display/VirtualDisplayAdapterTest.java b/services/tests/displayservicetests/src/com/android/server/display/VirtualDisplayAdapterTest.java
index 81e6cc3..3ac7fb0 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/VirtualDisplayAdapterTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/VirtualDisplayAdapterTest.java
@@ -21,71 +21,104 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.when;
 
-import android.content.Context;
 import android.hardware.display.IVirtualDisplayCallback;
 import android.hardware.display.VirtualDisplayConfig;
+import android.media.projection.IMediaProjection;
 import android.os.IBinder;
 import android.os.Process;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
+import android.testing.TestableContext;
+import android.view.Surface;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
 
+import com.android.internal.R;
 import com.android.server.display.feature.DisplayManagerFlags;
 import com.android.server.testutils.TestHandler;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.ArrayList;
+import java.util.List;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class VirtualDisplayAdapterTest {
 
-    @Mock
-    Context mContextMock;
+    private static final int MAX_DEVICES = 3;
+    private static final int MAX_DEVICES_PER_PACKAGE = 2;
+
+    @Rule
+    public final TestableContext mContext = new TestableContext(
+            InstrumentationRegistry.getInstrumentation().getContext());
 
     @Mock
-    VirtualDisplayAdapter.SurfaceControlDisplayFactory mMockSufaceControlDisplayFactory;
+    private VirtualDisplayAdapter.SurfaceControlDisplayFactory mMockSufaceControlDisplayFactory;
 
     @Mock
-    DisplayAdapter.Listener mMockListener;
+    private DisplayAdapter.Listener mMockListener;
 
     @Mock
-    IVirtualDisplayCallback mMockCallback;
+    private IVirtualDisplayCallback mMockCallback;
 
     @Mock
-    IBinder mMockBinder;
+    private IBinder mMockBinder;
+
+    @Mock
+    private IMediaProjection mMediaProjectionMock;
+
+    @Mock
+    private Surface mSurfaceMock;
+
+    @Mock
+    private VirtualDisplayConfig mVirtualDisplayConfigMock;
 
     private TestHandler mHandler;
 
-    private VirtualDisplayAdapter mVirtualDisplayAdapter;
-
     @Mock
     private DisplayManagerFlags mFeatureFlags;
 
+    private VirtualDisplayAdapter mAdapter;
+
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+
+        mContext.getOrCreateTestableResources().addOverride(R.integer.config_virtualDisplayLimit,
+                MAX_DEVICES);
+        mContext.getOrCreateTestableResources().addOverride(
+                R.integer.config_virtualDisplayLimitPerPackage, MAX_DEVICES_PER_PACKAGE);
+
         mHandler = new TestHandler(null);
-        mVirtualDisplayAdapter = new VirtualDisplayAdapter(new DisplayManagerService.SyncRoot(),
-                mContextMock, mHandler, mMockListener, mMockSufaceControlDisplayFactory,
-                mFeatureFlags);
+        mAdapter = new VirtualDisplayAdapter(new DisplayManagerService.SyncRoot(), mContext,
+                mHandler, mMockListener, mMockSufaceControlDisplayFactory, mFeatureFlags);
 
         when(mMockCallback.asBinder()).thenReturn(mMockBinder);
     }
 
     @Test
-    public void testCreatesVirtualDisplay() {
+    public void testCreateAndReleaseVirtualDisplay() {
         VirtualDisplayConfig config = new VirtualDisplayConfig.Builder("test", /* width= */ 1,
                 /* height= */ 1, /* densityDpi= */ 1).build();
+        int ownerUid = 10;
 
-        DisplayDevice result = mVirtualDisplayAdapter.createVirtualDisplayLocked(mMockCallback,
-                /* projection= */ null, /* ownerUid= */ 10, /* packageName= */ "testpackage",
+        DisplayDevice result = mAdapter.createVirtualDisplayLocked(mMockCallback,
+                /* projection= */ null, ownerUid, /* packageName= */ "testpackage",
                 /* uniqueId= */ "uniqueId", /* surface= */ null, /* flags= */ 0, config);
+        assertNotNull(result);
 
+        result = mAdapter.releaseVirtualDisplayLocked(mMockBinder, ownerUid);
         assertNotNull(result);
     }
 
@@ -98,7 +131,7 @@
         final String displayUniqueId = VirtualDisplayAdapter.generateDisplayUniqueId(
                 packageName, Process.myUid(), config);
 
-        DisplayDevice result = mVirtualDisplayAdapter.createVirtualDisplayLocked(
+        DisplayDevice result = mAdapter.createVirtualDisplayLocked(
                 mMockCallback, /* projection= */ null, /* ownerUid= */ 10,
                 packageName, displayUniqueId, /* surface= */ null, /* flags= */ 0, config);
 
@@ -114,14 +147,194 @@
                 /* height= */ 1, /* densityDpi= */ 1).build();
         VirtualDisplayConfig config2 = new VirtualDisplayConfig.Builder("test2", /* width= */ 1,
                 /* height= */ 1, /* densityDpi= */ 1).build();
-        mVirtualDisplayAdapter.createVirtualDisplayLocked(mMockCallback, /* projection= */ null,
-                /* ownerUid= */ 10, /* packageName= */ "testpackage", /* uniqueId= */ "uniqueId1",
-                /* surface= */ null, /* flags= */ 0, config1);
 
-        DisplayDevice result = mVirtualDisplayAdapter.createVirtualDisplayLocked(mMockCallback,
+        DisplayDevice result = mAdapter.createVirtualDisplayLocked(mMockCallback,
+                /* projection= */ null, /* ownerUid= */ 10, /* packageName= */ "testpackage",
+                /* uniqueId= */ "uniqueId1", /* surface= */ null, /* flags= */ 0, config1);
+        assertNotNull(result);
+
+        result = mAdapter.createVirtualDisplayLocked(mMockCallback,
                 /* projection= */ null, /* ownerUid= */ 10, /* packageName= */ "testpackage",
                 /* uniqueId= */ "uniqueId2", /* surface= */ null, /* flags= */ 0, config2);
-
         assertNull(result);
     }
+
+    @Test
+    public void testCreateManyVirtualDisplays_LimitFlagDisabled() {
+        when(mFeatureFlags.isVirtualDisplayLimitEnabled()).thenReturn(false);
+
+        // Displays for the same package
+        for (int i = 0; i < MAX_DEVICES_PER_PACKAGE * 2; i++) {
+            // Same owner UID
+            IVirtualDisplayCallback callback = createCallback();
+            DisplayDevice device = mAdapter.createVirtualDisplayLocked(callback,
+                    mMediaProjectionMock, 1234, "test.package", "123",
+                    mSurfaceMock, /* flags= */ 0, mVirtualDisplayConfigMock);
+            assertNotNull(device);
+        }
+
+        // Displays for different packages
+        for (int i = 0; i < MAX_DEVICES * 2; i++) {
+            // Same owner UID
+            IVirtualDisplayCallback callback = createCallback();
+            DisplayDevice device = mAdapter.createVirtualDisplayLocked(callback,
+                    mMediaProjectionMock, 1234 + i, "test.package", "123",
+                    mSurfaceMock, /* flags= */ 0, mVirtualDisplayConfigMock);
+            assertNotNull(device);
+        }
+    }
+
+    @Test
+    public void testCreateVirtualDisplay_MaxDisplaysPerPackage() {
+        when(mFeatureFlags.isVirtualDisplayLimitEnabled()).thenReturn(true);
+        List<IVirtualDisplayCallback> callbacks = new ArrayList<>();
+        int ownerUid = 1234;
+
+        for (int i = 0; i < MAX_DEVICES_PER_PACKAGE * 2; i++) {
+            // Same owner UID
+            IVirtualDisplayCallback callback = createCallback();
+            DisplayDevice device = mAdapter.createVirtualDisplayLocked(callback,
+                    mMediaProjectionMock, ownerUid, "test.package", "123",
+                    mSurfaceMock, /* flags= */ 0, mVirtualDisplayConfigMock);
+            if (i < MAX_DEVICES_PER_PACKAGE) {
+                assertNotNull(device);
+                callbacks.add(callback);
+            } else {
+                assertNull(device);
+            }
+        }
+
+        // Release one display
+        DisplayDevice device = mAdapter.releaseVirtualDisplayLocked(callbacks.get(0).asBinder(),
+                ownerUid);
+        assertNotNull(device);
+        callbacks.remove(0);
+
+        // We should be able to create another display
+        IVirtualDisplayCallback callback = createCallback();
+        device = mAdapter.createVirtualDisplayLocked(callback, mMediaProjectionMock, ownerUid,
+                "test.package", "123", mSurfaceMock, /* flags= */ 0,
+                mVirtualDisplayConfigMock);
+        assertNotNull(device);
+        callbacks.add(callback);
+
+        // But only one
+        callback = createCallback();
+        device = mAdapter.createVirtualDisplayLocked(callback, mMediaProjectionMock, ownerUid,
+                "test.package", "123", mSurfaceMock, /* flags= */ 0,
+                mVirtualDisplayConfigMock);
+        assertNull(device);
+
+        // Release all the displays
+        for (IVirtualDisplayCallback cb : callbacks) {
+            device = mAdapter.releaseVirtualDisplayLocked(cb.asBinder(), ownerUid);
+            assertNotNull(device);
+        }
+        callbacks.clear();
+
+        // We should be able to create the max number of displays again
+        for (int i = 0; i < MAX_DEVICES_PER_PACKAGE * 2; i++) {
+            // Same owner UID
+            callback = createCallback();
+            device = mAdapter.createVirtualDisplayLocked(callback, mMediaProjectionMock, ownerUid,
+                    "test.package", "123", mSurfaceMock, /* flags= */ 0,
+                    mVirtualDisplayConfigMock);
+            if (i < MAX_DEVICES_PER_PACKAGE) {
+                assertNotNull(device);
+                callbacks.add(callback);
+            } else {
+                assertNull(device);
+            }
+        }
+
+        // We should be able to create a display for a different package
+        callback = createCallback();
+        device = mAdapter.createVirtualDisplayLocked(callback, mMediaProjectionMock, ownerUid + 1,
+                "test.package", "123", mSurfaceMock, /* flags= */ 0,
+                mVirtualDisplayConfigMock);
+        assertNotNull(device);
+        callbacks.add(callback);
+    }
+
+    @Test
+    public void testCreateVirtualDisplay_MaxDisplays() {
+        when(mFeatureFlags.isVirtualDisplayLimitEnabled()).thenReturn(true);
+        List<IVirtualDisplayCallback> callbacks = new ArrayList<>();
+        int firstOwnerUid = 1000;
+
+        for (int i = 0; i < MAX_DEVICES * 2; i++) {
+            // Different owner UID
+            IVirtualDisplayCallback callback = createCallback();
+            DisplayDevice device = mAdapter.createVirtualDisplayLocked(callback,
+                    mMediaProjectionMock, firstOwnerUid + i, "test.package",
+                    "123", mSurfaceMock, /* flags= */ 0, mVirtualDisplayConfigMock);
+            if (i < MAX_DEVICES) {
+                assertNotNull(device);
+                callbacks.add(callback);
+            } else {
+                assertNull(device);
+            }
+        }
+
+        // Release one display
+        DisplayDevice device = mAdapter.releaseVirtualDisplayLocked(callbacks.get(0).asBinder(),
+                firstOwnerUid);
+        assertNotNull(device);
+        callbacks.remove(0);
+
+        // We should be able to create another display
+        IVirtualDisplayCallback callback = createCallback();
+        device = mAdapter.createVirtualDisplayLocked(callback, mMediaProjectionMock,
+                firstOwnerUid, "test.package", "123", mSurfaceMock, /* flags= */ 0,
+                mVirtualDisplayConfigMock);
+        assertNotNull(device);
+        callbacks.add(callback);
+
+        // But only one
+        callback = createCallback();
+        device = mAdapter.createVirtualDisplayLocked(callback, mMediaProjectionMock,
+                firstOwnerUid, "test.package", "123", mSurfaceMock, /* flags= */ 0,
+                mVirtualDisplayConfigMock);
+        assertNull(device);
+
+        // Release all the displays
+        for (int i = 0; i < callbacks.size(); i++) {
+            device = mAdapter.releaseVirtualDisplayLocked(callbacks.get(i).asBinder(),
+                    firstOwnerUid + i);
+            assertNotNull(device);
+        }
+        callbacks.clear();
+
+        // We should be able to create the max number of displays again
+        for (int i = 0; i < MAX_DEVICES * 2; i++) {
+            // Different owner UID
+            callback = createCallback();
+            device = mAdapter.createVirtualDisplayLocked(callback, mMediaProjectionMock,
+                    firstOwnerUid + i, "test.package", "123", mSurfaceMock,
+                    /* flags= */ 0, mVirtualDisplayConfigMock);
+            if (i < MAX_DEVICES) {
+                assertNotNull(device);
+                callbacks.add(callback);
+            } else {
+                assertNull(device);
+            }
+        }
+    }
+
+    private IVirtualDisplayCallback createCallback() {
+        return new IVirtualDisplayCallback.Stub() {
+
+            @Override
+            public void onPaused() {
+            }
+
+            @Override
+            public void onResumed() {
+            }
+
+            @Override
+            public void onStopped() {
+            }
+        };
+    }
 }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java
index d91f154..58f0ab4 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayModeDirectorTest.java
@@ -1872,6 +1872,60 @@
     }
 
     @Test
+    public void testPeakRefreshRate_notAppliedToExternalDisplays() {
+        when(mDisplayManagerFlags.isBackUpSmoothDisplayAndForcePeakRefreshRateEnabled())
+                .thenReturn(true);
+        mInjector.mDisplayInfo.type = Display.TYPE_EXTERNAL;
+        DisplayModeDirector director =
+                new DisplayModeDirector(mContext, mHandler, mInjector,
+                        mDisplayManagerFlags, mDisplayDeviceConfigProvider);
+        director.getBrightnessObserver().setDefaultDisplayState(Display.STATE_ON);
+        director.getDisplayObserver().onDisplayAdded(DISPLAY_ID);
+        director.getDisplayObserver().onDisplayAdded(DISPLAY_ID_2);
+
+        Display.Mode[] modes1 = new Display.Mode[] {
+                new Display.Mode(/* modeId= */ 1, /* width= */ 1280, /* height= */ 720,
+                        /* refreshRate= */ 60),
+                new Display.Mode(/* modeId= */ 2, /* width= */ 1280, /* height= */ 720,
+                        /* refreshRate= */ 130),
+        };
+        Display.Mode[] modes2 = new Display.Mode[] {
+                new Display.Mode(/* modeId= */ 1, /* width= */ 1280, /* height= */ 720,
+                        /* refreshRate= */ 60),
+                new Display.Mode(/* modeId= */ 2, /* width= */ 1280, /* height= */ 720,
+                        /* refreshRate= */ 140),
+        };
+        SparseArray<Display.Mode[]> supportedModesByDisplay = new SparseArray<>();
+        supportedModesByDisplay.put(DISPLAY_ID, modes1);
+        supportedModesByDisplay.put(DISPLAY_ID_2, modes2);
+
+        Sensor lightSensor = createLightSensor();
+        SensorManager sensorManager = createMockSensorManager(lightSensor);
+        director.start(sensorManager);
+        director.injectSupportedModesByDisplay(supportedModesByDisplay);
+
+        // Disable Smooth Display
+        setPeakRefreshRate(RefreshRateSettingsUtils.DEFAULT_REFRESH_RATE);
+
+        Vote vote1 = director.getVote(DISPLAY_ID,
+                Vote.PRIORITY_USER_SETTING_PEAK_RENDER_FRAME_RATE);
+        Vote vote2 = director.getVote(DISPLAY_ID_2,
+                Vote.PRIORITY_USER_SETTING_PEAK_RENDER_FRAME_RATE);
+        assertThat(vote1).isNull();
+        assertThat(vote2).isNull();
+
+        // Enable Smooth Display
+        setPeakRefreshRate(Float.POSITIVE_INFINITY);
+
+        vote1 = director.getVote(DISPLAY_ID,
+                Vote.PRIORITY_USER_SETTING_PEAK_RENDER_FRAME_RATE);
+        vote2 = director.getVote(DISPLAY_ID_2,
+                Vote.PRIORITY_USER_SETTING_PEAK_RENDER_FRAME_RATE);
+        assertThat(vote1).isNull();
+        assertThat(vote2).isNull();
+    }
+
+    @Test
     public void testPeakRefreshRate_DisplayChanged() {
         when(mDisplayManagerFlags.isBackUpSmoothDisplayAndForcePeakRefreshRateEnabled())
                 .thenReturn(true);
@@ -1968,8 +2022,9 @@
     @Test
     @Parameters({
         "true, true, 60",
-        "false, true, 50",
-        "true, false, 50"
+        "false, true, 60",
+        "true, false, 50",
+        "false, false, 50"
     })
     public void testExternalDisplayMaxRefreshRate(boolean isRefreshRateSynchronizationEnabled,
             boolean isExternalDisplay, float expectedMaxRenderFrameRate) {
@@ -3810,6 +3865,7 @@
                 SensorManagerInternal sensorManagerInternal) {
             mDeviceConfig = new FakeDeviceConfig();
             mDisplayInfo = new DisplayInfo();
+            mDisplayInfo.type = Display.TYPE_INTERNAL;
             mDisplayInfo.defaultModeId = MODE_ID;
             mDisplayInfo.supportedModes = new Display.Mode[] {new Display.Mode(MODE_ID,
                     800, 600, /* refreshRate= */ 60)};
@@ -3856,6 +3912,7 @@
         @Override
         public boolean getDisplayInfo(int displayId, DisplayInfo displayInfo) {
             displayInfo.copyFrom(mDisplayInfo);
+            displayInfo.displayId = displayId;
             return mDisplayInfoValid;
         }
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java
index 5e240cf..e3f150e 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/DisplayObserverTest.java
@@ -189,6 +189,7 @@
     @Test
     public void testExternalDisplay_voteUserPreferredMode() {
         when(mDisplayManagerFlags.isUserPreferredModeVoteEnabled()).thenReturn(true);
+        when(mDisplayManagerFlags.isUserRefreshRateForExternalDisplayEnabled()).thenReturn(false);
         var preferredMode = mExternalDisplayModes[5];
         mExternalDisplayUserPreferredModeId = preferredMode.getModeId();
         var expectedVote = Vote.forSize(
@@ -229,6 +230,108 @@
                 .isEqualTo(null);
     }
 
+    /** Vote for user preferred mode */
+    @Test
+    public void testDefaultDisplay_voteUserPreferredMode() {
+        when(mDisplayManagerFlags.isUserPreferredModeVoteEnabled()).thenReturn(true);
+        when(mDisplayManagerFlags.isUserRefreshRateForExternalDisplayEnabled()).thenReturn(true);
+        var preferredMode = mInternalDisplayModes[5];
+        mInternalDisplayUserPreferredModeId = preferredMode.getModeId();
+        var expectedVote = Vote.forSize(
+                        preferredMode.getPhysicalWidth(),
+                        preferredMode.getPhysicalHeight());
+        init();
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
+        mObserver.onDisplayAdded(DEFAULT_DISPLAY);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(expectedVote);
+
+        mInternalDisplayUserPreferredModeId = INVALID_MODE_ID;
+        mObserver.onDisplayChanged(EXTERNAL_DISPLAY);
+        mObserver.onDisplayChanged(DEFAULT_DISPLAY);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+
+        preferredMode = mInternalDisplayModes[4];
+        mInternalDisplayUserPreferredModeId = preferredMode.getModeId();
+        expectedVote = Vote.forSize(
+                        preferredMode.getPhysicalWidth(),
+                        preferredMode.getPhysicalHeight());
+        mObserver.onDisplayChanged(EXTERNAL_DISPLAY);
+        mObserver.onDisplayChanged(DEFAULT_DISPLAY);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(expectedVote);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+
+        // Testing that the vote is removed.
+        mObserver.onDisplayRemoved(EXTERNAL_DISPLAY);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(expectedVote);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+    }
+
+    /** Vote for user preferred mode with refresh rate, synchronization vote must be disabled */
+    @Test
+    public void testExternalDisplay_voteUserPreferredMode_withRefreshRate() {
+        when(mDisplayManagerFlags.isUserPreferredModeVoteEnabled()).thenReturn(true);
+        when(mDisplayManagerFlags.isDisplaysRefreshRatesSynchronizationEnabled()).thenReturn(false);
+        when(mDisplayManagerFlags.isUserRefreshRateForExternalDisplayEnabled()).thenReturn(true);
+        var preferredMode = mExternalDisplayModes[5];
+        mExternalDisplayUserPreferredModeId = preferredMode.getModeId();
+        var expectedVote = Vote.forSizeAndPhysicalRefreshRatesRange(
+                        preferredMode.getPhysicalWidth(),
+                        preferredMode.getPhysicalHeight(),
+                        preferredMode.getPhysicalWidth(),
+                        preferredMode.getPhysicalHeight(),
+                        preferredMode.getRefreshRate(),
+                        preferredMode.getRefreshRate());
+        init();
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        mObserver.onDisplayAdded(EXTERNAL_DISPLAY);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(expectedVote);
+
+        mExternalDisplayUserPreferredModeId = INVALID_MODE_ID;
+        mObserver.onDisplayChanged(EXTERNAL_DISPLAY);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+
+        preferredMode = mExternalDisplayModes[4];
+        mExternalDisplayUserPreferredModeId = preferredMode.getModeId();
+        expectedVote = Vote.forSizeAndPhysicalRefreshRatesRange(
+                        preferredMode.getPhysicalWidth(),
+                        preferredMode.getPhysicalHeight(),
+                        preferredMode.getPhysicalWidth(),
+                        preferredMode.getPhysicalHeight(),
+                        preferredMode.getRefreshRate(),
+                        preferredMode.getRefreshRate());
+        mObserver.onDisplayChanged(EXTERNAL_DISPLAY);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(expectedVote);
+
+        // Testing that the vote is removed.
+        mObserver.onDisplayRemoved(EXTERNAL_DISPLAY);
+        assertThat(getVote(DEFAULT_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+        assertThat(getVote(EXTERNAL_DISPLAY, PRIORITY_USER_SETTING_DISPLAY_PREFERRED_SIZE))
+                .isEqualTo(null);
+    }
+
     /** External display: Do not apply limit to user preferred mode */
     @Test
     public void testExternalDisplay_doNotApplyLimitToUserPreferredMode() {
diff --git a/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java b/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java
index 2107406..412599d 100644
--- a/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java
+++ b/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java
@@ -171,6 +171,8 @@
     private static int sUiTierSize = 5;
 
     private Context mContext;
+    private ProcessStateController mProcessStateController;
+    private ActiveUids mActiveUids;
     private PackageManagerInternal mPackageManagerInternal;
     private ActivityManagerService mService;
     private OomAdjusterInjector mInjector = new OomAdjusterInjector();
@@ -229,15 +231,19 @@
         doCallRealMethod().when(mService).enqueueOomAdjTargetLocked(any(ProcessRecord.class));
         doCallRealMethod().when(mService).updateOomAdjPendingTargetsLocked(OOM_ADJ_REASON_ACTIVITY);
         setFieldValue(AppProfiler.class, profiler, "mProfilerLock", new Object());
+
         doNothing().when(pr).enqueueProcessChangeItemLocked(anyInt(), anyInt(), anyInt(),
                 anyInt());
         doNothing().when(pr).enqueueProcessChangeItemLocked(anyInt(), anyInt(), anyInt(),
                 anyBoolean());
-        mService.mOomAdjuster = mService.mConstants.ENABLE_NEW_OOMADJ
-                ? new OomAdjusterModernImpl(mService, mService.mProcessList,
-                        new ActiveUids(mService, false), mInjector)
-                : new OomAdjuster(mService, mService.mProcessList, new ActiveUids(mService, false),
-                        mInjector);
+        mActiveUids = new ActiveUids(mService, false);
+        mProcessStateController = new ProcessStateController.Builder(mService,
+                mService.mProcessList, mActiveUids)
+                .useModernOomAdjuster(mService.mConstants.ENABLE_NEW_OOMADJ)
+                .setOomAdjusterInjector(mInjector)
+                .build();
+        mService.mProcessStateController = mProcessStateController;
+        mService.mOomAdjuster = mService.mProcessStateController.getOomAdjuster();
         mService.mOomAdjuster.mAdjSeq = 10000;
         mService.mWakefulness = new AtomicInteger(PowerManagerInternal.WAKEFULNESS_AWAKE);
         mSetFlagsRule.enableFlags(Flags.FLAG_NEW_FGS_RESTRICTION_LOGIC);
@@ -246,8 +252,8 @@
     @SuppressWarnings("GuardedBy")
     @After
     public void tearDown() {
-        mService.mOomAdjuster.resetInternal();
-        mService.mOomAdjuster.mActiveUids.clear();
+        mProcessStateController.getOomAdjuster().resetInternal();
+        mActiveUids.clear();
         LocalServices.removeServiceForTest(PackageManagerInternal.class);
         mInjector.reset();
     }
@@ -293,7 +299,7 @@
     private void updateOomAdj(ProcessRecord... apps) {
         if (apps.length == 0) {
             updateProcessRecordNodes(mService.mProcessList.getLruProcessesLOSP());
-            mService.mOomAdjuster.updateOomAdjLocked(OOM_ADJ_REASON_NONE);
+            mProcessStateController.runFullUpdate(OOM_ADJ_REASON_NONE);
         } else {
             updateProcessRecordNodes(Arrays.asList(apps));
             if (apps.length == 1) {
@@ -301,10 +307,10 @@
                 if (!mService.mProcessList.getLruProcessesLOSP().contains(app)) {
                     mService.mProcessList.getLruProcessesLOSP().add(app);
                 }
-                mService.mOomAdjuster.updateOomAdjLocked(apps[0], OOM_ADJ_REASON_NONE);
+                mProcessStateController.runUpdate(apps[0], OOM_ADJ_REASON_NONE);
             } else {
                 setProcessesToLru(apps);
-                mService.mOomAdjuster.updateOomAdjLocked(OOM_ADJ_REASON_NONE);
+                mProcessStateController.runFullUpdate(OOM_ADJ_REASON_NONE);
                 mService.mProcessList.getLruProcessesLOSP().clear();
             }
         }
@@ -318,9 +324,9 @@
     private void updateOomAdjPending(ProcessRecord... apps) {
         setProcessesToLru(apps);
         for (ProcessRecord app : apps) {
-            mService.mOomAdjuster.enqueueOomAdjTargetLocked(app);
+            mProcessStateController.enqueueUpdateTarget(app);
         }
-        mService.mOomAdjuster.updateOomAdjPendingTargetsLocked(OOM_ADJ_REASON_NONE);
+        mProcessStateController.runPendingUpdate(OOM_ADJ_REASON_NONE);
         mService.mProcessList.getLruProcessesLOSP().clear();
     }
 
@@ -341,11 +347,10 @@
     public void testUpdateOomAdj_DoOne_Persistent_TopUi_Sleeping() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        app.mState.setHasTopUi(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_ASLEEP);
+        mProcessStateController.setMaxAdj(app, PERSISTENT_PROC_ADJ);
+        mProcessStateController.setHasTopUi(app, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_ASLEEP);
         updateOomAdj(app);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, PERSISTENT_PROC_ADJ,
                 SCHED_GROUP_RESTRICTED);
@@ -357,9 +362,9 @@
     public void testUpdateOomAdj_DoOne_Persistent_TopUi_Awake() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        app.mState.setHasTopUi(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setMaxAdj(app, PERSISTENT_PROC_ADJ);
+        mProcessStateController.setHasTopUi(app, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_PERSISTENT_UI, PERSISTENT_PROC_ADJ,
@@ -371,9 +376,9 @@
     public void testUpdateOomAdj_DoOne_Persistent_TopApp() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(app, PERSISTENT_PROC_ADJ);
         doReturn(app).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doReturn(null).when(mService).getTopApp();
 
@@ -388,7 +393,7 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(app).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doReturn(null).when(mService).getTopApp();
 
@@ -401,8 +406,8 @@
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         doReturn(PROCESS_STATE_TOP_SLEEPING).when(mService.mAtmInternal).getTopProcessState();
-        app.mState.setRunningRemoteAnimation(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setRunningRemoteAnimation(app, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
 
@@ -415,7 +420,7 @@
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         doReturn(mock(ActiveInstrumentation.class)).when(app).getActiveInstrumentation();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doCallRealMethod().when(app).getActiveInstrumentation();
 
@@ -431,7 +436,7 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         doReturn(true).when(mService).isReceivingBroadcastLocked(any(ProcessRecord.class),
                 any(int[].class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doReturn(false).when(mService).isReceivingBroadcastLocked(any(ProcessRecord.class),
                 any(int[].class));
@@ -466,8 +471,8 @@
     public void testUpdateOomAdj_DoOne_ExecutingService() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mServices.startExecutingService(mock(ServiceRecord.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.startExecutingService(app.mServices, mock(ServiceRecord.class));
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_SERVICE, FOREGROUND_APP_ADJ, SCHED_GROUP_BACKGROUND);
@@ -480,11 +485,11 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         doReturn(PROCESS_STATE_TOP_SLEEPING).when(mService.mAtmInternal).getTopProcessState();
         doReturn(app).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_ASLEEP);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_ASLEEP);
         updateOomAdj(app);
         doReturn(null).when(mService).getTopApp();
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
         assertProcStates(app, PROCESS_STATE_TOP_SLEEPING, FOREGROUND_APP_ADJ,
                 SCHED_GROUP_BACKGROUND);
@@ -498,7 +503,7 @@
         app.mState.setCurRawAdj(CACHED_APP_MIN_ADJ);
         app.mState.setCurAdj(CACHED_APP_MIN_ADJ);
         doReturn(null).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         final int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
@@ -516,7 +521,7 @@
         doReturn(true).when(wpc).hasActivities();
         doReturn(WindowProcessController.ACTIVITY_STATE_FLAG_IS_VISIBLE)
                 .when(wpc).getActivityStateFlags();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_TOP, VISIBLE_APP_ADJ, SCHED_GROUP_DEFAULT);
@@ -555,7 +560,7 @@
         WindowProcessController wpc = app.getWindowProcessController();
         doReturn(true).when(wpc).hasRecentTasks();
         app.mState.setLastTopTime(SystemClock.uptimeMillis());
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doCallRealMethod().when(wpc).hasRecentTasks();
 
@@ -567,9 +572,9 @@
     public void testUpdateOomAdj_DoOne_FgServiceLocation() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mServices.setHasForegroundServices(true, ServiceInfo.FOREGROUND_SERVICE_TYPE_LOCATION,
-                /* hasNoneType=*/false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(app.mServices, true,
+                ServiceInfo.FOREGROUND_SERVICE_TYPE_LOCATION, /* hasNoneType=*/false);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -582,8 +587,9 @@
     public void testUpdateOomAdj_DoOne_FgService() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(app.mServices, true, 0, /* hasNoneType=*/
+                true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -599,20 +605,20 @@
 
         ServiceRecord s = ServiceRecord.newEmptyInstanceForTest(mService);
         s.appInfo = new ApplicationInfo();
-        s.startRequested = true;
+        mProcessStateController.setStartRequested(s, true);
         s.isForeground = true;
         s.foregroundServiceType = FOREGROUND_SERVICE_TYPE_SHORT_SERVICE;
-        s.setShortFgsInfo(SystemClock.uptimeMillis());
+        mProcessStateController.setShortFgsInfo(s, SystemClock.uptimeMillis());
 
         // SHORT_SERVICE FGS will get IMP_FG and a slightly different recent-adjustment.
         {
             ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                     MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-            app.mServices.startService(s);
-            app.mServices.setHasForegroundServices(true,
+            mProcessStateController.startService(app.mServices, s);
+            mProcessStateController.setHasForegroundServices(app.mServices, true,
                     FOREGROUND_SERVICE_TYPE_SHORT_SERVICE, /* hasNoneType=*/false);
             app.mState.setLastTopTime(SystemClock.uptimeMillis());
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
             updateOomAdj(app);
 
@@ -625,9 +631,9 @@
         {
             ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                     MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-            app.mServices.setHasForegroundServices(true,
+            mProcessStateController.setHasForegroundServices(app.mServices, true,
                     FOREGROUND_SERVICE_TYPE_SHORT_SERVICE, /* hasNoneType=*/false);
-            app.mServices.startService(s);
+            mProcessStateController.startService(app.mServices, s);
             app.mState.setLastTopTime(SystemClock.uptimeMillis()
                     - mService.mConstants.TOP_TO_FGS_GRACE_DURATION);
             mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
@@ -642,21 +648,21 @@
         // SHORT_SERVICE, timed out already.
         s = ServiceRecord.newEmptyInstanceForTest(mService);
         s.appInfo = new ApplicationInfo();
-        s.startRequested = true;
+        mProcessStateController.setStartRequested(s, true);
         s.isForeground = true;
         s.foregroundServiceType = FOREGROUND_SERVICE_TYPE_SHORT_SERVICE;
-        s.setShortFgsInfo(SystemClock.uptimeMillis()
+        mProcessStateController.setShortFgsInfo(s, SystemClock.uptimeMillis()
                 - mService.mConstants.mShortFgsTimeoutDuration
                 - mService.mConstants.mShortFgsProcStateExtraWaitDuration);
         {
             ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                     MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-            app.mServices.setHasForegroundServices(true,
+            mProcessStateController.setHasForegroundServices(app.mServices, true,
                     FOREGROUND_SERVICE_TYPE_SHORT_SERVICE, /* hasNoneType=*/false);
-            app.mServices.startService(s);
+            mProcessStateController.startService(app.mServices, s);
             app.mState.setLastTopTime(SystemClock.uptimeMillis()
                     - mService.mConstants.TOP_TO_FGS_GRACE_DURATION);
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
             updateOomAdj(app);
 
@@ -671,8 +677,8 @@
     public void testUpdateOomAdj_DoOne_OverlayUi() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mState.setHasOverlayUi(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasOverlayUi(app, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_IMPORTANT_FOREGROUND, PERCEPTIBLE_APP_ADJ,
@@ -684,9 +690,10 @@
     public void testUpdateOomAdj_DoOne_PerceptibleRecent_FgService() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(app.mServices, true, 0, /* hasNoneType=*/
+                true);
         app.mState.setLastTopTime(SystemClock.uptimeMillis());
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE,
@@ -699,7 +706,7 @@
         verify(mService.mHandler).sendEmptyMessageAtTime(
                 eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG), followUpTimeCaptor.capture());
         mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
-        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+        mProcessStateController.runFollowUpUpdate();
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
                 SCHED_GROUP_DEFAULT, "fg-service");
@@ -722,9 +729,9 @@
             // Simulate the system starting and binding to a service in the app.
             ServiceRecord s = bindService(app, system,
                     null, null, Context.BIND_ALMOST_PERCEPTIBLE, mock(IBinder.class));
-            s.lastTopAlmostPerceptibleBindRequestUptimeMs = nowUptime;
+            mProcessStateController.setLastTopAlmostPerceptibleBindRequest(s, nowUptime);
             s.getConnections().clear();
-            app.mServices.updateHasTopStartedAlmostPerceptibleServices();
+            mProcessStateController.updateHasTopStartedAlmostPerceptibleServices(app.mServices);
             mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
             updateOomAdj(app);
 
@@ -736,7 +743,7 @@
             verify(mService.mHandler).sendEmptyMessageAtTime(
                     eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG), followUpTimeCaptor.capture());
             mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
-            mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+            mProcessStateController.runFollowUpUpdate();
 
             final int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
                     ? sFirstUiCachedAdj : sFirstCachedAdj;
@@ -758,15 +765,15 @@
             // Simulate the system starting and binding to a service in the app.
             ServiceRecord s = bindService(app, system,
                     null, null, Context.BIND_ALMOST_PERCEPTIBLE + 2, mock(IBinder.class));
-            s.lastTopAlmostPerceptibleBindRequestUptimeMs =
-                    nowUptime - 2 * mService.mConstants.mServiceBindAlmostPerceptibleTimeoutMs;
-            app.mServices.updateHasTopStartedAlmostPerceptibleServices();
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            mProcessStateController.setLastTopAlmostPerceptibleBindRequest(s,
+                    nowUptime - 2 * mService.mConstants.mServiceBindAlmostPerceptibleTimeoutMs);
+            mProcessStateController.updateHasTopStartedAlmostPerceptibleServices(app.mServices);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
             updateOomAdj(app);
 
             assertEquals(PERCEPTIBLE_RECENT_FOREGROUND_APP_ADJ + 2, app.mState.getSetAdj());
 
-            mService.mOomAdjuster.resetInternal();
+            mProcessStateController.getOomAdjuster().resetInternal();
         }
 
         // Out of grace period and no valid binding so no adjustment.
@@ -780,16 +787,16 @@
             // Simulate the system starting and binding to a service in the app.
             ServiceRecord s = bindService(app, system,
                     null, null, Context.BIND_ALMOST_PERCEPTIBLE, mock(IBinder.class));
-            s.lastTopAlmostPerceptibleBindRequestUptimeMs =
-                    nowUptime - 2 * mService.mConstants.mServiceBindAlmostPerceptibleTimeoutMs;
+            mProcessStateController.setLastTopAlmostPerceptibleBindRequest(s,
+                    nowUptime - 2 * mService.mConstants.mServiceBindAlmostPerceptibleTimeoutMs);
             s.getConnections().clear();
-            app.mServices.updateHasTopStartedAlmostPerceptibleServices();
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            mProcessStateController.updateHasTopStartedAlmostPerceptibleServices(app.mServices);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
             updateOomAdj(app);
 
             assertNotEquals(PERCEPTIBLE_RECENT_FOREGROUND_APP_ADJ + 2, app.mState.getSetAdj());
 
-            mService.mOomAdjuster.resetInternal();
+            mProcessStateController.getOomAdjuster().resetInternal();
         }
     }
 
@@ -800,12 +807,12 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, true));
-        system.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        system.mState.setHasTopUi(true);
+        mProcessStateController.setMaxAdj(system, PERSISTENT_PROC_ADJ);
+        mProcessStateController.setHasTopUi(system, true);
         // Simulate the system starting and binding to a service in the app.
         ServiceRecord s = bindService(app, system,
                 null, null, Context.BIND_ALMOST_PERCEPTIBLE, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(system, app);
 
         assertProcStates(app, PROCESS_STATE_IMPORTANT_FOREGROUND,
@@ -817,8 +824,8 @@
     public void testUpdateOomAdj_DoOne_Toast() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mState.setForcingToImportant(new Object());
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setForcingToImportant(app, new Object());
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_TRANSIENT_BACKGROUND, PERCEPTIBLE_APP_ADJ,
@@ -832,7 +839,7 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         WindowProcessController wpc = app.getWindowProcessController();
         doReturn(true).when(wpc).isHeavyWeightProcess();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doReturn(false).when(wpc).isHeavyWeightProcess();
 
@@ -847,7 +854,7 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         WindowProcessController wpc = app.getWindowProcessController();
         doReturn(true).when(wpc).isHomeProcess();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_HOME, HOME_APP_ADJ, SCHED_GROUP_BACKGROUND);
@@ -861,7 +868,7 @@
         WindowProcessController wpc = app.getWindowProcessController();
         doReturn(true).when(wpc).isPreviousProcess();
         doReturn(true).when(wpc).hasActivities();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_LAST_ACTIVITY, PREVIOUS_APP_ADJ,
@@ -873,7 +880,7 @@
         verify(mService.mHandler).sendEmptyMessageAtTime(eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG),
                 followUpTimeCaptor.capture());
         mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
-        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+        mProcessStateController.runFollowUpUpdate();
 
         int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
                 ? sFirstUiCachedAdj : CACHED_APP_MIN_ADJ;
@@ -896,9 +903,9 @@
             doReturn(true).when(wpc).isPreviousProcess();
             doReturn(true).when(wpc).hasActivities();
         }
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         setProcessesToLru(apps);
-        mService.mOomAdjuster.updateOomAdjLocked(OOM_ADJ_REASON_NONE);
+        mProcessStateController.runFullUpdate(OOM_ADJ_REASON_NONE);
 
         for (int i = 0; i < numberOfApps; i++) {
             assertProcStates(apps[i], PROCESS_STATE_LAST_ACTIVITY, PREVIOUS_APP_ADJ,
@@ -914,7 +921,7 @@
             mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
         }
 
-        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+        mProcessStateController.runFollowUpUpdate();
 
         for (int i = 0; i < numberOfApps; i++) {
             final int mruIndex = numberOfApps - i - 1;
@@ -938,10 +945,8 @@
     public void testUpdateOomAdj_DoOne_Backup() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        BackupRecord backupTarget = new BackupRecord(null, 0, 0, 0);
-        backupTarget.app = app;
-        doReturn(backupTarget).when(mService.mBackupTargets).get(anyInt());
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setBackupTarget(app);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
         doReturn(null).when(mService.mBackupTargets).get(anyInt());
 
@@ -954,8 +959,8 @@
     public void testUpdateOomAdj_DoOne_ClientActivities() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mServices.setHasClientActivities(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasClientActivities(app.mServices, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertEquals(PROCESS_STATE_CACHED_ACTIVITY_CLIENT, app.mState.getSetProcState());
@@ -966,8 +971,8 @@
     public void testUpdateOomAdj_DoOne_TreatLikeActivity() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
-        app.mServices.setTreatLikeActivity(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setTreatLikeActivity(app.mServices, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertEquals(PROCESS_STATE_CACHED_ACTIVITY, app.mState.getSetProcState());
@@ -981,10 +986,10 @@
         app.mState.setServiceB(true);
         ServiceRecord s = mock(ServiceRecord.class);
         doReturn(new ArrayMap<IBinder, ArrayList<ConnectionRecord>>()).when(s).getConnections();
-        s.startRequested = true;
-        s.lastActivity = SystemClock.uptimeMillis();
-        app.mServices.startService(s);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, SystemClock.uptimeMillis());
+        mProcessStateController.startService(app.mServices, s);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_B_ADJ, SCHED_GROUP_BACKGROUND);
@@ -995,8 +1000,8 @@
     public void testUpdateOomAdj_DoOne_MaxAdj() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
-        app.mState.setMaxAdj(PERCEPTIBLE_LOW_APP_ADJ);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setMaxAdj(app, PERCEPTIBLE_LOW_APP_ADJ);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_CACHED_EMPTY, PERCEPTIBLE_LOW_APP_ADJ,
@@ -1011,7 +1016,7 @@
         app.mState.setCurRawAdj(SERVICE_ADJ);
         app.mState.setCurAdj(SERVICE_ADJ);
         doReturn(null).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertTrue(ProcessList.CACHED_APP_MIN_ADJ <= app.mState.getSetAdj());
@@ -1025,10 +1030,10 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         ServiceRecord s = mock(ServiceRecord.class);
         doReturn(new ArrayMap<IBinder, ArrayList<ConnectionRecord>>()).when(s).getConnections();
-        s.startRequested = true;
-        s.lastActivity = SystemClock.uptimeMillis();
-        app.mServices.startService(s);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, SystemClock.uptimeMillis());
+        mProcessStateController.startService(app.mServices, s);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND);
@@ -1043,8 +1048,8 @@
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         ServiceRecord s = bindService(app, client, null, null, Context.BIND_WAIVE_PRIORITY,
                 mock(IBinder.class));
-        s.startRequested = true;
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(client).when(mService).getTopApp();
         updateOomAdj(client, app);
@@ -1062,10 +1067,10 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
-        client.mServices.setTreatLikeActivity(true);
+        mProcessStateController.setTreatLikeActivity(client.mServices, true);
         bindService(app, client, null, null, Context.BIND_WAIVE_PRIORITY
                 | Context.BIND_TREAT_LIKE_ACTIVITY, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PROCESS_STATE_CACHED_ACTIVITY, app.mState.getSetProcState());
@@ -1086,7 +1091,7 @@
                 mock(ActivityServiceConnectionsHolder.class));
         doReturn(client).when(mService).getTopApp();
         doReturn(true).when(cr.activity).isActivityVisible();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(FOREGROUND_APP_ADJ, app.mState.getSetAdj());
@@ -1099,7 +1104,7 @@
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         bindService(app, app, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         final int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
@@ -1114,9 +1119,9 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
-        client.mServices.setTreatLikeActivity(true);
+        mProcessStateController.setTreatLikeActivity(client.mServices, true);
         bindService(app, client, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PROCESS_STATE_CACHED_EMPTY, app.mState.getSetProcState());
@@ -1137,7 +1142,7 @@
                 mock(IBinder.class));
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(client).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
         doReturn(null).when(mService).getTopApp();
 
@@ -1152,9 +1157,9 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_FOREGROUND_SERVICE, mock(IBinder.class));
-        client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        client.mState.setHasTopUi(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+        mProcessStateController.setHasTopUi(client, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, VISIBLE_APP_ADJ,
@@ -1170,8 +1175,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_IMPORTANT, mock(IBinder.class));
-        client.mServices.startExecutingService(mock(ServiceRecord.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.startExecutingService(client.mServices, mock(ServiceRecord.class));
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(FOREGROUND_APP_ADJ, app.mState.getSetAdj());
@@ -1188,7 +1193,7 @@
         bindService(app, client, null, null, 0, mock(IBinder.class));
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(client).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
         doReturn(null).when(mService).getTopApp();
 
@@ -1203,8 +1208,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_FOREGROUND_SERVICE, mock(IBinder.class));
-        client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PROCESS_STATE_BOUND_FOREGROUND_SERVICE, app.mState.getSetProcState());
@@ -1222,9 +1227,9 @@
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_FOREGROUND_SERVICE, mock(IBinder.class));
         client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_ASLEEP);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_ASLEEP);
         updateOomAdj(client, app);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, VISIBLE_APP_ADJ,
                 SCHED_GROUP_RESTRICTED);
@@ -1240,8 +1245,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_NOT_FOREGROUND, mock(IBinder.class));
-        client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PROCESS_STATE_TRANSIENT_BACKGROUND, app.mState.getSetProcState());
@@ -1256,8 +1261,9 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, 0, mock(IBinder.class));
-        client.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client.mServices, true,
+                0, /* hasNoneType=*/true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PROCESS_STATE_FOREGROUND_SERVICE, client.mState.getSetProcState());
@@ -1279,16 +1285,16 @@
         // In order to trick OomAdjuster to think it has a short-service, we need this logic.
         ServiceRecord s = ServiceRecord.newEmptyInstanceForTest(mService);
         s.appInfo = new ApplicationInfo();
-        s.startRequested = true;
-        s.isForeground = true;
-        s.foregroundServiceType = FOREGROUND_SERVICE_TYPE_SHORT_SERVICE;
-        s.setShortFgsInfo(SystemClock.uptimeMillis());
-        client.mServices.startService(s);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setIsForegroundService(s, true);
+        mProcessStateController.setForegroundServiceType(s, FOREGROUND_SERVICE_TYPE_SHORT_SERVICE);
+        mProcessStateController.setShortFgsInfo(s, SystemClock.uptimeMillis());
+        mProcessStateController.startService(client.mServices, s);
         client.mState.setLastTopTime(SystemClock.uptimeMillis());
 
-        client.mServices.setHasForegroundServices(true, FOREGROUND_SERVICE_TYPE_SHORT_SERVICE,
-                /* hasNoneType=*/false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client.mServices, true,
+                FOREGROUND_SERVICE_TYPE_SHORT_SERVICE, /* hasNoneType=*/false);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         // Client only has a SHORT_FGS, so it doesn't have BFSL, and that's propagated.
@@ -1310,16 +1316,16 @@
         // In order to trick OomAdjuster to think it has a short-service, we need this logic.
         ServiceRecord s = ServiceRecord.newEmptyInstanceForTest(mService);
         s.appInfo = new ApplicationInfo();
-        s.startRequested = true;
-        s.isForeground = true;
-        s.foregroundServiceType = FOREGROUND_SERVICE_TYPE_SHORT_SERVICE;
-        s.setShortFgsInfo(SystemClock.uptimeMillis());
-        app2.mServices.startService(s);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setIsForegroundService(s, true);
+        mProcessStateController.setForegroundServiceType(s, FOREGROUND_SERVICE_TYPE_SHORT_SERVICE);
+        mProcessStateController.setShortFgsInfo(s, SystemClock.uptimeMillis());
+        mProcessStateController.startService(app2.mServices, s);
         app2.mState.setLastTopTime(SystemClock.uptimeMillis());
 
-        app2.mServices.setHasForegroundServices(true, FOREGROUND_SERVICE_TYPE_SHORT_SERVICE,
-                /* hasNoneType=*/false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(app2.mServices, true,
+                FOREGROUND_SERVICE_TYPE_SHORT_SERVICE, /* hasNoneType=*/false);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app2);
 
         // Client only has a SHORT_FGS, so it doesn't have BFSL, and that's propagated.
@@ -1331,7 +1337,7 @@
         // Persistent process
         ProcessRecord pers = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
-        pers.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(pers, PERSISTENT_PROC_ADJ);
 
         // app1, which is bound by pers (which makes it BFGS)
         ProcessRecord app1 = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
@@ -1358,18 +1364,16 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_ABOVE_CLIENT, mock(IBinder.class));
-        BackupRecord backupTarget = new BackupRecord(null, 0, 0, 0);
-        backupTarget.app = client;
-        doReturn(backupTarget).when(mService.mBackupTargets).get(anyInt());
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setBackupTarget(client);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
-        doReturn(null).when(mService.mBackupTargets).get(anyInt());
+        mProcessStateController.stopBackupTarget(UserHandle.getUserId(MOCKAPP2_UID));
 
         assertEquals(BACKUP_APP_ADJ, app.mState.getSetAdj());
         assertNoBfsl(app);
 
-        client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
         updateOomAdj(client, app);
 
         assertEquals(PERSISTENT_SERVICE_ADJ, app.mState.getSetAdj());
@@ -1384,8 +1388,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_NOT_PERCEPTIBLE, mock(IBinder.class));
-        client.mState.setRunningRemoteAnimation(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setRunningRemoteAnimation(client, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PERCEPTIBLE_LOW_APP_ADJ, app.mState.getSetAdj());
@@ -1399,8 +1403,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_NOT_VISIBLE, mock(IBinder.class));
-        client.mState.setRunningRemoteAnimation(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setRunningRemoteAnimation(client, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PERCEPTIBLE_APP_ADJ, app.mState.getSetAdj());
@@ -1414,8 +1418,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, 0, mock(IBinder.class));
-        client.mState.setHasOverlayUi(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasOverlayUi(client, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PERCEPTIBLE_APP_ADJ, app.mState.getSetAdj());
@@ -1432,13 +1436,13 @@
             bindService(app, client, null, null,
                     Context.BIND_ALMOST_PERCEPTIBLE | Context.BIND_NOT_FOREGROUND,
                     mock(IBinder.class));
-            client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
             updateOomAdj(client, app);
 
             assertEquals(PERCEPTIBLE_MEDIUM_APP_ADJ + 2, app.mState.getSetAdj());
 
-            mService.mOomAdjuster.resetInternal();
+            mProcessStateController.getOomAdjuster().resetInternal();
         }
 
         {
@@ -1451,14 +1455,14 @@
             bindService(app, client, null, null,
                     Context.BIND_ALMOST_PERCEPTIBLE | Context.BIND_NOT_FOREGROUND,
                     mock(IBinder.class));
-            client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
             updateOomAdj(client, app);
             doReturn(false).when(wpc).isHeavyWeightProcess();
 
             assertEquals(PERCEPTIBLE_MEDIUM_APP_ADJ + 2, app.mState.getSetAdj());
 
-            mService.mOomAdjuster.resetInternal();
+            mProcessStateController.getOomAdjuster().resetInternal();
         }
 
         {
@@ -1469,13 +1473,13 @@
             bindService(app, client, null, null,
                     Context.BIND_ALMOST_PERCEPTIBLE,
                     mock(IBinder.class));
-            client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
             updateOomAdj(client, app);
 
             assertEquals(PERCEPTIBLE_APP_ADJ + 1, app.mState.getSetAdj());
 
-            mService.mOomAdjuster.resetInternal();
+            mProcessStateController.getOomAdjuster().resetInternal();
         }
 
         {
@@ -1488,14 +1492,14 @@
             bindService(app, client, null, null,
                     Context.BIND_ALMOST_PERCEPTIBLE,
                     mock(IBinder.class));
-            client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-            mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+            mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+            setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
             updateOomAdj(client, app);
             doReturn(false).when(wpc).isHeavyWeightProcess();
 
             assertEquals(PERCEPTIBLE_APP_ADJ + 1, app.mState.getSetAdj());
 
-            mService.mOomAdjuster.resetInternal();
+            mProcessStateController.getOomAdjuster().resetInternal();
         }
     }
 
@@ -1507,8 +1511,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, 0, mock(IBinder.class));
-        client.mState.setRunningRemoteAnimation(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setRunningRemoteAnimation(client, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(VISIBLE_APP_ADJ, app.mState.getSetAdj());
@@ -1523,8 +1527,8 @@
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindService(app, client, null, null, Context.BIND_IMPORTANT_BACKGROUND,
                 mock(IBinder.class));
-        client.mState.setHasOverlayUi(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasOverlayUi(client, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertEquals(PROCESS_STATE_IMPORTANT_BACKGROUND, app.mState.getSetProcState());
@@ -1552,8 +1556,8 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindProvider(app, client, null, null, false);
-        client.mServices.setTreatLikeActivity(true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setTreatLikeActivity(client.mServices, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client);
 
         final int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
@@ -1571,7 +1575,7 @@
         bindProvider(app, client, null, null, false);
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(client).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
         doReturn(null).when(mService).getTopApp();
 
@@ -1585,9 +1589,9 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
-        client.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(client.mServices, true, 0, true);
         bindProvider(app, client, null, null, false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1607,17 +1611,17 @@
         // In order to trick OomAdjuster to think it has a short-service, we need this logic.
         ServiceRecord s = ServiceRecord.newEmptyInstanceForTest(mService);
         s.appInfo = new ApplicationInfo();
-        s.startRequested = true;
+        mProcessStateController.setStartRequested(s, true);
         s.isForeground = true;
         s.foregroundServiceType = FOREGROUND_SERVICE_TYPE_SHORT_SERVICE;
-        s.setShortFgsInfo(SystemClock.uptimeMillis());
-        client.mServices.startService(s);
+        mProcessStateController.setShortFgsInfo(s, SystemClock.uptimeMillis());
+        mProcessStateController.startService(client.mServices, s);
         client.mState.setLastTopTime(SystemClock.uptimeMillis());
 
-        client.mServices.setHasForegroundServices(true, FOREGROUND_SERVICE_TYPE_SHORT_SERVICE,
-                /* hasNoneType=*/false);
+        mProcessStateController.setHasForegroundServices(client.mServices, true,
+                FOREGROUND_SERVICE_TYPE_SHORT_SERVICE, false);
         bindProvider(app, client, null, null, false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         // Client only has a SHORT_FGS, so it doesn't have BFSL, and that's propagated.
@@ -1639,7 +1643,7 @@
         ProcessRecord client = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         bindProvider(app, client, null, null, true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, app);
 
         assertProcStates(app, PROCESS_STATE_IMPORTANT_FOREGROUND, FOREGROUND_APP_ADJ,
@@ -1652,7 +1656,7 @@
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         app.mProviders.setLastProviderTime(SystemClock.uptimeMillis());
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_LAST_ACTIVITY, PREVIOUS_APP_ADJ,
@@ -1664,7 +1668,7 @@
         verify(mService.mHandler).sendEmptyMessageAtTime(eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG),
                 followUpTimeCaptor.capture());
         mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
-        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+        mProcessStateController.runFollowUpUpdate();
 
         final int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
                 ? sFirstNonUiCachedAdj : sFirstCachedAdj;
@@ -1688,7 +1692,7 @@
         bindService(client, client2, null, null, 0, mock(IBinder.class));
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(client2).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, client2, app);
         doReturn(null).when(mService).getTopApp();
 
@@ -1707,8 +1711,8 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(app, client2, null, null, 0, mock(IBinder.class));
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, client2, app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1727,8 +1731,8 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(client, client2, null, null, 0, mock(IBinder.class));
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, client2, app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1747,13 +1751,13 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(client, client2, null, null, 0, mock(IBinder.class));
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
         bindService(client2, app, null, null, 0, mock(IBinder.class));
 
         // Note: We add processes to LRU but still call updateOomAdjLocked() with a specific
         // processes.
         setProcessesToLru(app, client, client2);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1766,8 +1770,8 @@
         assertBfsl(client);
         assertBfsl(client2);
 
-        client2.mServices.setHasForegroundServices(false, 0, /* hasNoneType=*/false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client2.mServices, false, 0, false);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client2);
 
         assertEquals(PROCESS_STATE_CACHED_EMPTY, client2.mState.getSetProcState());
@@ -1790,8 +1794,8 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(client2, client, null, null, 0, mock(IBinder.class));
-        client.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1817,8 +1821,8 @@
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(client, client2, null, null, 0, mock(IBinder.class));
         bindService(client2, client, null, null, 0, mock(IBinder.class));
-        client.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1851,8 +1855,8 @@
                 MOCKAPP5_PROCESSNAME, MOCKAPP5_PACKAGENAME, false));
         bindService(client3, client4, null, null, 0, mock(IBinder.class));
         bindService(client4, client3, null, null, 0, mock(IBinder.class));
-        client.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2, client3, client4);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1883,13 +1887,13 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(client, client2, null, null, 0, mock(IBinder.class));
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
         bindService(client2, app, null, null, 0, mock(IBinder.class));
         ProcessRecord client3 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        client3.mState.setForcingToImportant(new Object());
+        mProcessStateController.setForcingToImportant(client3, new Object());
         bindService(app, client3, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2, client3);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1913,9 +1917,9 @@
         doReturn(true).when(wpc).isHomeProcess();
         ProcessRecord client3 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        client3.mState.setForcingToImportant(new Object());
+        mProcessStateController.setForcingToImportant(client3, new Object());
         bindService(app, client3, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2, client3);
 
         assertProcStates(app, PROCESS_STATE_TRANSIENT_BACKGROUND, PERCEPTIBLE_APP_ADJ,
@@ -1940,9 +1944,9 @@
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
         ProcessRecord client4 = spy(makeDefaultProcessRecord(MOCKAPP5_PID, MOCKAPP5_UID,
                 MOCKAPP5_PROCESSNAME, MOCKAPP5_PACKAGENAME, false));
-        client4.mState.setForcingToImportant(new Object());
+        mProcessStateController.setForcingToImportant(client4, new Object());
         bindService(app, client4, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2, client3, client4);
 
         assertProcStates(app, PROCESS_STATE_TRANSIENT_BACKGROUND, PERCEPTIBLE_APP_ADJ,
@@ -1965,13 +1969,13 @@
         doReturn(true).when(wpc).isHomeProcess();
         ProcessRecord client3 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        client3.mState.setForcingToImportant(new Object());
+        mProcessStateController.setForcingToImportant(client3, new Object());
         bindService(app, client3, null, null, 0, mock(IBinder.class));
         ProcessRecord client4 = spy(makeDefaultProcessRecord(MOCKAPP5_PID, MOCKAPP5_UID,
                 MOCKAPP5_PROCESSNAME, MOCKAPP5_PACKAGENAME, false));
-        client4.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(client4.mServices, true, 0, true);
         bindService(app, client4, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2, client3, client4);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -1992,12 +1996,12 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(app, client2, null, null, 0, mock(IBinder.class));
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
         ProcessRecord client3 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        client3.mState.setForcingToImportant(new Object());
+        mProcessStateController.setForcingToImportant(client3, new Object());
         bindService(app, client3, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, client2, client3, app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2016,8 +2020,8 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindProvider(client, client2, null, null, false);
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, client2, app);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2036,9 +2040,9 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindProvider(client, client2, null, null, false);
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
         bindService(client2, app, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2057,8 +2061,8 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindProvider(client, client2, null, null, false);
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client, client2, app);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2077,9 +2081,9 @@
         ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindProvider(client, client2, null, null, false);
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
         bindProvider(client2, app, null, null, false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2102,10 +2106,10 @@
                 mock(IBinder.class));
         bindService(app2, client2, null, null, Context.BIND_FOREGROUND_SERVICE_WHILE_AWAKE,
                 mock(IBinder.class));
-        client1.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        client2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setMaxAdj(client1, PERSISTENT_PROC_ADJ);
+        mProcessStateController.setHasForegroundServices(client2.mServices, true, 0, true);
 
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(client1, client2, app1, app2);
 
         assertProcStates(app1, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, VISIBLE_APP_ADJ,
@@ -2126,7 +2130,7 @@
         assertProcStates(app2, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
                 SCHED_GROUP_DEFAULT);
 
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_ASLEEP);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_ASLEEP);
         updateOomAdj(client1, client2, app1, app2);
         assertProcStates(app1, PROCESS_STATE_IMPORTANT_FOREGROUND, VISIBLE_APP_ADJ,
                 SCHED_GROUP_TOP_APP);
@@ -2136,7 +2140,7 @@
 
         bindService(client2, app1, null, null, 0, mock(IBinder.class));
         bindService(app1, client2, null, null, 0, mock(IBinder.class));
-        client2.mServices.setHasForegroundServices(false, 0, /* hasNoneType=*/false);
+        mProcessStateController.setHasForegroundServices(client2.mServices, false, 0, false);
         updateOomAdj(app1, client1, client2);
         assertProcStates(app1, PROCESS_STATE_IMPORTANT_FOREGROUND, VISIBLE_APP_ADJ,
                 SCHED_GROUP_TOP_APP);
@@ -2153,8 +2157,8 @@
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         final ProcessRecord client2 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        client1.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        client2.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(client1, PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(client2, PERSISTENT_PROC_ADJ);
 
         final ServiceRecord s1 = bindService(app1, client1, null, null,
                 Context.BIND_TREAT_LIKE_VISIBLE_FOREGROUND_SERVICE, mock(IBinder.class));
@@ -2178,10 +2182,10 @@
         s2.getConnections().clear();
         client1.mServices.removeAllConnections();
         client2.mServices.removeAllConnections();
-        client1.mState.setMaxAdj(UNKNOWN_ADJ);
-        client2.mState.setMaxAdj(UNKNOWN_ADJ);
-        client1.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        client2.mState.setHasOverlayUi(true);
+        mProcessStateController.setMaxAdj(client1, UNKNOWN_ADJ);
+        mProcessStateController.setMaxAdj(client2, UNKNOWN_ADJ);
+        mProcessStateController.setHasForegroundServices(client1.mServices, true, 0, true);
+        mProcessStateController.setHasOverlayUi(client2, true);
 
         bindService(app1, client1, null, s1, Context.BIND_TREAT_LIKE_VISIBLE_FOREGROUND_SERVICE,
                 mock(IBinder.class));
@@ -2196,10 +2200,10 @@
         assertProcStates(app2, PROCESS_STATE_IMPORTANT_FOREGROUND, PERCEPTIBLE_APP_ADJ,
                 SCHED_GROUP_DEFAULT);
 
-        client2.mState.setHasOverlayUi(false);
+        mProcessStateController.setHasOverlayUi(client2, false);
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(client2).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
         updateOomAdj(client2, app2);
         assertProcStates(app2, PROCESS_STATE_BOUND_TOP, VISIBLE_APP_ADJ,
@@ -2213,10 +2217,10 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         final ProcessRecord client1 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
-        client1.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(client1, PERSISTENT_PROC_ADJ);
 
-        app1.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(app1.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
         bindService(app1, client1, null, null, Context.BIND_NOT_PERCEPTIBLE, mock(IBinder.class));
 
@@ -2234,10 +2238,10 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
         final ProcessRecord client1 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
-        client1.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(client1, PERSISTENT_PROC_ADJ);
 
-        app1.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setHasForegroundServices(app1.mServices, true, 0, true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
         bindService(app1, client1, null, null, Context.BIND_ALMOST_PERCEPTIBLE,
                 mock(IBinder.class));
@@ -2254,10 +2258,10 @@
     public void testUpdateOomAdj_DoOne_PendingFinishAttach() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
-        app.setPendingFinishAttach(true);
+        mProcessStateController.setPendingFinishAttach(app, true);
         app.mState.setHasForegroundActivities(false);
 
-        mService.mOomAdjuster.setAttachingProcessStatesLSP(app);
+        mProcessStateController.setAttachingProcessStatesLSP(app);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_CACHED_EMPTY, FOREGROUND_APP_ADJ,
@@ -2269,11 +2273,11 @@
     public void testUpdateOomAdj_DoOne_TopApp_PendingFinishAttach() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
-        app.setPendingFinishAttach(true);
+        mProcessStateController.setPendingFinishAttach(app, true);
         app.mState.setHasForegroundActivities(true);
         doReturn(app).when(mService).getTopApp();
 
-        mService.mOomAdjuster.setAttachingProcessStatesLSP(app);
+        mProcessStateController.setAttachingProcessStatesLSP(app);
         updateOomAdj(app);
 
         assertProcStates(app, PROCESS_STATE_TOP, FOREGROUND_APP_ADJ,
@@ -2303,40 +2307,40 @@
         client1.setUidRecord(clientUidRecord);
         client2.setUidRecord(clientUidRecord);
 
-        client1.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
-        client2.mState.setForcingToImportant(new Object());
+        mProcessStateController.setHasForegroundServices(client1.mServices, true, 0, true);
+        mProcessStateController.setForcingToImportant(client2, new Object());
         setProcessesToLru(app1, app2, app3, client1, client2);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
 
         final ComponentName cn1 = ComponentName.unflattenFromString(
                 MOCKAPP_PACKAGENAME + "/.TestService");
         final ServiceRecord s1 = bindService(app1, client1, null, null, 0, mock(IBinder.class));
         setFieldValue(ServiceRecord.class, s1, "name", cn1);
-        s1.startRequested = true;
+        mProcessStateController.setStartRequested(s1, true);
 
         final ComponentName cn2 = ComponentName.unflattenFromString(
                 MOCKAPP2_PACKAGENAME + "/.TestService");
         final ServiceRecord s2 = bindService(app2, client2, null, null, 0, mock(IBinder.class));
         setFieldValue(ServiceRecord.class, s2, "name", cn2);
-        s2.startRequested = true;
+        mProcessStateController.setStartRequested(s2, true);
 
         final ComponentName cn3 = ComponentName.unflattenFromString(
                 MOCKAPP5_PACKAGENAME + "/.TestService");
         final ServiceRecord s3 = bindService(app3, client1, null, null, 0, mock(IBinder.class));
         setFieldValue(ServiceRecord.class, s3, "name", cn3);
-        s3.startRequested = true;
+        mProcessStateController.setStartRequested(s3, true);
 
         final ComponentName cn4 = ComponentName.unflattenFromString(
                 MOCKAPP3_PACKAGENAME + "/.TestService");
         final ServiceRecord c2s = makeServiceRecord(client2);
         setFieldValue(ServiceRecord.class, c2s, "name", cn4);
-        c2s.startRequested = true;
+        mProcessStateController.setStartRequested(c2s, true);
 
         try {
-            mService.mOomAdjuster.mActiveUids.put(MOCKAPP_UID, app1UidRecord);
-            mService.mOomAdjuster.mActiveUids.put(MOCKAPP2_UID, app2UidRecord);
-            mService.mOomAdjuster.mActiveUids.put(MOCKAPP5_UID, app3UidRecord);
-            mService.mOomAdjuster.mActiveUids.put(MOCKAPP3_UID, clientUidRecord);
+            mActiveUids.put(MOCKAPP_UID, app1UidRecord);
+            mActiveUids.put(MOCKAPP2_UID, app2UidRecord);
+            mActiveUids.put(MOCKAPP5_UID, app3UidRecord);
+            mActiveUids.put(MOCKAPP3_UID, clientUidRecord);
 
             setServiceMap(s1, MOCKAPP_UID, cn1);
             setServiceMap(s2, MOCKAPP2_UID, cn2);
@@ -2354,8 +2358,8 @@
             assertEquals(PROCESS_STATE_TRANSIENT_BACKGROUND, app2.mState.getSetProcState());
             assertEquals(PROCESS_STATE_TRANSIENT_BACKGROUND, client2.mState.getSetProcState());
 
-            client1.mServices.setHasForegroundServices(false, 0, /* hasNoneType=*/false);
-            client2.mState.setForcingToImportant(null);
+            mProcessStateController.setHasForegroundServices(client1.mServices, false, 0, false);
+            mProcessStateController.setForcingToImportant(client2, null);
             app1UidRecord.reset();
             app2UidRecord.reset();
             app3UidRecord.reset();
@@ -2379,7 +2383,7 @@
                     .getAppStartModeLOSP(anyInt(), any(String.class), anyInt(),
                             anyInt(), anyBoolean(), anyBoolean(), anyBoolean());
             mService.mServices.mServiceMap.clear();
-            mService.mOomAdjuster.mActiveUids.clear();
+            mActiveUids.clear();
         }
     }
 
@@ -2388,12 +2392,13 @@
     public void testUpdateOomAdj_DoAll_Unbound() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
-        app.mState.setForcingToImportant(new Object());
         ProcessRecord app2 = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
-        app2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
 
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setForcingToImportant(app, new Object());
+        mProcessStateController.setHasForegroundServices(app2.mServices, true, 0, /* hasNoneType=*/
+                true);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, app2);
 
         assertProcStates(app, PROCESS_STATE_TRANSIENT_BACKGROUND, PERCEPTIBLE_APP_ADJ,
@@ -2408,12 +2413,14 @@
     public void testUpdateOomAdj_DoAll_BoundFgService() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
-        app.mState.setForcingToImportant(new Object());
         ProcessRecord app2 = spy(makeDefaultProcessRecord(MOCKAPP2_PID, MOCKAPP2_UID,
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
-        app2.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+
+        mProcessStateController.setForcingToImportant(app, new Object());
+        mProcessStateController.setHasForegroundServices(app2.mServices, true, 0, /* hasNoneType=*/
+                true);
         bindService(app, app2, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, app2);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2435,9 +2442,9 @@
         ProcessRecord app3 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         bindService(app2, app3, null, null, 0, mock(IBinder.class));
-        app3.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(app3.mServices, true, 0, true);
         bindService(app3, app, null, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, app2, app3);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2476,13 +2483,13 @@
         doReturn(true).when(wpc).isHomeProcess();
         ProcessRecord app4 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        app4.mState.setHasOverlayUi(true);
+        mProcessStateController.setHasOverlayUi(app4, true);
         bindService(app, app4, null, s, 0, mock(IBinder.class));
         ProcessRecord app5 = spy(makeDefaultProcessRecord(MOCKAPP5_PID, MOCKAPP5_UID,
                 MOCKAPP5_PROCESSNAME, MOCKAPP5_PACKAGENAME, false));
-        app5.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(app5.mServices, true, 0, true);
         bindService(app, app5, null, s, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, app2, app3, app4, app5);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2518,13 +2525,13 @@
         doReturn(true).when(wpc).isHomeProcess();
         ProcessRecord app4 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        app4.mState.setHasOverlayUi(true);
+        mProcessStateController.setHasOverlayUi(app4, true);
         bindService(app, app4, null, s, 0, mock(IBinder.class));
         ProcessRecord app5 = spy(makeDefaultProcessRecord(MOCKAPP5_PID, MOCKAPP5_UID,
                 MOCKAPP5_PROCESSNAME, MOCKAPP5_PACKAGENAME, false));
-        app5.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(app5.mServices, true, 0, true);
         bindService(app, app5, null, s, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app5, app4, app3, app2, app);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2560,13 +2567,13 @@
         doReturn(true).when(wpc).isHomeProcess();
         ProcessRecord app4 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        app4.mState.setHasOverlayUi(true);
+        mProcessStateController.setHasOverlayUi(app4, true);
         bindService(app, app4, null, s, 0, mock(IBinder.class));
         ProcessRecord app5 = spy(makeDefaultProcessRecord(MOCKAPP5_PID, MOCKAPP5_UID,
                 MOCKAPP5_PROCESSNAME, MOCKAPP5_PACKAGENAME, false));
-        app5.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(app5.mServices, true, 0, true);
         bindService(app, app5, null, s, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app3, app4, app2, app, app5);
 
         assertProcStates(app, PROCESS_STATE_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2603,10 +2610,10 @@
                 mock(IBinder.class));
         ProcessRecord client3 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        client3.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
+        mProcessStateController.setMaxAdj(client3, PERSISTENT_PROC_ADJ);
         bindService(app, client3, null, null, Context.BIND_INCLUDE_CAPABILITIES,
                 mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, client, client2, client3);
 
         final int expected = PROCESS_CAPABILITY_ALL & ~PROCESS_CAPABILITY_BFSL;
@@ -2631,13 +2638,13 @@
         doReturn(true).when(wpc).isHomeProcess();
         ProcessRecord app4 = spy(makeDefaultProcessRecord(MOCKAPP4_PID, MOCKAPP4_UID,
                 MOCKAPP4_PROCESSNAME, MOCKAPP4_PACKAGENAME, false));
-        app4.mState.setHasOverlayUi(true);
+        mProcessStateController.setHasOverlayUi(app4, true);
         bindProvider(app, app4, cr, null, false);
         ProcessRecord app5 = spy(makeDefaultProcessRecord(MOCKAPP5_PID, MOCKAPP5_UID,
                 MOCKAPP5_PROCESSNAME, MOCKAPP5_PACKAGENAME, false));
-        app5.mServices.setHasForegroundServices(true, 0, /* hasNoneType=*/true);
+        mProcessStateController.setHasForegroundServices(app5.mServices, true, 0, true);
         bindProvider(app, app5, cr, null, false);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app, app2, app3, app4, app5);
 
         assertProcStates(app, PROCESS_STATE_BOUND_FOREGROUND_SERVICE, PERCEPTIBLE_APP_ADJ,
@@ -2666,22 +2673,22 @@
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         long now = SystemClock.uptimeMillis();
         ServiceRecord s = bindService(app, app2, null, null, 0, mock(IBinder.class));
-        s.startRequested = true;
-        s.lastActivity = now;
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, now);
         s = bindService(app2, app, null, null, 0, mock(IBinder.class));
-        s.startRequested = true;
-        s.lastActivity = now;
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, now);
         ProcessRecord app3 = spy(makeDefaultProcessRecord(MOCKAPP3_PID, MOCKAPP3_UID,
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         s = mock(ServiceRecord.class);
-        s.app = app3;
+        mProcessStateController.setHostProcess(s, app3);
         setFieldValue(ServiceRecord.class, s, "connections",
                 new ArrayMap<IBinder, ArrayList<ConnectionRecord>>());
-        app3.mServices.startService(s);
+        mProcessStateController.startService(app3.mServices, s);
         doCallRealMethod().when(s).getConnections();
-        s.startRequested = true;
-        s.lastActivity = now;
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, now);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         mService.mOomAdjuster.mNumServiceProcs = 3;
         updateOomAdj(app3, app2, app);
 
@@ -2702,8 +2709,8 @@
 
         // cachedAdj1 and cachedAdj2 will be read if USE_TIERED_CACHED_ADJ is disabled. Otherwise,
         // sFirstUiCachedAdj and sFirstNonUiCachedAdj are used instead.
-        final int cachedAdj1 = CACHED_APP_MIN_ADJ + ProcessList.CACHED_APP_IMPORTANCE_LEVELS;
-        final int cachedAdj2 = cachedAdj1 + ProcessList.CACHED_APP_IMPORTANCE_LEVELS * 2;
+        final int cachedAdj1 = CACHED_APP_MIN_ADJ + CACHED_APP_IMPORTANCE_LEVELS;
+        final int cachedAdj2 = cachedAdj1 + CACHED_APP_IMPORTANCE_LEVELS * 2;
         doReturn(userOwner).when(mService.mUserController).getCurrentUserId();
 
         final ArrayList<ProcessRecord> lru = mService.mProcessList.getLruProcessesLOSP();
@@ -2723,11 +2730,11 @@
         ServiceRecord s = spy(new ServiceRecord(mService, cn, cn, null, 0, null,
                 si, false, null));
         doReturn(new ArrayMap<IBinder, ArrayList<ConnectionRecord>>()).when(s).getConnections();
-        s.startRequested = true;
-        s.lastActivity = now;
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, now);
 
-        app.mServices.startService(s);
-        app.mState.setHasShownUi(true);
+        mProcessStateController.startService(app.mServices, s);
+        mProcessStateController.setHasShownUi(app, true);
 
         final ServiceInfo si2 = mock(ServiceInfo.class);
         si2.applicationInfo = mock(ApplicationInfo.class);
@@ -2735,13 +2742,14 @@
         ServiceRecord s2 = spy(new ServiceRecord(mService, cn2, cn2, null, 0, null,
                 si2, false, null));
         doReturn(new ArrayMap<IBinder, ArrayList<ConnectionRecord>>()).when(s2).getConnections();
-        s2.startRequested = true;
-        s2.lastActivity = now - mService.mConstants.MAX_SERVICE_INACTIVITY - 1;
+        mProcessStateController.setStartRequested(s2, true);
+        mProcessStateController.setServiceLastActivityTime(s2,
+                now - mService.mConstants.MAX_SERVICE_INACTIVITY - 1);
 
-        app2.mServices.startService(s2);
-        app2.mState.setHasShownUi(false);
+        mProcessStateController.startService(app2.mServices, s2);
+        mProcessStateController.setHasShownUi(app2, false);
 
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj();
 
         assertProcStates(app, PROCESS_STATE_SERVICE,
@@ -2754,7 +2762,7 @@
         app.mState.setSetProcState(PROCESS_STATE_NONEXISTENT);
         app.mState.setAdjType(null);
         app.mState.setSetAdj(UNKNOWN_ADJ);
-        app.mState.setHasShownUi(false);
+        mProcessStateController.setHasShownUi(app, false);
         updateOomAdj();
 
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND,
@@ -2763,27 +2771,28 @@
         app.mState.setSetProcState(PROCESS_STATE_NONEXISTENT);
         app.mState.setAdjType(null);
         app.mState.setSetAdj(UNKNOWN_ADJ);
-        s.lastActivity = now - mService.mConstants.MAX_SERVICE_INACTIVITY - 1;
+        mProcessStateController.setServiceLastActivityTime(s,
+                now - mService.mConstants.MAX_SERVICE_INACTIVITY - 1);
         updateOomAdj();
 
         assertProcStates(app, PROCESS_STATE_SERVICE,
                 mService.mConstants.USE_TIERED_CACHED_ADJ ? sFirstNonUiCachedAdj : cachedAdj1,
                 SCHED_GROUP_BACKGROUND, "cch-started-services", true);
 
-        app.mServices.stopService(s);
+        mProcessStateController.stopService(app.mServices, s);
         app.mState.setSetProcState(PROCESS_STATE_NONEXISTENT);
         app.mState.setAdjType(null);
         app.mState.setSetAdj(UNKNOWN_ADJ);
-        app.mState.setHasShownUi(true);
+        mProcessStateController.setHasShownUi(app, true);
         mService.mConstants.KEEP_WARMING_SERVICES.add(cn);
         mService.mConstants.KEEP_WARMING_SERVICES.add(cn2);
         s = spy(new ServiceRecord(mService, cn, cn, null, 0, null,
                 si, false, null));
         doReturn(new ArrayMap<IBinder, ArrayList<ConnectionRecord>>()).when(s).getConnections();
-        s.startRequested = true;
-        s.lastActivity = now;
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, now);
 
-        app.mServices.startService(s);
+        mProcessStateController.startService(app.mServices, s);
         updateOomAdj();
 
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND,
@@ -2795,8 +2804,9 @@
         app.mState.setSetProcState(PROCESS_STATE_NONEXISTENT);
         app.mState.setAdjType(null);
         app.mState.setSetAdj(UNKNOWN_ADJ);
-        app.mState.setHasShownUi(false);
-        s.lastActivity = now - mService.mConstants.MAX_SERVICE_INACTIVITY - 1;
+        mProcessStateController.setHasShownUi(app, false);
+        mProcessStateController.setServiceLastActivityTime(s,
+                now - mService.mConstants.MAX_SERVICE_INACTIVITY - 1);
         updateOomAdj();
 
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND,
@@ -2825,7 +2835,7 @@
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, true));
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(app).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertEquals(FOREGROUND_APP_ADJ, app.mState.getSetAdj());
@@ -2847,7 +2857,7 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(app).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertEquals(FOREGROUND_APP_ADJ, app.mState.getSetAdj());
@@ -2873,14 +2883,14 @@
                 MOCKAPP3_PROCESSNAME, MOCKAPP3_PACKAGENAME, false));
         long now = SystemClock.uptimeMillis();
         ServiceRecord s = bindService(app, app2, null, null, 0, mock(IBinder.class));
-        s.startRequested = true;
-        s.lastActivity = now;
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, now);
         s = bindService(app2, app3, null, null, 0, mock(IBinder.class));
-        s.lastActivity = now;
+        mProcessStateController.setServiceLastActivityTime(s, now);
         s = bindService(app3, app2, null, null, 0, mock(IBinder.class));
-        s.lastActivity = now;
+        mProcessStateController.setServiceLastActivityTime(s, now);
 
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         mService.mOomAdjuster.mNumServiceProcs = 3;
         updateOomAdj(app, app2, app3);
 
@@ -2898,14 +2908,14 @@
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
         doReturn(PROCESS_STATE_TOP).when(mService.mAtmInternal).getTopProcessState();
         doReturn(app).when(mService).getTopApp();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj(app);
 
         assertEquals(FOREGROUND_APP_ADJ, app.mState.getSetAdj());
 
         // Start binding to a service that isn't running yet.
         ServiceRecord sr = makeServiceRecord(app);
-        sr.app = null;
+        mProcessStateController.setHostProcess(sr, null);
         bindService(null, app, null, sr, Context.BIND_ABOVE_CLIENT, mock(IBinder.class));
 
         // Since sr.app is null, this service cannot be in the same process as the
@@ -2924,13 +2934,13 @@
 
         setProcessesToLru(app);
         ServiceRecord s = makeServiceRecord(app);
-        s.startRequested = true;
-        s.lastActivity = SystemClock.uptimeMillis();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, SystemClock.uptimeMillis());
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj();
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND);
 
-        app.mServices.stopService(s);
+        mProcessStateController.stopService(app.mServices, s);
         updateOomAdj();
         // isolated process should be killed immediately after service stop.
         verify(app).killLocked("isolated not needed", ApplicationExitInfo.REASON_OTHER,
@@ -2944,13 +2954,13 @@
                 MOCKAPP_ISOLATED_PROCESSNAME, MOCKAPP_PACKAGENAME, false));
 
         ServiceRecord s = makeServiceRecord(app);
-        s.startRequested = true;
-        s.lastActivity = SystemClock.uptimeMillis();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, SystemClock.uptimeMillis());
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdjPending(app);
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND);
 
-        app.mServices.stopService(s);
+        mProcessStateController.stopService(app.mServices, s);
         updateOomAdjPending(app);
         // isolated process should be killed immediately after service stop.
         verify(app).killLocked("isolated not needed", ApplicationExitInfo.REASON_OTHER,
@@ -2966,13 +2976,13 @@
 
         setProcessesToLru(app);
         ServiceRecord s = makeServiceRecord(app);
-        s.startRequested = true;
-        s.lastActivity = SystemClock.uptimeMillis();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, SystemClock.uptimeMillis());
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj();
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND);
 
-        app.mServices.stopService(s);
+        mProcessStateController.stopService(app.mServices, s);
         updateOomAdj();
         // isolated process with entry point should not be killed
         verify(app, never()).killLocked("isolated not needed", ApplicationExitInfo.REASON_OTHER,
@@ -2987,13 +2997,13 @@
         app.setIsolatedEntryPoint("test");
 
         ServiceRecord s = makeServiceRecord(app);
-        s.startRequested = true;
-        s.lastActivity = SystemClock.uptimeMillis();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, SystemClock.uptimeMillis());
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdjPending(app);
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND);
 
-        app.mServices.stopService(s);
+        mProcessStateController.stopService(app.mServices, s);
         updateOomAdjPending(app);
         // isolated process with entry point should not be killed
         verify(app, never()).killLocked("isolated not needed", ApplicationExitInfo.REASON_OTHER,
@@ -3014,10 +3024,10 @@
 
         setProcessesToLru(sandboxService, client, attributedClient);
 
-        client.mState.setMaxAdj(PERSISTENT_PROC_ADJ);
-        attributedClient.mServices.setHasForegroundServices(true, 0, true);
+        mProcessStateController.setMaxAdj(client, PERSISTENT_PROC_ADJ);
+        mProcessStateController.setHasForegroundServices(attributedClient.mServices, true, 0, true);
         bindService(sandboxService, client, attributedClient, null, 0, mock(IBinder.class));
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj();
         assertProcStates(client, PROCESS_STATE_PERSISTENT, PERSISTENT_PROC_ADJ,
                 SCHED_GROUP_DEFAULT);
@@ -3038,13 +3048,13 @@
 
         // App1 binds to app2 and gets temp allowlisted.
         bindService(app2, app, null, null, 0, mock(IBinder.class));
-        mService.mOomAdjuster.setUidTempAllowlistStateLSP(MOCKAPP_UID, true);
+        mProcessStateController.setUidTempAllowlistStateLSP(MOCKAPP_UID, true);
 
         assertEquals(true, app.getUidRecord().isSetAllowListed());
         assertEquals(true, app.mOptRecord.shouldNotFreeze());
         assertEquals(true, app2.mOptRecord.shouldNotFreeze());
 
-        mService.mOomAdjuster.setUidTempAllowlistStateLSP(MOCKAPP_UID, false);
+        mProcessStateController.setUidTempAllowlistStateLSP(MOCKAPP_UID, false);
         assertEquals(false, app.getUidRecord().isSetAllowListed());
         assertEquals(false, app.mOptRecord.shouldNotFreeze());
         assertEquals(false, app2.mOptRecord.shouldNotFreeze());
@@ -3064,8 +3074,8 @@
         // App1 and app2 both bind to app3 and get temp allowlisted.
         bindService(app3, app, null, null, 0, mock(IBinder.class));
         bindService(app3, app2, null, null, 0, mock(IBinder.class));
-        mService.mOomAdjuster.setUidTempAllowlistStateLSP(MOCKAPP_UID, true);
-        mService.mOomAdjuster.setUidTempAllowlistStateLSP(MOCKAPP2_UID, true);
+        mProcessStateController.setUidTempAllowlistStateLSP(MOCKAPP_UID, true);
+        mProcessStateController.setUidTempAllowlistStateLSP(MOCKAPP2_UID, true);
 
         assertEquals(true, app.getUidRecord().isSetAllowListed());
         assertEquals(true, app2.getUidRecord().isSetAllowListed());
@@ -3074,7 +3084,7 @@
         assertEquals(true, app3.mOptRecord.shouldNotFreeze());
 
         // Remove app1 from allowlist.
-        mService.mOomAdjuster.setUidTempAllowlistStateLSP(MOCKAPP_UID, false);
+        mProcessStateController.setUidTempAllowlistStateLSP(MOCKAPP_UID, false);
         assertEquals(false, app.getUidRecord().isSetAllowListed());
         assertEquals(true, app2.getUidRecord().isSetAllowListed());
         assertEquals(false, app.mOptRecord.shouldNotFreeze());
@@ -3082,7 +3092,7 @@
         assertEquals(true, app3.mOptRecord.shouldNotFreeze());
 
         // Now remove app2 from allowlist.
-        mService.mOomAdjuster.setUidTempAllowlistStateLSP(MOCKAPP2_UID, false);
+        mProcessStateController.setUidTempAllowlistStateLSP(MOCKAPP2_UID, false);
         assertEquals(false, app.getUidRecord().isSetAllowListed());
         assertEquals(false, app2.getUidRecord().isSetAllowListed());
         assertEquals(false, app.mOptRecord.shouldNotFreeze());
@@ -3098,9 +3108,9 @@
 
         setProcessesToLru(app);
         ServiceRecord s = makeServiceRecord(app);
-        s.startRequested = true;
-        s.lastActivity = SystemClock.uptimeMillis();
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        mProcessStateController.setStartRequested(s, true);
+        mProcessStateController.setServiceLastActivityTime(s, SystemClock.uptimeMillis());
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         updateOomAdj();
         assertProcStates(app, PROCESS_STATE_SERVICE, SERVICE_ADJ, SCHED_GROUP_BACKGROUND,
                 "started-services");
@@ -3111,7 +3121,7 @@
         verify(mService.mHandler).sendEmptyMessageAtTime(
                 eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG), followUpTimeCaptor.capture());
         mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
-        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+        mProcessStateController.runFollowUpUpdate();
 
         final int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
                 ? sFirstNonUiCachedAdj : sFirstCachedAdj;
@@ -3131,9 +3141,9 @@
                 MOCKAPP2_PROCESSNAME, MOCKAPP2_PACKAGENAME, false));
         app1.mProviders.setLastProviderTime(SystemClock.uptimeMillis());
         app2.mProviders.setLastProviderTime(SystemClock.uptimeMillis() + 2000);
-        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setWakefulness(PowerManagerInternal.WAKEFULNESS_AWAKE);
         setProcessesToLru(app1, app2);
-        mService.mOomAdjuster.updateOomAdjLocked(OOM_ADJ_REASON_NONE);
+        mProcessStateController.runFullUpdate(OOM_ADJ_REASON_NONE);
 
         assertProcStates(app1, PROCESS_STATE_LAST_ACTIVITY, PREVIOUS_APP_ADJ,
                 SCHED_GROUP_BACKGROUND, "recent-provider");
@@ -3146,7 +3156,7 @@
         verify(mService.mHandler, atLeastOnce()).sendEmptyMessageAtTime(
                 eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG), followUpTimeCaptor.capture());
         mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
-        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+        mProcessStateController.runFollowUpUpdate();
 
         final int expectedAdj = mService.mConstants.USE_TIERED_CACHED_ADJ
                 ? sFirstNonUiCachedAdj : sFirstCachedAdj;
@@ -3156,7 +3166,7 @@
         verify(mService.mHandler, atLeastOnce()).sendEmptyMessageAtTime(
                 eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG), followUpTimeCaptor.capture());
         mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
-        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+        mProcessStateController.runFollowUpUpdate();
         assertProcStates(app2, PROCESS_STATE_CACHED_EMPTY, expectedAdj, SCHED_GROUP_BACKGROUND,
                 "cch-empty");
     }
@@ -3169,12 +3179,12 @@
 
     private ServiceRecord makeServiceRecord(ProcessRecord app) {
         final ServiceRecord record = mock(ServiceRecord.class);
-        record.app = app;
+        mProcessStateController.setHostProcess(record, app);
         setFieldValue(ServiceRecord.class, record, "connections",
                 new ArrayMap<IBinder, ArrayList<ConnectionRecord>>());
         doCallRealMethod().when(record).getConnections();
         setFieldValue(ServiceRecord.class, record, "packageName", app.info.packageName);
-        app.mServices.startService(record);
+        mProcessStateController.startService(app.mServices, record);
         record.appInfo = app.info;
         setFieldValue(ServiceRecord.class, record, "bindings", new ArrayMap<>());
         setFieldValue(ServiceRecord.class, record, "pendingStarts", new ArrayList<>());
@@ -3213,17 +3223,36 @@
         doCallRealMethod().when(record).addConnection(any(IBinder.class),
                 any(ConnectionRecord.class));
         record.addConnection(binder, cr);
-        client.mServices.addConnection(cr);
+        mProcessStateController.addConnection(client.mServices, cr);
         binding.connections.add(cr);
         doNothing().when(cr).trackProcState(anyInt(), anyInt());
         return record;
     }
 
+    private void setWakefulness(int state) {
+        if (Flags.pushGlobalStateToOomadjuster()) {
+            mProcessStateController.setWakefulness(state);
+        } else {
+            mService.mWakefulness.set(state);
+        }
+    }
+
+    @SuppressWarnings("GuardedBy")
+    private void setBackupTarget(ProcessRecord app) {
+        if (Flags.pushGlobalStateToOomadjuster()) {
+            mProcessStateController.setBackupTarget(app, app.userId);
+        } else {
+            BackupRecord backupTarget = new BackupRecord(null, 0, 0, 0);
+            backupTarget.app = app;
+            doReturn(backupTarget).when(mService.mBackupTargets).get(anyInt());
+        }
+    }
+
     private ContentProviderRecord bindProvider(ProcessRecord publisher, ProcessRecord client,
             ContentProviderRecord record, String name, boolean hasExternalProviders) {
         if (record == null) {
             record = mock(ContentProviderRecord.class);
-            publisher.mProviders.installProvider(name, record);
+            mProcessStateController.addPublishedProvider(publisher, name, record);
             record.proc = publisher;
             setFieldValue(ContentProviderRecord.class, record, "connections",
                     new ArrayList<ContentProviderConnection>());
@@ -3232,7 +3261,7 @@
         ContentProviderConnection conn = spy(new ContentProviderConnection(record, client,
                 client.info.packageName, UserHandle.getUserId(client.uid)));
         record.connections.add(conn);
-        client.mProviders.addProviderConnection(conn);
+        mProcessStateController.addProviderConnection(client, conn);
         return record;
     }
 
@@ -3405,10 +3434,10 @@
             }
             providers.setLastProviderTime(mLastProviderTime);
 
-            UidRecord uidRec = mService.mOomAdjuster.mActiveUids.get(mUid);
+            UidRecord uidRec = mActiveUids.get(mUid);
             if (uidRec == null) {
                 uidRec = new UidRecord(mUid, mService);
-                mService.mOomAdjuster.mActiveUids.put(mUid, uidRec);
+                mActiveUids.put(mUid, uidRec);
             }
             uidRec.addProcess(app);
             app.setUidRecord(uidRec);
diff --git a/services/tests/mockingservicestests/src/com/android/server/am/ServiceBindingOomAdjPolicyTest.java b/services/tests/mockingservicestests/src/com/android/server/am/ServiceBindingOomAdjPolicyTest.java
index 1ff4a27..59302ee 100644
--- a/services/tests/mockingservicestests/src/com/android/server/am/ServiceBindingOomAdjPolicyTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/am/ServiceBindingOomAdjPolicyTest.java
@@ -50,7 +50,6 @@
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
-import android.app.IApplicationThread;
 import android.app.IServiceConnection;
 import android.app.usage.UsageStatsManagerInternal;
 import android.content.ComponentName;
@@ -169,6 +168,7 @@
         realAtm.initialize(null, null, mContext.getMainLooper());
         realAms.mActivityTaskManager = spy(realAtm);
         realAms.mAtmInternal = spy(realAms.mActivityTaskManager.getAtmInternal());
+        realAms.mProcessStateController = spy(realAms.mProcessStateController);
         realAms.mOomAdjuster = spy(realAms.mOomAdjuster);
         realAms.mOomAdjuster.mCachedAppOptimizer = spy(realAms.mOomAdjuster.mCachedAppOptimizer);
         realAms.mPackageManagerInt = mPackageManagerInt;
@@ -242,14 +242,14 @@
                 USER_SYSTEM              // userId
         ));
 
-        verify(mAms.mOomAdjuster, bindMode).updateOomAdjPendingTargetsLocked(anyInt());
-        clearInvocations(mAms.mOomAdjuster);
+        verify(mAms.mProcessStateController, bindMode).runPendingUpdate(anyInt());
+        clearInvocations(mAms.mProcessStateController);
 
         // Unbind the service.
         mAms.unbindService(serviceConnection);
 
-        verify(mAms.mOomAdjuster, unbindMode).updateOomAdjPendingTargetsLocked(anyInt());
-        clearInvocations(mAms.mOomAdjuster);
+        verify(mAms.mProcessStateController, unbindMode).runPendingUpdate(anyInt());
+        clearInvocations(mAms.mProcessStateController);
 
         removeProcessRecord(app);
     }
@@ -496,8 +496,8 @@
                 USER_SYSTEM            // userId
         ));
 
-        verify(mAms.mOomAdjuster, bindMode).updateOomAdjPendingTargetsLocked(anyInt());
-        clearInvocations(mAms.mOomAdjuster);
+        verify(mAms.mProcessStateController, bindMode).runPendingUpdate(anyInt());
+        clearInvocations(mAms.mProcessStateController);
 
         if (clientApp.isFreezable()) {
             verify(mAms.mOomAdjuster.mCachedAppOptimizer,
@@ -509,8 +509,8 @@
         // Unbind the service.
         mAms.unbindService(serviceConnection);
 
-        verify(mAms.mOomAdjuster, unbindMode).updateOomAdjPendingTargetsLocked(anyInt());
-        clearInvocations(mAms.mOomAdjuster);
+        verify(mAms.mProcessStateController, unbindMode).runPendingUpdate(anyInt());
+        clearInvocations(mAms.mProcessStateController);
 
         removeProcessRecord(clientApp);
         removeProcessRecord(serviceApp);
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index d6ca10a..d9e071f 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -27,6 +27,9 @@
         "servicestests-utils",
         "platform-test-annotations",
         "flag-junit",
+        "statsdprotolite",
+        "StatsdTestUtils",
+        "platformprotoslite",
     ],
 
     libs: [
@@ -74,6 +77,10 @@
         "src/com/android/server/power/stats/format/*.java",
         "src/com/android/server/power/stats/processor/*.java",
     ],
+    // TODO(b/372292543): Enable this test.
+    exclude_srcs: [
+        "src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java",
+    ],
     java_resources: [
         "res/xml/power_profile*.xml",
     ],
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java
new file mode 100644
index 0000000..cb644db
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WakelockStatsFrameworkEventsTest.java
@@ -0,0 +1,433 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.power.stats;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.os.WakeLockLevelEnum;
+import android.util.StatsEvent;
+import android.util.StatsEventTestUtils;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.os.AtomsProto;
+import com.android.os.framework.FrameworkExtensionAtoms;
+import com.android.os.framework.FrameworkExtensionAtoms.FrameworkWakelockInfo;
+
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.CodedOutputStream;
+import com.google.protobuf.ExtensionRegistryLite;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.util.ArrayList;
+import java.util.List;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class WakelockStatsFrameworkEventsTest {
+    private WakelockStatsFrameworkEvents mEvents;
+    private ExtensionRegistryLite mRegistry;
+
+    @Before
+    public void setup() {
+        mEvents = new WakelockStatsFrameworkEvents();
+        mRegistry = ExtensionRegistryLite.newInstance();
+        FrameworkExtensionAtoms.registerAllExtensions(mRegistry);
+    }
+
+    private static final int UID_1 = 1;
+    private static final int UID_2 = 2;
+
+    private static final String TAG_1 = "TAG1";
+    private static final String TAG_2 = "TAG2";
+
+    private static final WakeLockLevelEnum WAKELOCK_TYPE_1 = WakeLockLevelEnum.PARTIAL_WAKE_LOCK;
+    private static final WakeLockLevelEnum WAKELOCK_TYPE_2 = WakeLockLevelEnum.DOZE_WAKE_LOCK;
+
+    private static final long TS_1 = 1000;
+    private static final long TS_2 = 2000;
+    private static final long TS_3 = 3000;
+    private static final long TS_4 = 4000;
+    private static final long TS_5 = 5000;
+
+    // Assumes that mEvents is empty.
+    private void makeMetricsAlmostOverflow() throws Exception {
+        for (int i = 0; i < mEvents.SUMMARY_THRESHOLD - 1; i++) {
+            String tag = "forceOverflow" + i;
+            mEvents.noteStartWakeLock(UID_1, tag, WAKELOCK_TYPE_1.getNumber(), TS_1);
+            mEvents.noteStopWakeLock(UID_1, tag, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        }
+
+        assertFalse("not overflow", mEvents.inOverflow());
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo notOverflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("not overflow", notOverflowInfo, null);
+
+        // Add one more to hit an overflow state.
+        String lastTag = "forceOverflowLast";
+        mEvents.noteStartWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, lastTag, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        assertTrue("overflow", mEvents.inOverflow());
+        info = pullResults(TS_4);
+
+        FrameworkWakelockInfo tag1Info =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(lastTag))
+                        .findFirst()
+                        .orElse(null);
+
+        assertTrue("lastTag found", tag1Info != null);
+        assertEquals("uid", UID_1, tag1Info.getAttributionUid());
+        assertEquals("tag", lastTag, tag1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_2, tag1Info.getType());
+        assertEquals("duration", TS_2 - TS_1, tag1Info.getUptimeMillis());
+        assertEquals("count", 1, tag1Info.getCompletedCount());
+    }
+
+    // Assumes that mEvents is empty.
+    private void makeMetricsAlmostHardCap() throws Exception {
+        for (int i = 0; i < mEvents.MAX_WAKELOCK_DIMENSIONS - 1; i++) {
+            mEvents.noteStartWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+            mEvents.noteStopWakeLock(i /* uid */, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        }
+
+        assertFalse("not hard capped", mEvents.inHardCap());
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo notOverflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("not overflow", notOverflowInfo, null);
+
+        // Add one more to hit an hardcap state.
+        int hardCapUid = mEvents.MAX_WAKELOCK_DIMENSIONS;
+        mEvents.noteStartWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(hardCapUid, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        assertTrue("hard capped", mEvents.inHardCap());
+        info = pullResults(TS_4);
+
+        FrameworkWakelockInfo tag2Info =
+                info.stream()
+                        .filter(i -> i.getAttributionUid() == hardCapUid)
+                        .findFirst()
+                        .orElse(null);
+
+        assertTrue("hardCapUid found", tag2Info != null);
+        assertEquals("uid", hardCapUid, tag2Info.getAttributionUid());
+        assertEquals("tag", mEvents.OVERFLOW_TAG, tag2Info.getAttributionTag());
+        assertEquals(
+                "type", WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL), tag2Info.getType());
+        assertEquals("duration", TS_2 - TS_1, tag2Info.getUptimeMillis());
+        assertEquals("count", 1, tag2Info.getCompletedCount());
+    }
+
+    private ArrayList<FrameworkWakelockInfo> pullResults(long timestamp) throws Exception {
+        ArrayList<FrameworkWakelockInfo> result = new ArrayList<>();
+        List<StatsEvent> events = mEvents.pullFrameworkWakelockInfoAtoms(timestamp);
+
+        for (StatsEvent e : events) {
+            // The returned atom does not have external extensions registered.
+            // So we serialize and then deserialize with extensions registered.
+            AtomsProto.Atom atom = StatsEventTestUtils.convertToAtom(e);
+
+            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+            CodedOutputStream codedos = CodedOutputStream.newInstance(outputStream);
+            atom.writeTo(codedos);
+            codedos.flush();
+
+            ByteArrayInputStream inputStream = new ByteArrayInputStream(outputStream.toByteArray());
+            CodedInputStream codedis = CodedInputStream.newInstance(inputStream);
+            AtomsProto.Atom atomWithExtensions = AtomsProto.Atom.parseFrom(codedis, mRegistry);
+
+            assertTrue(
+                    atomWithExtensions.hasExtension(FrameworkExtensionAtoms.frameworkWakelockInfo));
+            FrameworkWakelockInfo info =
+                    atomWithExtensions.getExtension(FrameworkExtensionAtoms.frameworkWakelockInfo);
+            result.add(info);
+        }
+
+        return result;
+    }
+
+    @Test
+    public void singleWakelock() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_3);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_2 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 1, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void wakelockOpen() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_3);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_3 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 0, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void wakelockOpenOverlap() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_4 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 0, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void testOverflow() throws Exception {
+        makeMetricsAlmostOverflow();
+
+        // This one gets tagged as an overflow.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo overflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", UID_1, overflowInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                overflowInfo.getType());
+        assertEquals("duration", TS_2 - TS_1, overflowInfo.getUptimeMillis());
+        assertEquals("count", 1, overflowInfo.getCompletedCount());
+    }
+
+    @Test
+    public void testOverflowOpen() throws Exception {
+        makeMetricsAlmostOverflow();
+
+        // This is the open wakelock that overflows.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo overflowInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.OVERFLOW_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", UID_1, overflowInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                overflowInfo.getType());
+        assertEquals("duration", (TS_4 - TS_1), overflowInfo.getUptimeMillis());
+        assertEquals("count", 0, overflowInfo.getCompletedCount());
+    }
+
+    @Test
+    public void testHardCap() throws Exception {
+        makeMetricsAlmostHardCap();
+
+        // This one gets tagged as a hard cap.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_2);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo hardCapInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                hardCapInfo.getType());
+        assertEquals("duration", TS_2 - TS_1, hardCapInfo.getUptimeMillis());
+        assertEquals("count", 1, hardCapInfo.getCompletedCount());
+    }
+
+    @Test
+    public void testHardCapOpen() throws Exception {
+        makeMetricsAlmostHardCap();
+
+        // This is the open wakelock that overflows.
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_2.getNumber(), TS_1);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_4);
+        FrameworkWakelockInfo hardCapInfo =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(mEvents.HARD_CAP_TAG))
+                        .findFirst()
+                        .orElse(null);
+
+        assertEquals("uid", mEvents.HARD_CAP_UID, hardCapInfo.getAttributionUid());
+        assertEquals(
+                "type",
+                WakeLockLevelEnum.forNumber(mEvents.OVERFLOW_LEVEL),
+                hardCapInfo.getType());
+        assertEquals("duration", (TS_4 - TS_1), hardCapInfo.getUptimeMillis());
+        assertEquals("count", 0, hardCapInfo.getCompletedCount());
+    }
+
+    @Test
+    public void overlappingWakelocks() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+
+        assertEquals("size", 1, info.size());
+        assertEquals("uid", UID_1, info.get(0).getAttributionUid());
+        assertEquals("tag", TAG_1, info.get(0).getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, info.get(0).getType());
+        assertEquals("duration", TS_4 - TS_1, info.get(0).getUptimeMillis());
+        assertEquals("count", 1, info.get(0).getCompletedCount());
+    }
+
+    @Test
+    public void diffUid() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_2, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        assertEquals("size", 2, info.size());
+
+        FrameworkWakelockInfo uid1Info =
+                info.stream().filter(i -> i.getAttributionUid() == UID_1).findFirst().orElse(null);
+
+        assertTrue("UID_1 found", uid1Info != null);
+        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
+        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
+        assertEquals("count", 1, uid1Info.getCompletedCount());
+
+        FrameworkWakelockInfo uid2Info =
+                info.stream().filter(i -> i.getAttributionUid() == UID_2).findFirst().orElse(null);
+        assertTrue("UID_2 found", uid2Info != null);
+        assertEquals("uid", UID_2, uid2Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid2Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.getType());
+        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
+        assertEquals("count", 1, uid2Info.getCompletedCount());
+    }
+
+    @Test
+    public void diffTag() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_2, WAKELOCK_TYPE_1.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        assertEquals("size", 2, info.size());
+
+        FrameworkWakelockInfo uid1Info =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(TAG_1))
+                        .findFirst()
+                        .orElse(null);
+
+        assertTrue("TAG_1 found", uid1Info != null);
+        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
+        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
+        assertEquals("count", 1, uid1Info.getCompletedCount());
+
+        FrameworkWakelockInfo uid2Info =
+                info.stream()
+                        .filter(i -> i.getAttributionTag().equals(TAG_2))
+                        .findFirst()
+                        .orElse(null);
+        assertTrue("TAG_2 found", uid2Info != null);
+        assertEquals("uid", UID_1, uid2Info.getAttributionUid());
+        assertEquals("tag", TAG_2, uid2Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid2Info.getType());
+        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
+        assertEquals("count", 1, uid2Info.getCompletedCount());
+    }
+
+    @Test
+    public void diffType() throws Exception {
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_1);
+        mEvents.noteStartWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2.getNumber(), TS_2);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_2.getNumber(), TS_3);
+        mEvents.noteStopWakeLock(UID_1, TAG_1, WAKELOCK_TYPE_1.getNumber(), TS_4);
+
+        ArrayList<FrameworkWakelockInfo> info = pullResults(TS_5);
+        assertEquals("size", 2, info.size());
+
+        FrameworkWakelockInfo uid1Info =
+                info.stream().filter(i -> i.getType() == WAKELOCK_TYPE_1).findFirst().orElse(null);
+
+        assertTrue("WAKELOCK_TYPE_1 found", uid1Info != null);
+        assertEquals("uid", UID_1, uid1Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid1Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_1, uid1Info.getType());
+        assertEquals("duration", TS_4 - TS_1, uid1Info.getUptimeMillis());
+        assertEquals("count", 1, uid1Info.getCompletedCount());
+
+        FrameworkWakelockInfo uid2Info =
+                info.stream().filter(i -> i.getType() == WAKELOCK_TYPE_2).findFirst().orElse(null);
+        assertTrue("WAKELOCK_TYPE_2 found", uid2Info != null);
+        assertEquals("uid", UID_1, uid2Info.getAttributionUid());
+        assertEquals("tag", TAG_1, uid2Info.getAttributionTag());
+        assertEquals("type", WAKELOCK_TYPE_2, uid2Info.getType());
+        assertEquals("duration", TS_3 - TS_2, uid2Info.getUptimeMillis());
+        assertEquals("count", 1, uid2Info.getCompletedCount());
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/am/OomAdjusterTests.java b/services/tests/servicestests/src/com/android/server/am/OomAdjusterTests.java
index 31bf5f0..4981ceb 100644
--- a/services/tests/servicestests/src/com/android/server/am/OomAdjusterTests.java
+++ b/services/tests/servicestests/src/com/android/server/am/OomAdjusterTests.java
@@ -61,6 +61,7 @@
     private static final long USAGE_STATS_INTERACTION = 10 * 60 * 1000L;
     private static final long SERVICE_USAGE_INTERACTION = 60 * 1000;
 
+    @SuppressWarnings("GuardedBy")
     @BeforeClass
     public static void setUpOnce() {
         sContext = getInstrumentation().getTargetContext();
@@ -92,8 +93,11 @@
                     return true;
                 }
             };
-            sService.mOomAdjuster = new OomAdjuster(sService, sService.mProcessList, null,
-                    injector);
+            sService.mProcessStateController = new ProcessStateController.Builder(sService,
+                    sService.mProcessList, null)
+                    .setOomAdjusterInjector(injector)
+                    .build();
+            sService.mOomAdjuster = sService.mProcessStateController.getOomAdjuster();
             LocalServices.addService(UsageStatsManagerInternal.class,
                     mock(UsageStatsManagerInternal.class));
             sService.mUsageStatsService = LocalServices.getService(UsageStatsManagerInternal.class);
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
index e6c34ca..62f5edc 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
@@ -1999,7 +1999,8 @@
                 /* tag= */ null, MacAddress.BROADCAST_ADDRESS, displayName, deviceProfile,
                 /* associatedDevice= */ null, /* selfManaged= */ true,
                 /* notifyOnDeviceNearby= */ false, /* revoked= */ false, /* pending= */ false,
-                /* timeApprovedMs= */0, /* lastTimeConnectedMs= */0, /* systemDataSyncFlags= */ -1);
+                /* timeApprovedMs= */0, /* lastTimeConnectedMs= */0,
+                /* systemDataSyncFlags= */ -1, /* deviceIcon= */ null);
     }
 
     /** Helper class to drop permissions temporarily and restore them at the end of a test. */
diff --git a/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java
index 733f056..b565f4b 100644
--- a/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicestate/DeviceStateManagerServiceTest.java
@@ -24,13 +24,9 @@
 
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
-import static org.testng.Assert.assertEquals;
-import static org.testng.Assert.assertFalse;
-import static org.testng.Assert.assertNotNull;
-import static org.testng.Assert.assertNull;
 import static org.testng.Assert.assertThrows;
 
-import android.annotation.NonNull;
+import android.content.Context;
 import android.hardware.devicestate.DeviceState;
 import android.hardware.devicestate.DeviceStateInfo;
 import android.hardware.devicestate.DeviceStateRequest;
@@ -40,9 +36,10 @@
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
 
-import androidx.test.InstrumentationRegistry;
+import androidx.annotation.NonNull;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.FlakyTest;
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
 
 import com.android.compatibility.common.util.PollingCheck;
 import com.android.server.wm.ActivityTaskManagerInternal;
@@ -63,16 +60,25 @@
 
 /**
  * Unit tests for {@link DeviceStateManagerService}.
- * <p/>
- * Run with <code>atest DeviceStateManagerServiceTest</code>.
+ *
+ * <p> Build/Install/Run:
+ * atest FrameworksServicesTests:DeviceStateManagerServiceTest
  */
 @Presubmit
 @RunWith(AndroidJUnit4.class)
 public final class DeviceStateManagerServiceTest {
     private static final DeviceState DEFAULT_DEVICE_STATE = new DeviceState(
             new DeviceState.Configuration.Builder(0, "DEFAULT").build());
+    private static final int DEFAULT_DEVICE_STATE_IDENTIFIER = DEFAULT_DEVICE_STATE.getIdentifier();
+    private static final String DEFAULT_DEVICE_STATE_TRACE_STRING =
+            DEFAULT_DEVICE_STATE_IDENTIFIER + ":" + DEFAULT_DEVICE_STATE.getName();
+
     private static final DeviceState OTHER_DEVICE_STATE = new DeviceState(
             new DeviceState.Configuration.Builder(1, "DEFAULT").build());
+    private static final int OTHER_DEVICE_STATE_IDENTIFIER = OTHER_DEVICE_STATE.getIdentifier();
+    private static final String OTHER_DEVICE_STATE_TRACE_STRING =
+            OTHER_DEVICE_STATE_IDENTIFIER + ":" + OTHER_DEVICE_STATE.getName();
+
     private static final DeviceState DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP =
             new DeviceState(new DeviceState.Configuration.Builder(2,
                     "DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP")
@@ -93,24 +99,29 @@
 
     private static final int TIMEOUT = 2000;
 
+    private final Context mContext = InstrumentationRegistry.getInstrumentation().getContext();
+    @NonNull
     private TestDeviceStatePolicy mPolicy;
+    @NonNull
     private TestDeviceStateProvider mProvider;
+    @NonNull
     private DeviceStateManagerService mService;
+    @NonNull
     private TestSystemPropertySetter mSysPropSetter;
+    @NonNull
     private WindowProcessController mWindowProcessController;
 
     @Before
     public void setup() {
         mProvider = new TestDeviceStateProvider();
-        mPolicy = new TestDeviceStatePolicy(mProvider);
+        mPolicy = new TestDeviceStatePolicy(mContext, mProvider);
         mSysPropSetter = new TestSystemPropertySetter();
         setupDeviceStateManagerService();
         flushHandler(); // Flush the handler to ensure the initial values are committed.
     }
 
     private void setupDeviceStateManagerService() {
-        mService = new DeviceStateManagerService(InstrumentationRegistry.getContext(), mPolicy,
-                mSysPropSetter);
+        mService = new DeviceStateManagerService(mContext, mPolicy, mSysPropSetter);
 
         // Necessary to allow us to check for top app process id in tests
         mService.mActivityTaskManagerInternal = mock(ActivityTaskManagerInternal.class);
@@ -136,56 +147,53 @@
 
     @Test
     public void baseStateChanged() {
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
 
-        mProvider.setState(OTHER_DEVICE_STATE.getIdentifier());
+        mProvider.setState(OTHER_DEVICE_STATE_IDENTIFIER);
         flushHandler();
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
     public void baseStateChanged_withStatePendingPolicyCallback() {
         mPolicy.blockConfigure();
 
-        mProvider.setState(OTHER_DEVICE_STATE.getIdentifier());
+        mProvider.setState(OTHER_DEVICE_STATE_IDENTIFIER);
         flushHandler();
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
-        mProvider.setState(DEFAULT_DEVICE_STATE.getIdentifier());
+        mProvider.setState(DEFAULT_DEVICE_STATE_IDENTIFIER);
         flushHandler();
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
         mPolicy.resumeConfigure();
         flushHandler();
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
@@ -194,13 +202,12 @@
             mProvider.setState(UNSUPPORTED_DEVICE_STATE.getIdentifier());
         });
 
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
@@ -208,25 +215,23 @@
         assertThrows(IllegalArgumentException.class,
                 () -> mProvider.setState(INVALID_DEVICE_STATE_IDENTIFIER));
 
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
     public void supportedStatesChanged() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
 
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
         assertThat(mService.getSupportedStates()).containsExactly(DEFAULT_DEVICE_STATE,
                 OTHER_DEVICE_STATE, DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP);
 
@@ -235,30 +240,29 @@
 
         // The current committed and requests states do not change because the current state remains
         // supported.
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
         assertThat(mService.getSupportedStates()).containsExactly(DEFAULT_DEVICE_STATE);
 
-        assertEquals(callback.getLastNotifiedInfo().supportedStates, List.of(DEFAULT_DEVICE_STATE));
+        assertThat(callback.getLastNotifiedInfo().supportedStates)
+                .containsExactly(DEFAULT_DEVICE_STATE);
     }
 
     @Test
     public void supportedStatesChanged_statesRemainSame() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
 
         // An initial callback will be triggered on registration, so we clear it here.
         flushHandler();
         callback.clearLastNotifiedInfo();
 
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
         assertThat(mService.getSupportedStates()).containsExactly(DEFAULT_DEVICE_STATE,
                 OTHER_DEVICE_STATE, DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP);
 
@@ -268,26 +272,26 @@
 
         // The current committed and requests states do not change because the current state remains
         // supported.
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
         assertThat(mService.getSupportedStates()).containsExactly(DEFAULT_DEVICE_STATE,
                 OTHER_DEVICE_STATE, DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP);
 
         // The callback wasn't notified about a change in supported states as the states have not
         // changed.
-        assertNull(callback.getLastNotifiedInfo());
+        assertThat(callback.getLastNotifiedInfo()).isNull();
     }
 
     @Test
     public void getDeviceStateInfo() throws RemoteException {
-        DeviceStateInfo info = mService.getBinderService().getDeviceStateInfo();
-        assertNotNull(info);
-        assertEquals(info.supportedStates, SUPPORTED_DEVICE_STATES);
-        assertEquals(info.baseState, DEFAULT_DEVICE_STATE);
-        assertEquals(info.currentState, DEFAULT_DEVICE_STATE);
+        final DeviceStateInfo info = mService.getBinderService().getDeviceStateInfo();
+        assertThat(info).isNotNull();
+        assertThat(info.supportedStates)
+                .containsExactlyElementsIn(SUPPORTED_DEVICE_STATES).inOrder();
+        assertThat(info.baseState).isEqualTo(DEFAULT_DEVICE_STATE);
+        assertThat(info.currentState).isEqualTo(DEFAULT_DEVICE_STATE);
     }
 
     @FlakyTest(bugId = 297949293)
@@ -295,105 +299,109 @@
     public void getDeviceStateInfo_baseStateAndCommittedStateNotSet() throws RemoteException {
         // Create a provider and a service without an initial base state.
         mProvider = new TestDeviceStateProvider(null /* initialState */);
-        mPolicy = new TestDeviceStatePolicy(mProvider);
+        mPolicy = new TestDeviceStatePolicy(mContext, mProvider);
         setupDeviceStateManagerService();
         flushHandler(); // Flush the handler to ensure the initial values are committed.
 
-        DeviceStateInfo info = mService.getBinderService().getDeviceStateInfo();
+        final DeviceStateInfo info = mService.getBinderService().getDeviceStateInfo();
 
-        assertEquals(info.supportedStates, SUPPORTED_DEVICE_STATES);
-        assertEquals(info.baseState.getIdentifier(), INVALID_DEVICE_STATE_IDENTIFIER);
-        assertEquals(info.currentState.getIdentifier(), INVALID_DEVICE_STATE_IDENTIFIER);
+        assertThat(info.supportedStates)
+                .containsExactlyElementsIn(SUPPORTED_DEVICE_STATES).inOrder();
+        assertThat(info.baseState.getIdentifier()).isEqualTo(INVALID_DEVICE_STATE_IDENTIFIER);
+        assertThat(info.currentState.getIdentifier()).isEqualTo(INVALID_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
     public void registerCallback() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
 
-        mProvider.setState(OTHER_DEVICE_STATE.getIdentifier());
+        mProvider.setState(OTHER_DEVICE_STATE_IDENTIFIER);
         waitAndAssert(() -> callback.getLastNotifiedInfo().baseState.getIdentifier()
-                == OTHER_DEVICE_STATE.getIdentifier());
+                == OTHER_DEVICE_STATE_IDENTIFIER);
         waitAndAssert(() -> callback.getLastNotifiedInfo().currentState.getIdentifier()
-                == OTHER_DEVICE_STATE.getIdentifier());
+                == OTHER_DEVICE_STATE_IDENTIFIER);
 
-        mProvider.setState(DEFAULT_DEVICE_STATE.getIdentifier());
+        mProvider.setState(DEFAULT_DEVICE_STATE_IDENTIFIER);
         waitAndAssert(() -> callback.getLastNotifiedInfo().baseState.getIdentifier()
-                == DEFAULT_DEVICE_STATE.getIdentifier());
-
+                == DEFAULT_DEVICE_STATE_IDENTIFIER);
         waitAndAssert(() -> callback.getLastNotifiedInfo().currentState.getIdentifier()
-                == DEFAULT_DEVICE_STATE.getIdentifier());
+                == DEFAULT_DEVICE_STATE_IDENTIFIER);
 
         mPolicy.blockConfigure();
-        mProvider.setState(OTHER_DEVICE_STATE.getIdentifier());
+        mProvider.setState(OTHER_DEVICE_STATE_IDENTIFIER);
         // The callback should not have been notified of the state change as the policy is still
         // pending callback.
         waitAndAssert(() -> callback.getLastNotifiedInfo().baseState.getIdentifier()
-                == DEFAULT_DEVICE_STATE.getIdentifier());
+                == DEFAULT_DEVICE_STATE_IDENTIFIER);
         waitAndAssert(() -> callback.getLastNotifiedInfo().currentState.getIdentifier()
-                == DEFAULT_DEVICE_STATE.getIdentifier());
+                == DEFAULT_DEVICE_STATE_IDENTIFIER);
 
         mPolicy.resumeConfigure();
         // Now that the policy is finished processing the callback should be notified of the state
         // change.
         waitAndAssert(() -> callback.getLastNotifiedInfo().baseState.getIdentifier()
-                == OTHER_DEVICE_STATE.getIdentifier());
+                == OTHER_DEVICE_STATE_IDENTIFIER);
         waitAndAssert(() -> callback.getLastNotifiedInfo().currentState.getIdentifier()
-                == OTHER_DEVICE_STATE.getIdentifier());
+                == OTHER_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
-    public void registerCallback_emitsInitialValue() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+    public void registerCallback_initialValueAvailable_emitsDeviceState() throws RemoteException {
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+
         mService.getBinderService().registerCallback(callback);
         flushHandler();
-        assertNotNull(callback.getLastNotifiedInfo());
-        assertEquals(callback.getLastNotifiedInfo().baseState, DEFAULT_DEVICE_STATE);
-        assertEquals(callback.getLastNotifiedInfo().currentState, DEFAULT_DEVICE_STATE);
+        final DeviceStateInfo stateInfo = callback.getLastNotifiedInfo();
+
+        assertThat(stateInfo).isNotNull();
+        assertThat(stateInfo.baseState).isEqualTo(DEFAULT_DEVICE_STATE);
+        assertThat(stateInfo.currentState).isEqualTo(DEFAULT_DEVICE_STATE);
     }
 
     @Test
-    public void registerCallback_initialValueUnavailable() throws RemoteException {
+    public void registerCallback_initialValueUnavailable_nullDeviceState() throws RemoteException {
         // Create a provider and a service without an initial base state.
         mProvider = new TestDeviceStateProvider(null /* initialState */);
-        mPolicy = new TestDeviceStatePolicy(mProvider);
+        mPolicy = new TestDeviceStatePolicy(mContext, mProvider);
         setupDeviceStateManagerService();
         flushHandler(); // Flush the handler to ensure the initial values are committed.
 
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
+        final DeviceStateInfo stateInfo = callback.getLastNotifiedInfo();
+
         // The callback should never be called when the base state is not set yet.
-        assertNull(callback.getLastNotifiedInfo());
+        assertThat(stateInfo).isNull();
     }
 
     @Test
     public void requestState() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
 
         final IBinder token = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
-        mService.getBinderService().requestState(token, OTHER_DEVICE_STATE.getIdentifier(),
-                0 /* flags */);
+        mService.getBinderService()
+                .requestState(token, OTHER_DEVICE_STATE_IDENTIFIER, 0 /* flags */);
 
         waitAndAssert(() -> callback.getLastNotifiedStatus(token)
                 == TestDeviceStateManagerCallback.STATUS_ACTIVE);
         // Committed state changes as there is a requested override.
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getOverrideState().get(), OTHER_DEVICE_STATE);
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getOverrideState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
-        assertNotNull(callback.getLastNotifiedInfo());
-        assertEquals(callback.getLastNotifiedInfo().baseState, DEFAULT_DEVICE_STATE);
-        assertEquals(callback.getLastNotifiedInfo().currentState, OTHER_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo()).isNotNull();
+        assertThat(callback.getLastNotifiedInfo().baseState).isEqualTo(DEFAULT_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo().currentState).isEqualTo(OTHER_DEVICE_STATE);
 
         mService.getBinderService().cancelStateRequest();
 
@@ -401,21 +409,20 @@
                 == TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set back to the requested state once the override is cleared.
         waitAndAssert(() -> mService.getCommittedState().equals(Optional.of(DEFAULT_DEVICE_STATE)));
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertFalse(mService.getOverrideState().isPresent());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getOverrideState()).isEmpty();
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
 
-        assertEquals(callback.getLastNotifiedInfo().baseState, DEFAULT_DEVICE_STATE);
-        assertEquals(callback.getLastNotifiedInfo().currentState, DEFAULT_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo().baseState).isEqualTo(DEFAULT_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo().currentState).isEqualTo(DEFAULT_DEVICE_STATE);
     }
 
     @FlakyTest(bugId = 200332057)
     @Test
     public void requestState_pendingStateAtRequest() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
 
@@ -423,134 +430,129 @@
 
         final IBinder firstRequestToken = new Binder();
         final IBinder secondRequestToken = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(firstRequestToken),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
-        assertEquals(callback.getLastNotifiedStatus(secondRequestToken),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(firstRequestToken))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(secondRequestToken))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
-        mService.getBinderService().requestState(firstRequestToken,
-                OTHER_DEVICE_STATE.getIdentifier(), 0 /* flags */);
+        mService.getBinderService()
+                .requestState(firstRequestToken, OTHER_DEVICE_STATE_IDENTIFIER, 0 /* flags */);
         // Flush the handler twice. The first flush ensures the request is added and the policy is
         // notified, while the second flush ensures the callback is notified once the change is
         // committed.
         flushHandler(2 /* count */);
 
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
-        mService.getBinderService().requestState(secondRequestToken,
-                DEFAULT_DEVICE_STATE.getIdentifier(), 0 /* flags */);
+        mService.getBinderService()
+                .requestState(secondRequestToken, DEFAULT_DEVICE_STATE_IDENTIFIER, 0 /* flags */);
         mPolicy.resumeConfigureOnce();
         flushHandler();
 
         // First request status is now canceled as there is another pending request.
-        assertEquals(callback.getLastNotifiedStatus(firstRequestToken),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        assertThat(callback.getLastNotifiedStatus(firstRequestToken))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Second request status still unknown because the service is still awaiting policy
         // callback.
-        assertEquals(callback.getLastNotifiedStatus(secondRequestToken),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(callback.getLastNotifiedStatus(secondRequestToken))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getPendingState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
 
         mPolicy.resumeConfigure();
         flushHandler();
 
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
 
         // Now cancel the second request to make the first request active.
         mService.getBinderService().cancelStateRequest();
         flushHandler();
 
-        assertEquals(callback.getLastNotifiedStatus(firstRequestToken),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
-        assertEquals(callback.getLastNotifiedStatus(secondRequestToken),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        assertThat(callback.getLastNotifiedStatus(firstRequestToken))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_CANCELED);
+        assertThat(callback.getLastNotifiedStatus(secondRequestToken))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_CANCELED);
 
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getPendingState(), Optional.empty());
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getPendingState()).isEmpty();
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
     public void requestState_sameAsBaseState() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
 
         final IBinder token = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
-        mService.getBinderService().requestState(token, DEFAULT_DEVICE_STATE.getIdentifier(),
-                0 /* flags */);
+        mService.getBinderService()
+                .requestState(token, DEFAULT_DEVICE_STATE_IDENTIFIER, 0 /* flags */);
         flushHandler();
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_ACTIVE);
     }
 
     @Test
     public void requestState_flagCancelWhenBaseChanges() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
 
         final IBinder token = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
-        mService.getBinderService().requestState(token, OTHER_DEVICE_STATE.getIdentifier(),
+        mService.getBinderService().requestState(token, OTHER_DEVICE_STATE_IDENTIFIER,
                 DeviceStateRequest.FLAG_CANCEL_WHEN_BASE_CHANGES);
         // Flush the handler twice. The first flush ensures the request is added and the policy is
         // notified, while the second flush ensures the callback is notified once the change is
         // committed.
         flushHandler(2 /* count */);
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_ACTIVE);
 
         // Committed state changes as there is a requested override.
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertEquals(mService.getOverrideState().get(), OTHER_DEVICE_STATE);
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getOverrideState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
-        mProvider.setState(OTHER_DEVICE_STATE.getIdentifier());
+        mProvider.setState(OTHER_DEVICE_STATE_IDENTIFIER);
         flushHandler();
 
         // Request is canceled because the base state changed.
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set back to the requested state once the override is cleared.
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mService.getBaseState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertFalse(mService.getOverrideState().isPresent());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getBaseState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getOverrideState()).isEmpty();
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
@@ -581,8 +583,8 @@
         requestState_flagCancelWhenRequesterNotOnTop_common(
                 // When the app is foreground, the state should not change
                 () -> {
-                    int pid = Binder.getCallingPid();
-                    int uid = Binder.getCallingUid();
+                    final int pid = Binder.getCallingPid();
+                    final int uid = Binder.getCallingUid();
                     try {
                         mService.mProcessObserver.onForegroundActivitiesChanged(pid, uid,
                                 true /* foregroundActivities */);
@@ -594,8 +596,8 @@
                 () -> {
                     when(mWindowProcessController.getPid()).thenReturn(FAKE_PROCESS_ID);
                     try {
-                        int pid = Binder.getCallingPid();
-                        int uid = Binder.getCallingUid();
+                        final int pid = Binder.getCallingPid();
+                        final int uid = Binder.getCallingUid();
                         mService.mProcessObserver.onForegroundActivitiesChanged(pid, uid,
                                 false /* foregroundActivities */);
 
@@ -609,68 +611,66 @@
     @FlakyTest(bugId = 200332057)
     @Test
     public void requestState_becomesUnsupported() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
 
         final IBinder token = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
-        mService.getBinderService().requestState(token, OTHER_DEVICE_STATE.getIdentifier(),
-                0 /* flags */);
+        mService.getBinderService()
+                .requestState(token, OTHER_DEVICE_STATE_IDENTIFIER, 0 /* flags */);
         flushHandler();
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_ACTIVE);
         // Committed state changes as there is a requested override.
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mService.getOverrideState().get(), OTHER_DEVICE_STATE);
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getOverrideState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
         mProvider.notifySupportedDeviceStates(
                 new DeviceState[]{DEFAULT_DEVICE_STATE});
         flushHandler();
 
         // Request is canceled because the state is no longer supported.
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set back to the requested state as the override state is no longer
         // supported.
-        assertEquals(mService.getCommittedState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertFalse(mService.getOverrideState().isPresent());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getOverrideState()).isEmpty();
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
     }
 
     @Test
     public void requestState_unsupportedState() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
 
         assertThrows(IllegalArgumentException.class, () -> {
             final IBinder token = new Binder();
-            mService.getBinderService().requestState(token,
-                    UNSUPPORTED_DEVICE_STATE.getIdentifier(), 0 /* flags */);
+            mService.getBinderService()
+                    .requestState(token, UNSUPPORTED_DEVICE_STATE.getIdentifier(), 0 /* flags */);
         });
     }
 
     @Test
     public void requestState_invalidState() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
 
         assertThrows(IllegalArgumentException.class, () -> {
             final IBinder token = new Binder();
-            mService.getBinderService().requestState(token, INVALID_DEVICE_STATE_IDENTIFIER,
-                    0 /* flags */);
+            mService.getBinderService()
+                    .requestState(token, INVALID_DEVICE_STATE_IDENTIFIER, 0 /* flags */);
         });
     }
 
@@ -678,40 +678,39 @@
     public void requestState_beforeRegisteringCallback() {
         assertThrows(IllegalStateException.class, () -> {
             final IBinder token = new Binder();
-            mService.getBinderService().requestState(token, DEFAULT_DEVICE_STATE.getIdentifier(),
-                    0 /* flags */);
+            mService.getBinderService()
+                    .requestState(token, DEFAULT_DEVICE_STATE_IDENTIFIER, 0 /* flags */);
         });
     }
 
     @Test
     public void requestBaseStateOverride() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
 
         final IBinder token = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
         mService.getBinderService().requestBaseStateOverride(token,
-                OTHER_DEVICE_STATE.getIdentifier(),
+                OTHER_DEVICE_STATE_IDENTIFIER,
                 0 /* flags */);
 
         waitAndAssert(() -> callback.getLastNotifiedStatus(token)
                 == TestDeviceStateManagerCallback.STATUS_ACTIVE);
         // Committed state changes as there is a requested override.
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mService.getOverrideBaseState().get(), OTHER_DEVICE_STATE);
-        assertFalse(mService.getOverrideState().isPresent());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getOverrideBaseState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getOverrideState()).isEmpty();
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
-        assertNotNull(callback.getLastNotifiedInfo());
-        assertEquals(callback.getLastNotifiedInfo().baseState, OTHER_DEVICE_STATE);
-        assertEquals(callback.getLastNotifiedInfo().currentState, OTHER_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo()).isNotNull();
+        assertThat(callback.getLastNotifiedInfo().baseState).isEqualTo(OTHER_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo().currentState).isEqualTo(OTHER_DEVICE_STATE);
 
         mService.getBinderService().cancelBaseStateOverride();
 
@@ -719,52 +718,50 @@
                 == TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set back to the requested state once the override is cleared.
         waitAndAssert(() -> mService.getCommittedState().equals(Optional.of(DEFAULT_DEVICE_STATE)));
-        assertEquals(mSysPropSetter.getValue(),
-                DEFAULT_DEVICE_STATE.getIdentifier() + ":" + DEFAULT_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(DEFAULT_DEVICE_STATE));
-        assertFalse(mService.getOverrideBaseState().isPresent());
-        assertFalse(mService.getOverrideState().isPresent());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                DEFAULT_DEVICE_STATE.getIdentifier());
+        assertThat(mSysPropSetter.getValue()).isEqualTo(DEFAULT_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(DEFAULT_DEVICE_STATE);
+        assertThat(mService.getOverrideBaseState()).isEmpty();
+        assertThat(mService.getOverrideState()).isEmpty();
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(DEFAULT_DEVICE_STATE_IDENTIFIER);
 
         waitAndAssert(() -> callback.getLastNotifiedInfo().baseState.getIdentifier()
-                == DEFAULT_DEVICE_STATE.getIdentifier());
-        assertEquals(callback.getLastNotifiedInfo().currentState, DEFAULT_DEVICE_STATE);
+                == DEFAULT_DEVICE_STATE_IDENTIFIER);
+        assertThat(callback.getLastNotifiedInfo().currentState).isEqualTo(DEFAULT_DEVICE_STATE);
     }
 
     @Test
     public void requestBaseStateOverride_cancelledByBaseStateUpdate() throws RemoteException {
-        final DeviceState testDeviceState = new DeviceState(new DeviceState.Configuration.Builder(2,
-                "TEST").build());
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final DeviceState testDeviceState = new DeviceState(
+                new DeviceState.Configuration.Builder(2, "TEST").build());
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         mProvider.notifySupportedDeviceStates(new DeviceState[]{DEFAULT_DEVICE_STATE,
                 OTHER_DEVICE_STATE, testDeviceState});
         flushHandler();
 
         final IBinder token = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
         mService.getBinderService().requestBaseStateOverride(token,
-                OTHER_DEVICE_STATE.getIdentifier(),
+                OTHER_DEVICE_STATE_IDENTIFIER,
                 0 /* flags */);
 
         waitAndAssert(() -> callback.getLastNotifiedStatus(token)
                 == TestDeviceStateManagerCallback.STATUS_ACTIVE);
         // Committed state changes as there is a requested override.
-        assertEquals(mService.getCommittedState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mSysPropSetter.getValue(),
-                OTHER_DEVICE_STATE.getIdentifier() + ":" + OTHER_DEVICE_STATE.getName());
-        assertEquals(mService.getBaseState(), Optional.of(OTHER_DEVICE_STATE));
-        assertEquals(mService.getOverrideBaseState().get(), OTHER_DEVICE_STATE);
-        assertFalse(mService.getOverrideState().isPresent());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                OTHER_DEVICE_STATE.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mSysPropSetter.getValue()).isEqualTo(OTHER_DEVICE_STATE_TRACE_STRING);
+        assertThat(mService.getBaseState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getOverrideBaseState()).hasValue(OTHER_DEVICE_STATE);
+        assertThat(mService.getOverrideState()).isEmpty();
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(OTHER_DEVICE_STATE_IDENTIFIER);
 
-        assertNotNull(callback.getLastNotifiedInfo());
-        assertEquals(callback.getLastNotifiedInfo().baseState, OTHER_DEVICE_STATE);
-        assertEquals(callback.getLastNotifiedInfo().currentState, OTHER_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo()).isNotNull();
+        assertThat(callback.getLastNotifiedInfo().baseState).isEqualTo(OTHER_DEVICE_STATE);
+        assertThat(callback.getLastNotifiedInfo().currentState).isEqualTo(OTHER_DEVICE_STATE);
 
         mProvider.setState(testDeviceState.getIdentifier());
 
@@ -772,22 +769,22 @@
                 == TestDeviceStateManagerCallback.STATUS_CANCELED);
         // Committed state is set to the new base state once the override is cleared.
         waitAndAssert(() -> mService.getCommittedState().equals(Optional.of(testDeviceState)));
-        assertEquals(mSysPropSetter.getValue(),
-                testDeviceState.getIdentifier() + ":" + testDeviceState.getName());
-        assertEquals(mService.getBaseState(), Optional.of(testDeviceState));
-        assertFalse(mService.getOverrideBaseState().isPresent());
-        assertFalse(mService.getOverrideState().isPresent());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                testDeviceState.getIdentifier());
+        assertThat(mSysPropSetter.getValue())
+                .isEqualTo(testDeviceState.getIdentifier() + ":" + testDeviceState.getName());
+        assertThat(mService.getBaseState()).hasValue(testDeviceState);
+        assertThat(mService.getOverrideBaseState()).isEmpty();
+        assertThat(mService.getOverrideState()).isEmpty();
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(testDeviceState.getIdentifier());
 
         waitAndAssert(() -> callback.getLastNotifiedInfo().baseState.getIdentifier()
                 == testDeviceState.getIdentifier());
-        assertEquals(callback.getLastNotifiedInfo().currentState, testDeviceState);
+        assertThat(callback.getLastNotifiedInfo().currentState).isEqualTo(testDeviceState);
     }
 
     @Test
     public void requestBaseState_unsupportedState() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
 
         assertThrows(IllegalArgumentException.class, () -> {
@@ -799,7 +796,7 @@
 
     @Test
     public void requestBaseState_invalidState() throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
 
         assertThrows(IllegalArgumentException.class, () -> {
@@ -814,7 +811,7 @@
         assertThrows(IllegalStateException.class, () -> {
             final IBinder token = new Binder();
             mService.getBinderService().requestBaseStateOverride(token,
-                    DEFAULT_DEVICE_STATE.getIdentifier(),
+                    DEFAULT_DEVICE_STATE_IDENTIFIER,
                     0 /* flags */);
         });
     }
@@ -834,21 +831,21 @@
             Runnable noChangeEvent,
             Runnable autoCancelEvent
     ) throws RemoteException {
-        TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
+        final TestDeviceStateManagerCallback callback = new TestDeviceStateManagerCallback();
         mService.getBinderService().registerCallback(callback);
         flushHandler();
 
         final IBinder token = new Binder();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_UNKNOWN);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_UNKNOWN);
 
         mService.getBinderService().requestState(token,
                 DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP.getIdentifier(),
                 0 /* flags */);
         flushHandler(2 /* count */);
 
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_ACTIVE);
 
         // Committed state changes as there is a requested override.
         assertDeviceStateConditions(
@@ -858,8 +855,8 @@
 
         noChangeEvent.run();
         flushHandler();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_ACTIVE);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_ACTIVE);
         assertDeviceStateConditions(
                 DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP,
                 DEFAULT_DEVICE_STATE, /* base state */
@@ -867,8 +864,8 @@
 
         autoCancelEvent.run();
         flushHandler();
-        assertEquals(callback.getLastNotifiedStatus(token),
-                TestDeviceStateManagerCallback.STATUS_CANCELED);
+        assertThat(callback.getLastNotifiedStatus(token))
+                .isEqualTo(TestDeviceStateManagerCallback.STATUS_CANCELED);
         assertDeviceStateConditions(DEFAULT_DEVICE_STATE, DEFAULT_DEVICE_STATE,
                 false /* isOverrideState */);
     }
@@ -881,20 +878,20 @@
      * @param isOverrideState whether a state override is active.
      */
     private void assertDeviceStateConditions(
-            DeviceState state, DeviceState baseState,
+            @NonNull DeviceState state, @NonNull DeviceState baseState,
             boolean isOverrideState) {
-        assertEquals(mService.getCommittedState(), Optional.of(state));
-        assertEquals(mService.getBaseState(), Optional.of(baseState));
-        assertEquals(mSysPropSetter.getValue(),
-                state.getIdentifier() + ":" + state.getName());
-        assertEquals(mPolicy.getMostRecentRequestedStateToConfigure(),
-                state.getIdentifier());
+        assertThat(mService.getCommittedState()).hasValue(state);
+        assertThat(mService.getBaseState()).hasValue(baseState);
+        assertThat(mSysPropSetter.getValue())
+                .isEqualTo(state.getIdentifier() + ":" + state.getName());
+        assertThat(mPolicy.getMostRecentRequestedStateToConfigure())
+                .isEqualTo(state.getIdentifier());
         if (isOverrideState) {
             // When a state override is active, the committed state should batch the override state.
-            assertEquals(mService.getOverrideState().get(), state);
+            assertThat(mService.getOverrideState()).hasValue(state);
         } else {
             // When there is no state override, the override state should be empty.
-            assertFalse(mService.getOverrideState().isPresent());
+            assertThat(mService.getOverrideState()).isEmpty();
         }
     }
 
@@ -902,10 +899,11 @@
         private final DeviceStateProvider mProvider;
         private int mLastDeviceStateRequestedToConfigure = INVALID_DEVICE_STATE_IDENTIFIER;
         private boolean mConfigureBlocked = false;
+        @Nullable
         private Runnable mPendingConfigureCompleteRunnable;
 
-        TestDeviceStatePolicy(DeviceStateProvider provider) {
-            super(InstrumentationRegistry.getContext());
+        TestDeviceStatePolicy(@NonNull Context context, @NonNull DeviceStateProvider provider) {
+            super(context);
             mProvider = provider;
         }
 
@@ -921,7 +919,7 @@
         public void resumeConfigure() {
             mConfigureBlocked = false;
             if (mPendingConfigureCompleteRunnable != null) {
-                Runnable onComplete = mPendingConfigureCompleteRunnable;
+                final Runnable onComplete = mPendingConfigureCompleteRunnable;
                 mPendingConfigureCompleteRunnable = null;
                 onComplete.run();
             }
@@ -929,7 +927,7 @@
 
         public void resumeConfigureOnce() {
             if (mPendingConfigureCompleteRunnable != null) {
-                Runnable onComplete = mPendingConfigureCompleteRunnable;
+                final Runnable onComplete = mPendingConfigureCompleteRunnable;
                 mPendingConfigureCompleteRunnable = null;
                 onComplete.run();
             }
@@ -940,7 +938,7 @@
         }
 
         @Override
-        public void configureDeviceForState(int state, Runnable onComplete) {
+        public void configureDeviceForState(int state, @NonNull Runnable onComplete) {
             if (mPendingConfigureCompleteRunnable != null) {
                 throw new IllegalStateException("configureDeviceForState() called while configure"
                         + " is pending");
@@ -966,7 +964,9 @@
                         OTHER_DEVICE_STATE,
                         DEVICE_STATE_CANCEL_WHEN_REQUESTER_NOT_ON_TOP};
 
-        @Nullable private final DeviceState mInitialState;
+        @Nullable
+        private final DeviceState mInitialState;
+        @Nullable
         private Listener mListener;
 
         private TestDeviceStateProvider() {
@@ -991,7 +991,7 @@
             }
         }
 
-        public void notifySupportedDeviceStates(DeviceState[] supportedDeviceStates) {
+        public void notifySupportedDeviceStates(@NonNull DeviceState[] supportedDeviceStates) {
             mSupportedDeviceStates = supportedDeviceStates;
             mListener.onSupportedDeviceStatesChanged(supportedDeviceStates,
                     SUPPORTED_DEVICE_STATES_CHANGED_INITIALIZED);
@@ -1018,17 +1018,17 @@
         private final HashMap<IBinder, Integer> mLastNotifiedStatus = new HashMap<>();
 
         @Override
-        public void onDeviceStateInfoChanged(DeviceStateInfo info) {
+        public void onDeviceStateInfoChanged(@NonNull DeviceStateInfo info) {
             mLastNotifiedInfo = info;
         }
 
         @Override
-        public void onRequestActive(IBinder token) {
+        public void onRequestActive(@NonNull IBinder token) {
             mLastNotifiedStatus.put(token, STATUS_ACTIVE);
         }
 
         @Override
-        public void onRequestCanceled(IBinder token) {
+        public void onRequestCanceled(@NonNull IBinder token) {
             mLastNotifiedStatus.put(token, STATUS_CANCELED);
         }
 
@@ -1041,20 +1041,22 @@
             mLastNotifiedInfo = null;
         }
 
-        int getLastNotifiedStatus(IBinder requestToken) {
+        int getLastNotifiedStatus(@NonNull IBinder requestToken) {
             return mLastNotifiedStatus.getOrDefault(requestToken, STATUS_UNKNOWN);
         }
     }
 
     private static final class TestSystemPropertySetter implements
             DeviceStateManagerService.SystemPropertySetter {
+        @NonNull
         private String mValue;
 
         @Override
-        public void setDebugTracingDeviceStateProperty(String value) {
+        public void setDebugTracingDeviceStateProperty(@NonNull String value) {
             mValue = value;
         }
 
+        @NonNull
         public String getValue() {
             return mValue;
         }
diff --git a/services/tests/servicestests/src/com/android/server/stats/pull/netstats/NetworkStatsAccumulatorTest.kt b/services/tests/servicestests/src/com/android/server/stats/pull/netstats/NetworkStatsAccumulatorTest.kt
new file mode 100644
index 0000000..8cf0e82
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/stats/pull/netstats/NetworkStatsAccumulatorTest.kt
@@ -0,0 +1,189 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.stats.pull.netstats
+
+import android.net.NetworkStats
+import android.net.NetworkStats.DEFAULT_NETWORK_YES
+import android.net.NetworkStats.METERED_NO
+import android.net.NetworkStats.ROAMING_NO
+import android.net.NetworkStats.SET_DEFAULT
+import android.net.NetworkStats.TAG_NONE
+import android.net.NetworkTemplate
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.testutils.assertNetworkStatsEquals
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class NetworkStatsAccumulatorTest {
+
+    @Test
+    fun hasEqualParameters_differentParameters_returnsFalse() {
+        val wifiTemplate = NetworkTemplate.Builder(NetworkTemplate.MATCH_WIFI).build()
+        val mobileTemplate = NetworkTemplate.Builder(NetworkTemplate.MATCH_MOBILE).build()
+
+        val snapshot = NetworkStatsAccumulator(wifiTemplate, false, 0, 0)
+
+        assertThat(snapshot.hasEqualParameters(mobileTemplate, false)).isFalse()
+        assertThat(snapshot.hasEqualParameters(wifiTemplate, true)).isFalse()
+    }
+
+    @Test
+    fun hasSameParameters_equalParameters_returnsTrue() {
+        val wifiTemplate1 = NetworkTemplate.Builder(NetworkTemplate.MATCH_WIFI).build()
+        val wifiTemplate2 = NetworkTemplate.Builder(NetworkTemplate.MATCH_WIFI).build()
+
+        val snapshot = NetworkStatsAccumulator(wifiTemplate1, false, 0, 0)
+
+        assertThat(snapshot.hasEqualParameters(wifiTemplate1, false)).isTrue()
+        assertThat(snapshot.hasEqualParameters(wifiTemplate2, false)).isTrue()
+    }
+
+    @Test
+    fun queryStats_lessThanOneBucketFromSnapshotEndTime_returnsAllStats() {
+        val snapshot = NetworkStatsAccumulator(TEMPLATE, false, 200, 1000)
+
+        // Accumulator has data until 1000 (= 0), and its end-point is still in the history period.
+        // Current time is less than one bucket away from snapshot end-point: 1050 - 1000 < 200
+        val stats = snapshot.queryStats(1050, FakeStats(500, 1050, 1))
+
+        // After the query at 1050, accumulator should have 1 * (1050 - 1000) = 50 bytes.
+        assertNetworkStatsEquals(stats, networkStatsWithBytes(50))
+    }
+
+    @Test
+    fun queryStats_oneBucketFromSnapshotEndTime_returnsCumulativeStats() {
+        val snapshot = NetworkStatsAccumulator(TEMPLATE, false, 200, 1000)
+
+        // Accumulator has data until 1000 (= 0), and its end-point is still in the history period.
+        // Current time is one bucket away from snapshot end-point: 1250 - 1000 > 200
+        val stats = snapshot.queryStats(1250, FakeStats(550, 1250, 2))
+
+        // After the query at 1250, accumulator should have 2 * (1250 - 1000) = 500 bytes.
+        assertNetworkStatsEquals(stats, networkStatsWithBytes(500))
+    }
+
+    @Test
+    fun queryStats_twoBucketsFromSnapshotEndTime_returnsCumulativeStats() {
+        val snapshot = NetworkStatsAccumulator(TEMPLATE, false, 200, 1000)
+
+        // Accumulator has data until 1000 (= 0), and its end-point is in the history period.
+        // Current time is two buckets away from snapshot end-point: 1450 - 1000 > 2*200
+        val stats = snapshot.queryStats(1450, FakeStats(600, 1450, 3))
+
+        // After the query at 1450, accumulator should have 3 * (1450 - 1000) = 1350 bytes.
+        assertNetworkStatsEquals(stats, networkStatsWithBytes(1350))
+    }
+
+    @Test
+    fun queryStats_manyBucketsFromSnapshotEndTime_returnsCumulativeStats() {
+        val snapshot = NetworkStatsAccumulator(TEMPLATE, false, 200, 1000)
+
+        // Accumulator has data until 1000 (= 0), and its end-point is still in the history period.
+        // Current time is many buckets away from snapshot end-point
+        val stats = snapshot.queryStats(6100, FakeStats(900, 6100, 1))
+
+        // After the query at 6100, accumulator should have 1 * (6100 - 1000) = 5100 bytes.
+        assertNetworkStatsEquals(stats, networkStatsWithBytes(5100))
+    }
+
+    @Test
+    fun queryStats_multipleQueriesAndSameHistoryWindow_returnsCumulativeStats() {
+        val snapshot = NetworkStatsAccumulator(TEMPLATE, false, 200, 1000)
+
+        // Accumulator is queried within the history period, whose starting point stays the same.
+        // After each query, accumulator should contain bytes from the initial end-point until now.
+        val stats1 = snapshot.queryStats(5100, FakeStats(900, 5100, 1))
+        val stats2 = snapshot.queryStats(10100, FakeStats(900, 10100, 1))
+        val stats3 = snapshot.queryStats(15100, FakeStats(900, 15100, 1))
+
+        assertNetworkStatsEquals(stats1, networkStatsWithBytes(4100))
+        assertNetworkStatsEquals(stats2, networkStatsWithBytes(9100))
+        assertNetworkStatsEquals(stats3, networkStatsWithBytes(14100))
+    }
+
+    @Test
+    fun queryStats_multipleQueriesAndSlidingHistoryWindow_returnsCumulativeStats() {
+        val snapshot = NetworkStatsAccumulator(TEMPLATE, false, 200, 1000)
+
+        // Accumulator is queried within the history period, whose starting point is moving.
+        // After each query, accumulator should contain bytes from the initial end-point until now.
+        val stats1 = snapshot.queryStats(5100, FakeStats(900, 5100, 1))
+        val stats2 = snapshot.queryStats(10100, FakeStats(4000, 10100, 1))
+        val stats3 = snapshot.queryStats(15100, FakeStats(7000, 15100, 1))
+
+        assertNetworkStatsEquals(stats1, networkStatsWithBytes(4100))
+        assertNetworkStatsEquals(stats2, networkStatsWithBytes(9100))
+        assertNetworkStatsEquals(stats3, networkStatsWithBytes(14100))
+    }
+
+    @Test
+    fun queryStats_withSnapshotEndTimeBeforeHistoryStart_addsOnlyStatsWithinHistory() {
+        val snapshot = NetworkStatsAccumulator(TEMPLATE, false, 200, 1900)
+
+        // Accumulator has data until 1000 (= 0), but its end-point is not in the history period.
+        // After the query, accumulator should add only those bytes that are covered by the history.
+        val stats = snapshot.queryStats(2700, FakeStats(2200, 2700, 1))
+
+        assertNetworkStatsEquals(stats, networkStatsWithBytes(500))
+    }
+
+    /**
+     * Simulates equally distributed traffic stats persisted over a set period of time.
+     */
+    private class FakeStats(
+        val historyStartMillis: Long, val currentTimeMillis: Long, val bytesPerMilli: Long
+    ) : NetworkStatsAccumulator.StatsQueryFunction {
+
+        override fun queryNetworkStats(
+            template: NetworkTemplate, includeTags: Boolean, startTime: Long, endTime: Long
+        ): NetworkStats {
+            val overlap = overlap(startTime, endTime, historyStartMillis, currentTimeMillis)
+            return networkStatsWithBytes(overlap * bytesPerMilli)
+        }
+    }
+
+    companion object {
+
+        private val TEMPLATE = NetworkTemplate.Builder(NetworkTemplate.MATCH_WIFI).build()
+
+        fun networkStatsWithBytes(bytes: Long): NetworkStats {
+            val stats = NetworkStats(0, 1).addEntry(
+                NetworkStats.Entry(
+                    null,
+                    0,
+                    SET_DEFAULT,
+                    TAG_NONE,
+                    METERED_NO,
+                    ROAMING_NO,
+                    DEFAULT_NETWORK_YES,
+                    bytes,
+                    bytes / 100,
+                    bytes,
+                    bytes / 100,
+                    0
+                )
+            )
+            return stats
+        }
+
+        fun overlap(aStart: Long, aEnd: Long, bStart: Long, bEnd: Long): Long {
+            return maxOf(0L, minOf(aEnd, bEnd) - maxOf(aStart, bStart))
+        }
+    }
+}
\ No newline at end of file
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
index 3bbc6b2..48bc9d7 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
@@ -63,13 +63,11 @@
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.IInterface;
-import android.os.Looper;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.platform.test.annotations.EnableFlags;
 import android.provider.Settings;
-import android.testing.TestableLooper;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.ArraySet;
@@ -84,7 +82,6 @@
 
 import com.google.android.collect.Lists;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
@@ -106,7 +103,6 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 
-
 public class ManagedServicesTest extends UiServiceTestCase {
 
     @Mock
@@ -119,7 +115,6 @@
     private ManagedServices.UserProfiles mUserProfiles;
     @Mock private DevicePolicyManager mDpm;
     Object mLock = new Object();
-    private TestableLooper mTestableLooper;
 
     UserInfo mZero = new UserInfo(0, "zero", 0);
     UserInfo mTen = new UserInfo(10, "ten", 0);
@@ -147,7 +142,6 @@
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mTestableLooper = new TestableLooper(Looper.getMainLooper());
 
         mContext.setMockPackageManager(mPm);
         mContext.addMockSystemService(Context.USER_SERVICE, mUm);
@@ -205,11 +199,6 @@
                 mIpm, APPROVAL_BY_COMPONENT);
     }
 
-    @After
-    public void tearDown() throws Exception {
-        mTestableLooper.destroy();
-    }
-
     @Test
     public void testBackupAndRestore_migration() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
@@ -899,7 +888,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
         service.addApprovedList("a", 0, true);
 
         service.reregisterService(cn, 0);
@@ -930,7 +919,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
         service.addApprovedList("a", 0, false);
 
         service.reregisterService(cn, 0);
@@ -961,7 +950,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
         service.addApprovedList("a/a", 0, true);
 
         service.reregisterService(cn, 0);
@@ -992,7 +981,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
         service.addApprovedList("a/a", 0, false);
 
         service.reregisterService(cn, 0);
@@ -1064,77 +1053,6 @@
     }
 
     @Test
-    public void registerService_bindingDied_rebindIsClearedOnUserSwitch() throws Exception {
-        Context context = mock(Context.class);
-        PackageManager pm = mock(PackageManager.class);
-        ApplicationInfo ai = new ApplicationInfo();
-        ai.targetSdkVersion = Build.VERSION_CODES.CUR_DEVELOPMENT;
-
-        when(context.getPackageName()).thenReturn(mPkg);
-        when(context.getUserId()).thenReturn(mUser.getIdentifier());
-        when(context.getPackageManager()).thenReturn(pm);
-        when(pm.getApplicationInfo(anyString(), anyInt())).thenReturn(ai);
-
-        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles, mIpm,
-                APPROVAL_BY_PACKAGE);
-        service = spy(service);
-        ComponentName cn = ComponentName.unflattenFromString("a/a");
-
-        // Trigger onBindingDied for component when registering
-        //  => will schedule a rebind in 10 seconds
-        when(context.bindServiceAsUser(any(), any(), anyInt(), any())).thenAnswer(invocation -> {
-            Object[] args = invocation.getArguments();
-            ServiceConnection sc = (ServiceConnection) args[1];
-            sc.onBindingDied(cn);
-            return true;
-        });
-        service.registerService(cn, 0);
-        assertThat(service.isBound(cn, 0)).isFalse();
-
-        // Switch to user 10
-        service.onUserSwitched(10);
-
-        // Check that the scheduled rebind for user 0 was cleared
-        mTestableLooper.moveTimeForward(ManagedServices.ON_BINDING_DIED_REBIND_DELAY_MS);
-        mTestableLooper.processAllMessages();
-        verify(service, never()).reregisterService(any(), anyInt());
-    }
-
-    @Test
-    public void registerService_bindingDied_rebindIsExecutedAfterTimeout() throws Exception {
-        Context context = mock(Context.class);
-        PackageManager pm = mock(PackageManager.class);
-        ApplicationInfo ai = new ApplicationInfo();
-        ai.targetSdkVersion = Build.VERSION_CODES.CUR_DEVELOPMENT;
-
-        when(context.getPackageName()).thenReturn(mPkg);
-        when(context.getUserId()).thenReturn(mUser.getIdentifier());
-        when(context.getPackageManager()).thenReturn(pm);
-        when(pm.getApplicationInfo(anyString(), anyInt())).thenReturn(ai);
-
-        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles, mIpm,
-                APPROVAL_BY_PACKAGE);
-        service = spy(service);
-        ComponentName cn = ComponentName.unflattenFromString("a/a");
-
-        // Trigger onBindingDied for component when registering
-        //  => will schedule a rebind in 10 seconds
-        when(context.bindServiceAsUser(any(), any(), anyInt(), any())).thenAnswer(invocation -> {
-            Object[] args = invocation.getArguments();
-            ServiceConnection sc = (ServiceConnection) args[1];
-            sc.onBindingDied(cn);
-            return true;
-        });
-        service.registerService(cn, 0);
-        assertThat(service.isBound(cn, 0)).isFalse();
-
-        // Check that the scheduled rebind is run
-        mTestableLooper.moveTimeForward(ManagedServices.ON_BINDING_DIED_REBIND_DELAY_MS);
-        mTestableLooper.processAllMessages();
-        verify(service, times(1)).reregisterService(eq(cn), eq(0));
-    }
-
-    @Test
     public void testPackageUninstall_packageNoLongerInApprovedList() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1293,64 +1211,6 @@
     }
 
     @Test
-    public void testUpgradeAppNoIntentFilterNoRebind() throws Exception {
-        Context context = spy(getContext());
-        doReturn(true).when(context).bindServiceAsUser(any(), any(), anyInt(), any());
-
-        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles,
-                mIpm, APPROVAL_BY_COMPONENT);
-
-        List<String> packages = new ArrayList<>();
-        packages.add("package");
-        addExpectedServices(service, packages, 0);
-
-        final ComponentName unapprovedComponent = ComponentName.unflattenFromString("package/C1");
-        final ComponentName approvedComponent = ComponentName.unflattenFromString("package/C2");
-
-        // Both components are approved initially
-        mExpectedPrimaryComponentNames.clear();
-        mExpectedPrimaryPackages.clear();
-        mExpectedPrimaryComponentNames.put(0, "package/C1:package/C2");
-        mExpectedSecondaryComponentNames.clear();
-        mExpectedSecondaryPackages.clear();
-
-        loadXml(service);
-
-        //Component package/C1 loses serviceInterface intent filter
-        ManagedServices.Config config = service.getConfig();
-        when(mPm.queryIntentServicesAsUser(any(), anyInt(), anyInt()))
-                .thenAnswer(new Answer<List<ResolveInfo>>() {
-                    @Override
-                    public List<ResolveInfo> answer(InvocationOnMock invocationOnMock)
-                            throws Throwable {
-                        Object[] args = invocationOnMock.getArguments();
-                        Intent invocationIntent = (Intent) args[0];
-                        if (invocationIntent != null) {
-                            if (invocationIntent.getAction().equals(config.serviceInterface)
-                                    && packages.contains(invocationIntent.getPackage())) {
-                                List<ResolveInfo> dummyServices = new ArrayList<>();
-                                ResolveInfo resolveInfo = new ResolveInfo();
-                                ServiceInfo serviceInfo = new ServiceInfo();
-                                serviceInfo.packageName = invocationIntent.getPackage();
-                                serviceInfo.name = approvedComponent.getClassName();
-                                serviceInfo.permission = service.getConfig().bindPermission;
-                                resolveInfo.serviceInfo = serviceInfo;
-                                dummyServices.add(resolveInfo);
-                                return dummyServices;
-                            }
-                        }
-                        return new ArrayList<>();
-                    }
-                });
-
-        // Trigger package update
-        service.onPackagesChanged(false, new String[]{"package"}, new int[]{0});
-
-        assertFalse(service.isComponentEnabledForCurrentProfiles(unapprovedComponent));
-        assertTrue(service.isComponentEnabledForCurrentProfiles(approvedComponent));
-    }
-
-    @Test
     public void testSetPackageOrComponentEnabled() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1363,21 +1223,6 @@
                             "user10package1/K", "user10.3/Component", "user10package2/L",
                             "user10.4/Component"}));
 
-            // mock permissions for services
-            PackageManager pm = mock(PackageManager.class);
-            when(getContext().getPackageManager()).thenReturn(pm);
-            List<ComponentName> enabledComponents = List.of(
-                    ComponentName.unflattenFromString("package/Comp"),
-                    ComponentName.unflattenFromString("package/C2"),
-                    ComponentName.unflattenFromString("again/M4"),
-                    ComponentName.unflattenFromString("user10package/B"),
-                    ComponentName.unflattenFromString("user10/Component"),
-                    ComponentName.unflattenFromString("user10package1/K"),
-                    ComponentName.unflattenFromString("user10.3/Component"),
-                    ComponentName.unflattenFromString("user10package2/L"),
-                    ComponentName.unflattenFromString("user10.4/Component"));
-            mockServiceInfoWithMetaData(enabledComponents, service, pm, new ArrayMap<>());
-
             for (int userId : expectedEnabled.keySet()) {
                 ArrayList<String> expectedForUser = expectedEnabled.get(userId);
                 for (int i = 0; i < expectedForUser.size(); i++) {
@@ -2099,7 +1944,7 @@
         metaDataAutobindAllow.putBoolean(META_DATA_DEFAULT_AUTOBIND, true);
         metaDatas.put(cn_allowed, metaDataAutobindAllow);
 
-        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
 
         service.addApprovedList(cn_allowed.flattenToString(), 0, true);
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
@@ -2144,7 +1989,7 @@
         metaDataAutobindDisallow.putBoolean(META_DATA_DEFAULT_AUTOBIND, false);
         metaDatas.put(cn_disallowed, metaDataAutobindDisallow);
 
-        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
 
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
 
@@ -2183,7 +2028,7 @@
         metaDataAutobindDisallow.putBoolean(META_DATA_DEFAULT_AUTOBIND, false);
         metaDatas.put(cn_disallowed, metaDataAutobindDisallow);
 
-        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
 
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
 
@@ -2254,8 +2099,8 @@
     }
 
     private void mockServiceInfoWithMetaData(List<ComponentName> componentNames,
-            ManagedServices service, PackageManager packageManager,
-            ArrayMap<ComponentName, Bundle> metaDatas) throws RemoteException {
+            ManagedServices service, ArrayMap<ComponentName, Bundle> metaDatas)
+            throws RemoteException {
         when(mIpm.getServiceInfo(any(), anyLong(), anyInt())).thenAnswer(
                 (Answer<ServiceInfo>) invocation -> {
                     ComponentName invocationCn = invocation.getArgument(0);
@@ -2270,39 +2115,6 @@
                     return null;
                 }
         );
-
-        // add components to queryIntentServicesAsUser response
-        final List<String> packages = new ArrayList<>();
-        for (ComponentName cn: componentNames) {
-            packages.add(cn.getPackageName());
-        }
-        ManagedServices.Config config = service.getConfig();
-        when(packageManager.queryIntentServicesAsUser(any(), anyInt(), anyInt())).
-                thenAnswer(new Answer<List<ResolveInfo>>() {
-                @Override
-                public List<ResolveInfo> answer(InvocationOnMock invocationOnMock)
-                    throws Throwable {
-                    Object[] args = invocationOnMock.getArguments();
-                    Intent invocationIntent = (Intent) args[0];
-                    if (invocationIntent != null) {
-                        if (invocationIntent.getAction().equals(config.serviceInterface)
-                            && packages.contains(invocationIntent.getPackage())) {
-                            List<ResolveInfo> dummyServices = new ArrayList<>();
-                            for (ComponentName cn: componentNames) {
-                                ResolveInfo resolveInfo = new ResolveInfo();
-                                ServiceInfo serviceInfo = new ServiceInfo();
-                                serviceInfo.packageName = invocationIntent.getPackage();
-                                serviceInfo.name = cn.getClassName();
-                                serviceInfo.permission = service.getConfig().bindPermission;
-                                resolveInfo.serviceInfo = serviceInfo;
-                                dummyServices.add(resolveInfo);
-                            }
-                            return dummyServices;
-                        }
-                    }
-                    return new ArrayList<>();
-                }
-            });
     }
 
     private void resetComponentsAndPackages() {
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
index 7e4ae67..797b95b5 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
@@ -25,7 +25,6 @@
 import static junit.framework.Assert.assertTrue;
 
 import static org.junit.Assert.assertNull;
-
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.any;
@@ -194,8 +193,6 @@
     public void testWriteXml_userTurnedOffNAS() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
-        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
-
         mAssistants.loadDefaultsFromConfig(true);
 
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
@@ -401,10 +398,6 @@
     public void testSetPackageOrComponentEnabled_onlyOnePackage() throws Exception {
         ComponentName component1 = ComponentName.unflattenFromString("package/Component1");
         ComponentName component2 = ComponentName.unflattenFromString("package/Component2");
-
-        doReturn(true).when(mAssistants).isValidService(eq(component1), eq(mZero.id));
-        doReturn(true).when(mAssistants).isValidService(eq(component2), eq(mZero.id));
-
         mAssistants.setPackageOrComponentEnabled(component1.flattenToString(), mZero.id, true,
                 true, true);
         verify(mNm, never()).setNotificationAssistantAccessGrantedForUserInternal(
@@ -550,7 +543,6 @@
     public void testSetAdjustmentTypeSupportedState() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
-        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
@@ -574,7 +566,6 @@
     public void testSetAdjustmentTypeSupportedState_readWriteXml_entries() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
-        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
@@ -598,7 +589,6 @@
     public void testSetAdjustmentTypeSupportedState_readWriteXml_empty() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
-        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatUtilsTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatUtilsTest.java
index d8373c5..50c2c2f 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatUtilsTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatUtilsTest.java
@@ -18,16 +18,26 @@
 
 import static android.app.CameraCompatTaskInfo.CAMERA_COMPAT_FREEFORM_PORTRAIT_DEVICE_IN_LANDSCAPE;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doAnswer;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
 import android.app.CameraCompatTaskInfo.FreeformCameraCompatMode;
 import android.app.TaskInfo;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
+import android.view.DisplayInfo;
+import android.view.Surface;
 
 import androidx.annotation.NonNull;
 
+import com.android.window.flags.Flags;
+
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -174,9 +184,13 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     public void getTaskInfoPropagatesCameraCompatMode() {
         runTestScenario((robot) -> {
-            robot.applyOnActivity(AppCompatActivityRobot::createActivityWithComponentInNewTask);
+            robot.dw().allowEnterDesktopMode(/* isAllowed= */ true);
+            robot.applyOnActivity(
+                    AppCompatActivityRobot::createActivityWithComponentInNewTaskAndDisplay);
+            robot.setCameraCompatTreatmentEnabledForActivity(/* enabled= */ true);
 
             robot.setFreeformCameraCompatMode(CAMERA_COMPAT_FREEFORM_PORTRAIT_DEVICE_IN_LANDSCAPE);
             robot.checkTaskInfoFreeformCameraCompatMode(
@@ -212,6 +226,15 @@
             spyOn(activity.mAppCompatController.getAppCompatAspectRatioPolicy());
         }
 
+        @Override
+        void onPostDisplayContentCreation(@NonNull DisplayContent displayContent) {
+            super.onPostDisplayContentCreation(displayContent);
+            mockPortraitDisplay(displayContent);
+            if (displayContent.mAppCompatCameraPolicy.hasCameraCompatFreeformPolicy()) {
+                spyOn(displayContent.mAppCompatCameraPolicy.mCameraCompatFreeformPolicy);
+            }
+        }
+
         void transparentActivity(@NonNull Consumer<AppCompatTransparentActivityRobot> consumer) {
             // We always create at least an opaque activity in a Task.
             activity().createNewTaskWithBaseActivity();
@@ -235,8 +258,8 @@
         }
 
         void setFreeformCameraCompatMode(@FreeformCameraCompatMode int mode) {
-            activity().top().mAppCompatController.getAppCompatCameraOverrides()
-                    .setFreeformCameraCompatMode(mode);
+            doReturn(mode).when(activity().top().mDisplayContent.mAppCompatCameraPolicy
+                    .mCameraCompatFreeformPolicy).getCameraCompatMode(activity().top());
         }
 
         void checkTopActivityLetterboxReason(@NonNull String expected) {
@@ -258,6 +281,24 @@
             Assert.assertEquals(mode, getTopTaskInfo().appCompatTaskInfo
                     .cameraCompatTaskInfo.freeformCameraCompatMode);
         }
-    }
 
+        void setCameraCompatTreatmentEnabledForActivity(boolean enabled) {
+            doReturn(enabled).when(activity().displayContent().mAppCompatCameraPolicy
+                    .mCameraCompatFreeformPolicy).isTreatmentEnabledForActivity(
+                            eq(activity().top()), anyBoolean());
+        }
+
+        private void mockPortraitDisplay(DisplayContent displayContent) {
+            doAnswer(invocation -> {
+                DisplayInfo displayInfo = new DisplayInfo();
+                displayContent.getDisplay().getDisplayInfo(displayInfo);
+                displayInfo.rotation = Surface.ROTATION_90;
+                // Set height and width so that the natural orientation (when rotation is 0) is
+                // portrait.
+                displayInfo.logicalHeight = 600;
+                displayInfo.logicalWidth =  800;
+                return displayInfo;
+            }).when(displayContent.mWmService.mDisplayManagerInternal).getDisplayInfo(anyInt());
+        }
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java b/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java
index a8ccf95..a07fd23 100644
--- a/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java
@@ -223,10 +223,13 @@
         setDisplayRotation(Surface.ROTATION_270);
 
         mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
-        callOnActivityConfigurationChanging(mActivity);
+        callOnActivityConfigurationChanging(mActivity, /* letterboxNew= */ true,
+                /* lastLetterbox= */ false);
         mCameraAvailabilityCallback.onCameraClosed(CAMERA_ID_1);
         mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
-        callOnActivityConfigurationChanging(mActivity);
+        // Activity is letterboxed from the previous configuration change.
+        callOnActivityConfigurationChanging(mActivity, /* letterboxNew= */ true,
+                /* lastLetterbox= */ true);
 
         assertInCameraCompatMode(CAMERA_COMPAT_FREEFORM_PORTRAIT_DEVICE_IN_LANDSCAPE);
         assertActivityRefreshRequested(/* refreshRequested */ true);
@@ -264,6 +267,48 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    public void testShouldRefreshActivity_appBoundsChanged_returnsTrue() {
+        configureActivity(SCREEN_ORIENTATION_PORTRAIT);
+        Configuration oldConfiguration = createConfiguration(/* letterbox= */ false);
+        Configuration newConfiguration = createConfiguration(/* letterbox= */ true);
+        mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
+
+        assertTrue(mCameraCompatFreeformPolicy.shouldRefreshActivity(mActivity, newConfiguration,
+                oldConfiguration));
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    public void testShouldRefreshActivity_displayRotationChanged_returnsTrue() {
+        configureActivity(SCREEN_ORIENTATION_PORTRAIT);
+        Configuration oldConfiguration = createConfiguration(/* letterbox= */ true);
+        Configuration newConfiguration = createConfiguration(/* letterbox= */ true);
+
+        oldConfiguration.windowConfiguration.setDisplayRotation(0);
+        newConfiguration.windowConfiguration.setDisplayRotation(90);
+        mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
+
+        assertTrue(mCameraCompatFreeformPolicy.shouldRefreshActivity(mActivity, newConfiguration,
+                oldConfiguration));
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    public void testShouldRefreshActivity_appBoundsNorDisplayChanged_returnsFalse() {
+        configureActivity(SCREEN_ORIENTATION_PORTRAIT);
+        Configuration oldConfiguration = createConfiguration(/* letterbox= */ true);
+        Configuration newConfiguration = createConfiguration(/* letterbox= */ true);
+
+        oldConfiguration.windowConfiguration.setDisplayRotation(0);
+        newConfiguration.windowConfiguration.setDisplayRotation(0);
+        mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
+
+        assertFalse(mCameraCompatFreeformPolicy.shouldRefreshActivity(mActivity, newConfiguration,
+                oldConfiguration));
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     public void testOnActivityConfigurationChanging_refreshDisabledViaFlag_noRefresh()
             throws Exception {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
@@ -306,6 +351,7 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     public void testGetCameraCompatAspectRatio_activityNotInCameraCompat_returnsDefaultAspRatio() {
         configureActivity(SCREEN_ORIENTATION_FULL_USER);
 
@@ -318,6 +364,7 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     public void testGetCameraCompatAspectRatio_activityInCameraCompat_returnsConfigAspectRatio() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         final float configAspectRatio = 1.5f;
@@ -331,8 +378,8 @@
                 /* delta= */ 0.001);
     }
 
-
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     public void testGetCameraCompatAspectRatio_inCameraCompatPerAppOverride_returnDefAspectRatio() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         final float configAspectRatio = 1.5f;
@@ -411,9 +458,15 @@
     }
 
     private void callOnActivityConfigurationChanging(ActivityRecord activity) {
+        callOnActivityConfigurationChanging(activity, /* letterboxNew= */ true,
+                /* lastLetterbox= */false);
+    }
+
+    private void callOnActivityConfigurationChanging(ActivityRecord activity, boolean letterboxNew,
+            boolean lastLetterbox) {
         mActivityRefresher.onActivityConfigurationChanging(activity,
-                /* newConfig */ createConfiguration(/*letterbox=*/ true),
-                /* lastReportedConfig */ createConfiguration(/*letterbox=*/ false));
+                /* newConfig */ createConfiguration(letterboxNew),
+                /* lastReportedConfig */ createConfiguration(lastLetterbox));
     }
 
     private Configuration createConfiguration(boolean letterbox) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
index 6111a65..8bbba1b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
@@ -94,6 +94,7 @@
 import android.view.ContentRecordingSession;
 import android.view.IWindow;
 import android.view.InputChannel;
+import android.view.InputDevice;
 import android.view.InsetsSourceControl;
 import android.view.InsetsState;
 import android.view.Surface;
@@ -1275,6 +1276,48 @@
     }
 
     @Test
+    public void testInputDeviceNotifyConfigurationChanged() {
+        mSetFlagsRule.enableFlags(Flags.FLAG_FILTER_IRRELEVANT_INPUT_DEVICE_CHANGE);
+        spyOn(mDisplayContent);
+        doReturn(false).when(mDisplayContent).sendNewConfiguration();
+        final InputDevice deviceA = mock(InputDevice.class);
+        final InputDevice deviceB = mock(InputDevice.class);
+        doReturn("deviceA").when(deviceA).getDescriptor();
+        doReturn("deviceB").when(deviceB).getDescriptor();
+        final InputDevice[] devices1 = { deviceA };
+        final InputDevice[] devices2 = { deviceB, deviceA };
+        final Runnable verifySendNewConfiguration = () -> {
+            clearInvocations(mDisplayContent);
+            mWm.mInputManagerCallback.notifyConfigurationChanged();
+            verify(mDisplayContent).sendNewConfiguration();
+        };
+        doReturn(devices1).when(mWm.mInputManager).getInputDevices();
+        verifySendNewConfiguration.run();
+
+        doReturn(devices2).when(mWm.mInputManager).getInputDevices();
+        verifySendNewConfiguration.run();
+
+        doReturn(true).when(deviceB).isEnabled();
+        verifySendNewConfiguration.run();
+
+        doReturn(true).when(deviceA).isExternal();
+        verifySendNewConfiguration.run();
+
+        doReturn(1).when(deviceA).getSources();
+        verifySendNewConfiguration.run();
+
+        doReturn(1).when(deviceA).getAssociatedDisplayId();
+        verifySendNewConfiguration.run();
+
+        doReturn(1).when(deviceA).getKeyboardType();
+        verifySendNewConfiguration.run();
+
+        clearInvocations(mDisplayContent);
+        mWm.mInputManagerCallback.notifyConfigurationChanged();
+        verify(mDisplayContent, never()).sendNewConfiguration();
+    }
+
+    @Test
     public void testReportSystemGestureExclusionChanged_invalidWindow() {
         final Session session = mock(Session.class);
         final IWindow window = mock(IWindow.class);
diff --git a/tests/FlickerTests/Android.bp b/tests/FlickerTests/Android.bp
index 27e9ffa..1e997b3 100644
--- a/tests/FlickerTests/Android.bp
+++ b/tests/FlickerTests/Android.bp
@@ -47,7 +47,6 @@
 
 java_library {
     name: "wm-flicker-common-assertions",
-    platform_apis: true,
     optimize: {
         enabled: false,
     },
diff --git a/tests/FlickerTests/test-apps/app-helpers/Android.bp b/tests/FlickerTests/test-apps/app-helpers/Android.bp
index fc4d71c..e8bb64a 100644
--- a/tests/FlickerTests/test-apps/app-helpers/Android.bp
+++ b/tests/FlickerTests/test-apps/app-helpers/Android.bp
@@ -25,7 +25,6 @@
 
 java_library {
     name: "wm-flicker-common-app-helpers",
-    platform_apis: true,
     optimize: {
         enabled: false,
     },
diff --git a/tests/graphics/SilkFX/res/layout/view_blur_behind.xml b/tests/graphics/SilkFX/res/layout/view_blur_behind.xml
new file mode 100644
index 0000000..83b1fa4
--- /dev/null
+++ b/tests/graphics/SilkFX/res/layout/view_blur_behind.xml
@@ -0,0 +1,148 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<FrameLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent">
+
+    <LinearLayout
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:orientation="vertical">
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="wowwowwowwowwowwowwowwowwowwowwowwowwowwowwow" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="I'm a little teapot" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="Something. Something." />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="/\\/\\/\\/\\/\\/\\/\\/\\/\\/" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="wowwowwowwowwowwowwowwowwowwowwowwowwowwowwow" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="I'm a little teapot" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="Something. Something." />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="/\\/\\/\\/\\/\\/\\/\\/\\/\\/" />
+
+        <TextView
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:padding="8dp"
+            android:textSize="24dp"
+            android:text="^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^.^" />
+
+    </LinearLayout>
+
+    <ScrollView
+        android:layout_width="match_parent"
+        android:layout_height="match_parent">
+
+        <LinearLayout
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            android:orientation="vertical">
+
+            <View
+                android:layout_width="match_parent"
+                android:layout_height="300dp" />
+
+            <com.android.test.silkfx.materials.BlurBehindContainer
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:background="#33AAAAAA"
+                android:padding="32dp">
+
+                <TextView
+                    android:layout_width="wrap_content"
+                    android:layout_height="wrap_content"
+                    android:textSize="48dp"
+                    android:text="Blur!" />
+
+            </com.android.test.silkfx.materials.BlurBehindContainer>
+
+            <View
+                android:layout_width="match_parent"
+                android:layout_height="1024dp" />
+
+        </LinearLayout>
+
+    </ScrollView>
+
+</FrameLayout>
\ No newline at end of file
diff --git a/tests/graphics/SilkFX/src/com/android/test/silkfx/Main.kt b/tests/graphics/SilkFX/src/com/android/test/silkfx/Main.kt
index 59a6078..6b6d3b8 100644
--- a/tests/graphics/SilkFX/src/com/android/test/silkfx/Main.kt
+++ b/tests/graphics/SilkFX/src/com/android/test/silkfx/Main.kt
@@ -61,7 +61,8 @@
         )),
         DemoGroup("Materials", listOf(
                 Demo("Glass", GlassActivity::class),
-                Demo("Background Blur", BackgroundBlurActivity::class)
+                Demo("Background Blur", BackgroundBlurActivity::class),
+                Demo("View blur behind", R.layout.view_blur_behind, commonControls = false)
         ))
 )
 
diff --git a/tests/graphics/SilkFX/src/com/android/test/silkfx/materials/BlurBehindContainer.kt b/tests/graphics/SilkFX/src/com/android/test/silkfx/materials/BlurBehindContainer.kt
new file mode 100644
index 0000000..ce6348e
--- /dev/null
+++ b/tests/graphics/SilkFX/src/com/android/test/silkfx/materials/BlurBehindContainer.kt
@@ -0,0 +1,30 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.test.silkfx.materials
+
+import android.content.Context
+import android.graphics.RenderEffect
+import android.graphics.Shader
+import android.util.AttributeSet
+import android.widget.FrameLayout
+
+class BlurBehindContainer(context: Context, attributeSet: AttributeSet) : FrameLayout(context, attributeSet) {
+    override fun onFinishInflate() {
+        super.onFinishInflate()
+        setBackdropRenderEffect(
+            RenderEffect.createBlurEffect(16.0f, 16.0f, Shader.TileMode.CLAMP))
+    }
+}
\ No newline at end of file
diff --git a/wifi/java/src/android/net/wifi/WifiMigration.java b/wifi/java/src/android/net/wifi/WifiMigration.java
index f1850dd..28e9c45 100644
--- a/wifi/java/src/android/net/wifi/WifiMigration.java
+++ b/wifi/java/src/android/net/wifi/WifiMigration.java
@@ -38,6 +38,8 @@
 import android.util.Log;
 import android.util.SparseArray;
 
+import com.android.internal.os.BackgroundThread;
+
 import java.io.File;
 import java.io.FileNotFoundException;
 import java.io.InputStream;
@@ -48,6 +50,8 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.Executor;
+import java.util.function.IntConsumer;
 
 /**
  * Class used to provide one time hooks for existing OEM devices to migrate their config store
@@ -605,13 +609,35 @@
     /**
      * Migrate any certificates in Legacy Keystore to the newer WifiBlobstore database.
      *
-     * If there are no certificates to migrate, this method will return immediately.
+     * Operation will be handled on the BackgroundThread, and the result will be posted
+     * to the provided Executor.
+     *
+     * @param executor The executor on which callback will be invoked
+     * @param resultsCallback Callback to receive the status code
      *
      * @hide
      */
     @FlaggedApi(Flags.FLAG_LEGACY_KEYSTORE_TO_WIFI_BLOBSTORE_MIGRATION_READ_ONLY)
     @SystemApi(client = SystemApi.Client.MODULE_LIBRARIES)
-    public static @KeystoreMigrationStatus int migrateLegacyKeystoreToWifiBlobstore() {
+    public static void migrateLegacyKeystoreToWifiBlobstore(
+            @NonNull Executor executor, @NonNull IntConsumer resultsCallback) {
+        Objects.requireNonNull(executor, "executor cannot be null");
+        Objects.requireNonNull(resultsCallback, "resultsCallback cannot be null");
+        BackgroundThread.getHandler().post(() -> {
+            int status = migrateLegacyKeystoreToWifiBlobstoreInternal();
+            executor.execute(() -> {
+                resultsCallback.accept(status);
+            });
+        });
+    }
+
+    /**
+     * Synchronously perform the Keystore migration described in
+     * {@link #migrateLegacyKeystoreToWifiBlobstore(Executor, IntConsumer)}
+     *
+     * @hide
+     */
+    public static @KeystoreMigrationStatus int migrateLegacyKeystoreToWifiBlobstoreInternal() {
         if (!WifiBlobStore.supplicantCanAccessBlobstore()) {
             // Supplicant cannot access WifiBlobstore, so keep the certs in Legacy Keystore
             Log.i(TAG, "Avoiding migration since supplicant cannot access WifiBlobstore");
diff --git a/wifi/tests/src/android/net/wifi/WifiMigrationTest.java b/wifi/tests/src/android/net/wifi/WifiMigrationTest.java
index 0aa299f..ce5b60a 100644
--- a/wifi/tests/src/android/net/wifi/WifiMigrationTest.java
+++ b/wifi/tests/src/android/net/wifi/WifiMigrationTest.java
@@ -81,7 +81,7 @@
     public void testKeystoreMigrationAvoidedOnLegacyVendorPartition() {
         when(WifiBlobStore.supplicantCanAccessBlobstore()).thenReturn(false);
         assertEquals(WifiMigration.KEYSTORE_MIGRATION_SUCCESS_MIGRATION_NOT_NEEDED,
-                WifiMigration.migrateLegacyKeystoreToWifiBlobstore());
+                WifiMigration.migrateLegacyKeystoreToWifiBlobstoreInternal());
         verifyNoMoreInteractions(mLegacyKeystore, mWifiBlobStore);
     }
 
@@ -93,7 +93,7 @@
     public void testKeystoreMigrationNoLegacyAliases() throws Exception {
         when(mLegacyKeystore.list(anyString(), anyInt())).thenReturn(new String[0]);
         assertEquals(WifiMigration.KEYSTORE_MIGRATION_SUCCESS_MIGRATION_NOT_NEEDED,
-                WifiMigration.migrateLegacyKeystoreToWifiBlobstore());
+                WifiMigration.migrateLegacyKeystoreToWifiBlobstoreInternal());
         verify(mLegacyKeystore).list(anyString(), anyInt());
         verifyNoMoreInteractions(mLegacyKeystore, mWifiBlobStore);
     }
@@ -110,7 +110,7 @@
         when(mWifiBlobStore.list(anyString())).thenReturn(blobstoreAliases);
 
         assertEquals(WifiMigration.KEYSTORE_MIGRATION_SUCCESS_MIGRATION_COMPLETE,
-                WifiMigration.migrateLegacyKeystoreToWifiBlobstore());
+                WifiMigration.migrateLegacyKeystoreToWifiBlobstoreInternal());
         verify(mWifiBlobStore, times(legacyAliases.length)).put(anyString(), any(byte[].class));
     }
 
@@ -129,7 +129,7 @@
 
         // Expect that only the unique legacy alias is migrated to the blobstore
         assertEquals(WifiMigration.KEYSTORE_MIGRATION_SUCCESS_MIGRATION_COMPLETE,
-                WifiMigration.migrateLegacyKeystoreToWifiBlobstore());
+                WifiMigration.migrateLegacyKeystoreToWifiBlobstoreInternal());
         verify(mWifiBlobStore).list(anyString());
         verify(mWifiBlobStore).put(eq(uniqueLegacyAlias), any(byte[].class));
         verifyNoMoreInteractions(mWifiBlobStore);
@@ -146,7 +146,7 @@
         when(mLegacyKeystore.list(anyString(), anyInt())).thenThrow(
                 new ServiceSpecificException(ILegacyKeystore.ERROR_SYSTEM_ERROR));
         assertEquals(WifiMigration.KEYSTORE_MIGRATION_SUCCESS_MIGRATION_NOT_NEEDED,
-                WifiMigration.migrateLegacyKeystoreToWifiBlobstore());
+                WifiMigration.migrateLegacyKeystoreToWifiBlobstoreInternal());
     }
 
     /**
@@ -157,6 +157,6 @@
     public void testKeystoreMigrationFailsIfExceptionEncountered() throws Exception {
         when(mLegacyKeystore.list(anyString(), anyInt())).thenThrow(new RemoteException());
         assertEquals(WifiMigration.KEYSTORE_MIGRATION_FAILURE_ENCOUNTERED_EXCEPTION,
-                WifiMigration.migrateLegacyKeystoreToWifiBlobstore());
+                WifiMigration.migrateLegacyKeystoreToWifiBlobstoreInternal());
     }
 }