diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 80578e4..ac67909 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -1765,15 +1765,12 @@
   public final class InputManager {
     method public void addUniqueIdAssociation(@NonNull String, @NonNull String);
     method @RequiresPermission(android.Manifest.permission.REMAP_MODIFIER_KEYS) public void clearAllModifierKeyRemappings();
-    method @Nullable public String getCurrentKeyboardLayoutForInputDevice(@NonNull android.hardware.input.InputDeviceIdentifier);
-    method @NonNull public java.util.List<java.lang.String> getKeyboardLayoutDescriptorsForInputDevice(@NonNull android.view.InputDevice);
+    method @NonNull public java.util.List<java.lang.String> getKeyboardLayoutDescriptors();
     method @NonNull public String getKeyboardLayoutTypeForLayoutDescriptor(@NonNull String);
     method @NonNull @RequiresPermission(android.Manifest.permission.REMAP_MODIFIER_KEYS) public java.util.Map<java.lang.Integer,java.lang.Integer> getModifierKeyRemapping();
     method public int getMousePointerSpeed();
     method @RequiresPermission(android.Manifest.permission.REMAP_MODIFIER_KEYS) public void remapModifierKey(int, int);
-    method @RequiresPermission(android.Manifest.permission.SET_KEYBOARD_LAYOUT) public void removeKeyboardLayoutForInputDevice(@NonNull android.hardware.input.InputDeviceIdentifier, @NonNull String);
     method public void removeUniqueIdAssociation(@NonNull String);
-    method @RequiresPermission(android.Manifest.permission.SET_KEYBOARD_LAYOUT) public void setCurrentKeyboardLayoutForInputDevice(@NonNull android.hardware.input.InputDeviceIdentifier, @NonNull String);
     field public static final long BLOCK_UNTRUSTED_TOUCHES = 158002302L; // 0x96aec7eL
   }
 
diff --git a/core/java/android/hardware/input/IInputManager.aidl b/core/java/android/hardware/input/IInputManager.aidl
index 2816f77..1c37aa2 100644
--- a/core/java/android/hardware/input/IInputManager.aidl
+++ b/core/java/android/hardware/input/IInputManager.aidl
@@ -94,33 +94,8 @@
     // Keyboard layouts configuration.
     KeyboardLayout[] getKeyboardLayouts();
 
-    KeyboardLayout[] getKeyboardLayoutsForInputDevice(in InputDeviceIdentifier identifier);
-
     KeyboardLayout getKeyboardLayout(String keyboardLayoutDescriptor);
 
-    String getCurrentKeyboardLayoutForInputDevice(in InputDeviceIdentifier identifier);
-
-    @EnforcePermission("SET_KEYBOARD_LAYOUT")
-    @JavaPassthrough(annotation="@android.annotation.RequiresPermission(value = "
-            + "android.Manifest.permission.SET_KEYBOARD_LAYOUT)")
-    void setCurrentKeyboardLayoutForInputDevice(in InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor);
-
-    String[] getEnabledKeyboardLayoutsForInputDevice(in InputDeviceIdentifier identifier);
-
-    @EnforcePermission("SET_KEYBOARD_LAYOUT")
-    @JavaPassthrough(annotation="@android.annotation.RequiresPermission(value = "
-            + "android.Manifest.permission.SET_KEYBOARD_LAYOUT)")
-    void addKeyboardLayoutForInputDevice(in InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor);
-
-    @EnforcePermission("SET_KEYBOARD_LAYOUT")
-    @JavaPassthrough(annotation="@android.annotation.RequiresPermission(value = "
-            + "android.Manifest.permission.SET_KEYBOARD_LAYOUT)")
-    void removeKeyboardLayoutForInputDevice(in InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor);
-
-    // New Keyboard layout config APIs
     KeyboardLayoutSelectionResult getKeyboardLayoutForInputDevice(
             in InputDeviceIdentifier identifier, int userId, in InputMethodInfo imeInfo,
             in InputMethodSubtype imeSubtype);
diff --git a/core/java/android/hardware/input/InputManager.java b/core/java/android/hardware/input/InputManager.java
index a1242fb..f949158 100644
--- a/core/java/android/hardware/input/InputManager.java
+++ b/core/java/android/hardware/input/InputManager.java
@@ -473,22 +473,21 @@
     }
 
     /**
-     * Returns the descriptors of all supported keyboard layouts appropriate for the specified
-     * input device.
+     * Returns the descriptors of all supported keyboard layouts.
      * <p>
      * The input manager consults the built-in keyboard layouts as well as all keyboard layouts
      * advertised by applications using a {@link #ACTION_QUERY_KEYBOARD_LAYOUTS} broadcast receiver.
      * </p>
      *
-     * @param device The input device to query.
      * @return The ids of all keyboard layouts which are supported by the specified input device.
      *
      * @hide
      */
     @TestApi
     @NonNull
-    public List<String> getKeyboardLayoutDescriptorsForInputDevice(@NonNull InputDevice device) {
-        KeyboardLayout[] layouts = getKeyboardLayoutsForInputDevice(device.getIdentifier());
+    @SuppressLint("UnflaggedApi")
+    public List<String> getKeyboardLayoutDescriptors() {
+        KeyboardLayout[] layouts = getKeyboardLayouts();
         List<String> res = new ArrayList<>();
         for (KeyboardLayout kl : layouts) {
             res.add(kl.getDescriptor());
@@ -511,33 +510,18 @@
     @TestApi
     @NonNull
     public String getKeyboardLayoutTypeForLayoutDescriptor(@NonNull String layoutDescriptor) {
-        KeyboardLayout[] layouts = getKeyboardLayouts();
-        for (KeyboardLayout kl : layouts) {
-            if (layoutDescriptor.equals(kl.getDescriptor())) {
-                return kl.getLayoutType();
-            }
-        }
-        return "";
+        KeyboardLayout layout = getKeyboardLayout(layoutDescriptor);
+        return layout == null ? "" : layout.getLayoutType();
     }
 
     /**
-     * Gets information about all supported keyboard layouts appropriate
-     * for a specific input device.
-     * <p>
-     * The input manager consults the built-in keyboard layouts as well
-     * as all keyboard layouts advertised by applications using a
-     * {@link #ACTION_QUERY_KEYBOARD_LAYOUTS} broadcast receiver.
-     * </p>
-     *
-     * @return A list of all supported keyboard layouts for a specific
-     * input device.
-     *
+     * TODO(b/330517633): Cleanup the unsupported API
      * @hide
      */
     @NonNull
     public KeyboardLayout[] getKeyboardLayoutsForInputDevice(
             @NonNull InputDeviceIdentifier identifier) {
-        return mGlobal.getKeyboardLayoutsForInputDevice(identifier);
+        return new KeyboardLayout[0];
     }
 
     /**
@@ -549,6 +533,7 @@
      *
      * @hide
      */
+    @Nullable
     public KeyboardLayout getKeyboardLayout(String keyboardLayoutDescriptor) {
         if (keyboardLayoutDescriptor == null) {
             throw new IllegalArgumentException("keyboardLayoutDescriptor must not be null");
@@ -562,121 +547,45 @@
     }
 
     /**
-     * Gets the current keyboard layout descriptor for the specified input device.
-     *
-     * @param identifier Identifier for the input device
-     * @return The keyboard layout descriptor, or null if no keyboard layout has been set.
-     *
+     * TODO(b/330517633): Cleanup the unsupported API
      * @hide
      */
-    @TestApi
     @Nullable
     public String getCurrentKeyboardLayoutForInputDevice(
             @NonNull InputDeviceIdentifier identifier) {
-        try {
-            return mIm.getCurrentKeyboardLayoutForInputDevice(identifier);
-        } catch (RemoteException ex) {
-            throw ex.rethrowFromSystemServer();
-        }
+        return null;
     }
 
     /**
-     * Sets the current keyboard layout descriptor for the specified input device.
-     * <p>
-     * This method may have the side-effect of causing the input device in question to be
-     * reconfigured.
-     * </p>
-     *
-     * @param identifier The identifier for the input device.
-     * @param keyboardLayoutDescriptor The keyboard layout descriptor to use, must not be null.
-     *
+     * TODO(b/330517633): Cleanup the unsupported API
      * @hide
      */
-    @TestApi
-    @RequiresPermission(Manifest.permission.SET_KEYBOARD_LAYOUT)
     public void setCurrentKeyboardLayoutForInputDevice(@NonNull InputDeviceIdentifier identifier,
-            @NonNull String keyboardLayoutDescriptor) {
-        mGlobal.setCurrentKeyboardLayoutForInputDevice(identifier,
-                keyboardLayoutDescriptor);
-    }
+            @NonNull String keyboardLayoutDescriptor) {}
 
     /**
-     * Gets all keyboard layout descriptors that are enabled for the specified input device.
-     *
-     * @param identifier The identifier for the input device.
-     * @return The keyboard layout descriptors.
-     *
+     * TODO(b/330517633): Cleanup the unsupported API
      * @hide
      */
     public String[] getEnabledKeyboardLayoutsForInputDevice(InputDeviceIdentifier identifier) {
-        if (identifier == null) {
-            throw new IllegalArgumentException("inputDeviceDescriptor must not be null");
-        }
-
-        try {
-            return mIm.getEnabledKeyboardLayoutsForInputDevice(identifier);
-        } catch (RemoteException ex) {
-            throw ex.rethrowFromSystemServer();
-        }
+        return new String[0];
     }
 
     /**
-     * Adds the keyboard layout descriptor for the specified input device.
-     * <p>
-     * This method may have the side-effect of causing the input device in question to be
-     * reconfigured.
-     * </p>
-     *
-     * @param identifier The identifier for the input device.
-     * @param keyboardLayoutDescriptor The descriptor of the keyboard layout to add.
-     *
+     * TODO(b/330517633): Cleanup the unsupported API
      * @hide
      */
-    @RequiresPermission(Manifest.permission.SET_KEYBOARD_LAYOUT)
     public void addKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
             String keyboardLayoutDescriptor) {
-        if (identifier == null) {
-            throw new IllegalArgumentException("inputDeviceDescriptor must not be null");
-        }
-        if (keyboardLayoutDescriptor == null) {
-            throw new IllegalArgumentException("keyboardLayoutDescriptor must not be null");
-        }
-
-        try {
-            mIm.addKeyboardLayoutForInputDevice(identifier, keyboardLayoutDescriptor);
-        } catch (RemoteException ex) {
-            throw ex.rethrowFromSystemServer();
-        }
     }
 
     /**
-     * Removes the keyboard layout descriptor for the specified input device.
-     * <p>
-     * This method may have the side-effect of causing the input device in question to be
-     * reconfigured.
-     * </p>
-     *
-     * @param identifier The identifier for the input device.
-     * @param keyboardLayoutDescriptor The descriptor of the keyboard layout to remove.
-     *
+     * TODO(b/330517633): Cleanup the unsupported API
      * @hide
      */
-    @TestApi
     @RequiresPermission(Manifest.permission.SET_KEYBOARD_LAYOUT)
     public void removeKeyboardLayoutForInputDevice(@NonNull InputDeviceIdentifier identifier,
             @NonNull String keyboardLayoutDescriptor) {
-        if (identifier == null) {
-            throw new IllegalArgumentException("inputDeviceDescriptor must not be null");
-        }
-        if (keyboardLayoutDescriptor == null) {
-            throw new IllegalArgumentException("keyboardLayoutDescriptor must not be null");
-        }
-
-        try {
-            mIm.removeKeyboardLayoutForInputDevice(identifier, keyboardLayoutDescriptor);
-        } catch (RemoteException ex) {
-            throw ex.rethrowFromSystemServer();
-        }
     }
 
     /**
diff --git a/core/java/android/hardware/input/InputManagerGlobal.java b/core/java/android/hardware/input/InputManagerGlobal.java
index 7c104a0..7b29666 100644
--- a/core/java/android/hardware/input/InputManagerGlobal.java
+++ b/core/java/android/hardware/input/InputManagerGlobal.java
@@ -1068,36 +1068,21 @@
     }
 
     /**
-     * @see InputManager#getKeyboardLayoutsForInputDevice(InputDeviceIdentifier)
+     * TODO(b/330517633): Cleanup the unsupported API
      */
     @NonNull
     public KeyboardLayout[] getKeyboardLayoutsForInputDevice(
             @NonNull InputDeviceIdentifier identifier) {
-        try {
-            return mIm.getKeyboardLayoutsForInputDevice(identifier);
-        } catch (RemoteException ex) {
-            throw ex.rethrowFromSystemServer();
-        }
+        return new KeyboardLayout[0];
     }
 
     /**
-     * @see InputManager#setCurrentKeyboardLayoutForInputDevice
-     * (InputDeviceIdentifier, String)
+     * TODO(b/330517633): Cleanup the unsupported API
      */
-    @RequiresPermission(Manifest.permission.SET_KEYBOARD_LAYOUT)
     public void setCurrentKeyboardLayoutForInputDevice(
             @NonNull InputDeviceIdentifier identifier,
-            @NonNull String keyboardLayoutDescriptor) {
-        Objects.requireNonNull(identifier, "identifier must not be null");
-        Objects.requireNonNull(keyboardLayoutDescriptor,
-                "keyboardLayoutDescriptor must not be null");
-        try {
-            mIm.setCurrentKeyboardLayoutForInputDevice(identifier,
-                    keyboardLayoutDescriptor);
-        } catch (RemoteException ex) {
-            throw ex.rethrowFromSystemServer();
-        }
-    }
+            @NonNull String keyboardLayoutDescriptor) {}
+
 
     /**
      * @see InputDevice#getSensorManager()
diff --git a/core/java/android/service/wallpaper/WallpaperService.java b/core/java/android/service/wallpaper/WallpaperService.java
index 0fc51e7..582c90f 100644
--- a/core/java/android/service/wallpaper/WallpaperService.java
+++ b/core/java/android/service/wallpaper/WallpaperService.java
@@ -27,12 +27,15 @@
 import static android.view.View.SYSTEM_UI_FLAG_VISIBLE;
 import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
 
+import static com.android.window.flags.Flags.FLAG_OFFLOAD_COLOR_EXTRACTION;
 import static com.android.window.flags.Flags.noConsecutiveVisibilityEvents;
+import static com.android.window.flags.Flags.offloadColorExtraction;
 
 import android.animation.AnimationHandler;
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
+import android.annotation.FlaggedApi;
 import android.annotation.FloatRange;
 import android.annotation.MainThread;
 import android.annotation.NonNull;
@@ -832,6 +835,15 @@
         }
 
         /**
+         * Called when the dim amount of the wallpaper changed. This can be used to recompute the
+         * wallpaper colors based on the new dim, and call {@link #notifyColorsChanged()}.
+         * @hide
+         */
+        @FlaggedApi(FLAG_OFFLOAD_COLOR_EXTRACTION)
+        public void onDimAmountChanged(float dimAmount) {
+        }
+
+        /**
          * Called when an application has changed the desired virtual size of
          * the wallpaper.
          */
@@ -1043,6 +1055,10 @@
             }
 
             mPreviousWallpaperDimAmount = mWallpaperDimAmount;
+
+            // after the dim changes, allow colors to be immediately recomputed
+            mLastColorInvalidation = 0;
+            if (offloadColorExtraction()) onDimAmountChanged(mWallpaperDimAmount);
         }
 
         /**
diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index b300022..6caf4d6 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -196,11 +196,11 @@
         if (!super.setControl(control, showTypes, hideTypes)) {
             return false;
         }
-        final boolean hasLeash = control != null && control.getLeash() != null;
-        if (!hasLeash && !mIsRequestedVisibleAwaitingLeash) {
+        if (control == null && !mIsRequestedVisibleAwaitingLeash) {
             mController.setRequestedVisibleTypes(0 /* visibleTypes */, getType());
             removeSurface();
         }
+        final boolean hasLeash = control != null && control.getLeash() != null;
         if (hasLeash) {
             mIsRequestedVisibleAwaitingLeash = false;
         }
diff --git a/core/java/android/window/flags/lse_desktop_experience.aconfig b/core/java/android/window/flags/lse_desktop_experience.aconfig
index efe31ff..616004b 100644
--- a/core/java/android/window/flags/lse_desktop_experience.aconfig
+++ b/core/java/android/window/flags/lse_desktop_experience.aconfig
@@ -36,3 +36,10 @@
     description: "Enables a limit on the number of Tasks shown in Desktop Mode"
     bug: "332502912"
 }
+
+flag {
+    name: "enable_windowing_edge_drag_resize"
+    namespace: "lse_desktop_experience"
+    description: "Enables edge drag resizing for all input sources"
+    bug: "323383067"
+}
diff --git a/core/java/android/window/flags/wallpaper_manager.aconfig b/core/java/android/window/flags/wallpaper_manager.aconfig
index aa92af2..150b04e 100644
--- a/core/java/android/window/flags/wallpaper_manager.aconfig
+++ b/core/java/android/window/flags/wallpaper_manager.aconfig
@@ -21,4 +21,14 @@
   namespace: "systemui"
   description: "Prevent the system from sending consecutive onVisibilityChanged(false) events."
   bug: "285631818"
+}
+
+flag {
+  name: "offload_color_extraction"
+  namespace: "systemui"
+  description: "Let ImageWallpaper take care of its wallpaper color extraction, instead of system_server"
+  bug: "328791519"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
 }
\ No newline at end of file
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index 4d2abb0..87c47da 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -141,4 +141,15 @@
   description: "Configuration decoupled from insets"
   bug: "151861875"
   is_fixed_read_only: true
+}
+
+flag {
+  name: "keyguard_appear_transition"
+  namespace: "windowing_frontend"
+  description: "Add transition when keyguard appears"
+  bug: "327970608"
+  is_fixed_read_only: true
+  metadata {
+      purpose: PURPOSE_BUGFIX
+  }
 }
\ No newline at end of file
diff --git a/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java b/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java
index 6cc5485..8072d69 100644
--- a/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java
+++ b/core/tests/bugreports/src/com/android/os/bugreports/tests/BugreportManagerTest.java
@@ -16,6 +16,8 @@
 
 package com.android.os.bugreports.tests;
 
+import static android.content.Context.RECEIVER_EXPORTED;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertNotNull;
@@ -358,7 +360,10 @@
         // shell UID rather than our own.
         BugreportBroadcastReceiver br = new BugreportBroadcastReceiver();
         InstrumentationRegistry.getContext()
-                .registerReceiver(br, new IntentFilter(INTENT_BUGREPORT_FINISHED));
+                .registerReceiver(
+                        br,
+                        new IntentFilter(INTENT_BUGREPORT_FINISHED),
+                        RECEIVER_EXPORTED);
         UiDevice.getInstance(InstrumentationRegistry.getInstrumentation())
                 .executeShellCommand("am bug-report");
 
diff --git a/data/etc/core.protolog.pb b/data/etc/core.protolog.pb
index 826adc39..97147a0 100644
--- a/data/etc/core.protolog.pb
+++ b/data/etc/core.protolog.pb
Binary files differ
diff --git a/data/etc/services.core.protolog.json b/data/etc/services.core.protolog.json
index d410d5f..6cf12de 100644
--- a/data/etc/services.core.protolog.json
+++ b/data/etc/services.core.protolog.json
@@ -709,18 +709,6 @@
       "group": "WM_DEBUG_CONFIGURATION",
       "at": "com\/android\/server\/wm\/ActivityTaskManagerService.java"
     },
-    "2959074735946674755": {
-      "message": "Trying to update display configuration for system\/invalid process.",
-      "level": "WARN",
-      "group": "WM_DEBUG_CONFIGURATION",
-      "at": "com\/android\/server\/wm\/ActivityTaskManagerService.java"
-    },
-    "5668810920995272206": {
-      "message": "Trying to update display configuration for invalid process, pid=%d",
-      "level": "WARN",
-      "group": "WM_DEBUG_CONFIGURATION",
-      "at": "com\/android\/server\/wm\/ActivityTaskManagerService.java"
-    },
     "-1123414663662718691": {
       "message": "setVr2dDisplayId called for: %d",
       "level": "DEBUG",
@@ -3469,6 +3457,12 @@
       "group": "WM_DEBUG_WALLPAPER",
       "at": "com\/android\/server\/wm\/WallpaperController.java"
     },
+    "257349083882992098": {
+      "message": "updateWallpaperTokens requestedVisibility=%b on keyguardLocked=%b",
+      "level": "VERBOSE",
+      "group": "WM_DEBUG_WALLPAPER",
+      "at": "com\/android\/server\/wm\/WallpaperController.java"
+    },
     "7408402065665963407": {
       "message": "Wallpaper at display %d - visibility: %b, keyguardLocked: %b",
       "level": "VERBOSE",
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index 313d0d2..d295877 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -455,8 +455,7 @@
                         ProtoLog.d(WM_SHELL_BUBBLES,
                                 "onActivityRestartAttempt - taskId=%d selecting matching bubble=%s",
                                 task.taskId, b.getKey());
-                        mBubbleData.setSelectedBubble(b);
-                        mBubbleData.setExpanded(true);
+                        mBubbleData.setSelectedBubbleAndExpandStack(b);
                         return;
                     }
                 }
@@ -593,13 +592,6 @@
         }
     }
 
-    private void openBubbleOverflow() {
-        ensureBubbleViewsAndWindowCreated();
-        mBubbleData.setShowingOverflow(true);
-        mBubbleData.setSelectedBubble(mBubbleData.getOverflow());
-        mBubbleData.setExpanded(true);
-    }
-
     /**
      * Called when the status bar has become visible or invisible (either permanently or
      * temporarily).
@@ -1247,8 +1239,7 @@
         }
         if (mBubbleData.hasBubbleInStackWithKey(b.getKey())) {
             // already in the stack
-            mBubbleData.setSelectedBubble(b);
-            mBubbleData.setExpanded(true);
+            mBubbleData.setSelectedBubbleAndExpandStack(b);
         } else if (mBubbleData.hasOverflowBubbleWithKey(b.getKey())) {
             // promote it out of the overflow
             promoteBubbleFromOverflow(b);
@@ -1273,8 +1264,7 @@
             String key = entry.getKey();
             Bubble bubble = mBubbleData.getBubbleInStackWithKey(key);
             if (bubble != null) {
-                mBubbleData.setSelectedBubble(bubble);
-                mBubbleData.setExpanded(true);
+                mBubbleData.setSelectedBubbleAndExpandStack(bubble);
             } else {
                 bubble = mBubbleData.getOverflowBubbleWithKey(key);
                 if (bubble != null) {
@@ -1367,8 +1357,7 @@
             } else {
                 // App bubble is not selected, select it & expand
                 Log.i(TAG, "  showOrHideAppBubble, expand and select existing app bubble");
-                mBubbleData.setSelectedBubble(existingAppBubble);
-                mBubbleData.setExpanded(true);
+                mBubbleData.setSelectedBubbleAndExpandStack(existingAppBubble);
             }
         } else {
             // Check if it exists in the overflow
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
index 61f0ed2..ae3d0c5 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
@@ -365,6 +365,19 @@
         mSelectedBubble = bubble;
     }
 
+    /**
+     * Sets the selected bubble and expands it.
+     *
+     * <p>This dispatches a single state update for both changes and should be used instead of
+     * calling {@link #setSelectedBubble(BubbleViewProvider)} followed by
+     * {@link #setExpanded(boolean)} immediately after, which will generate 2 separate updates.
+     */
+    public void setSelectedBubbleAndExpandStack(BubbleViewProvider bubble) {
+        setSelectedBubbleInternal(bubble);
+        setExpandedInternal(true);
+        dispatchPendingChanges();
+    }
+
     public void setSelectedBubble(BubbleViewProvider bubble) {
         setSelectedBubbleInternal(bubble);
         dispatchPendingChanges();
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
index 48e396a..6be411d 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
@@ -1222,6 +1222,19 @@
         assertThat(update.bubbleBarLocation).isEqualTo(BubbleBarLocation.LEFT);
     }
 
+    @Test
+    public void setSelectedBubbleAndExpandStack() {
+        sendUpdatedEntryAtTime(mEntryA1, 1000);
+        sendUpdatedEntryAtTime(mEntryA2, 2000);
+        mBubbleData.setListener(mListener);
+
+        mBubbleData.setSelectedBubbleAndExpandStack(mBubbleA1);
+
+        verifyUpdateReceived();
+        assertSelectionChangedTo(mBubbleA1);
+        assertExpandedChangedTo(true);
+    }
+
     private void verifyUpdateReceived() {
         verify(mListener).applyUpdate(mUpdateCaptor.capture());
         reset(mListener);
diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index 76bbca6..5aa006b 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -1309,18 +1309,24 @@
             return;
         }
 
-        RoutingController newController;
-        if (sessionInfo.isSystemSession()) {
-            newController = getSystemController();
-            newController.setRoutingSessionInfo(sessionInfo);
+        RoutingController newController = addRoutingController(sessionInfo);
+        notifyTransfer(oldController, newController);
+    }
+
+    @NonNull
+    private RoutingController addRoutingController(@NonNull RoutingSessionInfo session) {
+        RoutingController controller;
+        if (session.isSystemSession()) {
+            // mSystemController is never released, so we only need to update its status.
+            mSystemController.setRoutingSessionInfo(session);
+            controller = mSystemController;
         } else {
-            newController = new RoutingController(sessionInfo);
+            controller = new RoutingController(session);
             synchronized (mLock) {
-                mNonSystemRoutingControllers.put(newController.getId(), newController);
+                mNonSystemRoutingControllers.put(controller.getId(), controller);
             }
         }
-
-        notifyTransfer(oldController, newController);
+        return controller;
     }
 
     void updateControllerOnHandler(RoutingSessionInfo sessionInfo) {
@@ -1329,40 +1335,12 @@
             return;
         }
 
-        if (sessionInfo.isSystemSession()) {
-            // The session info is sent from SystemMediaRoute2Provider.
-            RoutingController systemController = getSystemController();
-            systemController.setRoutingSessionInfo(sessionInfo);
-            notifyControllerUpdated(systemController);
-            return;
+        RoutingController controller =
+                getMatchingController(sessionInfo, /* logPrefix */ "updateControllerOnHandler");
+        if (controller != null) {
+            controller.setRoutingSessionInfo(sessionInfo);
+            notifyControllerUpdated(controller);
         }
-
-        RoutingController matchingController;
-        synchronized (mLock) {
-            matchingController = mNonSystemRoutingControllers.get(sessionInfo.getId());
-        }
-
-        if (matchingController == null) {
-            Log.w(
-                    TAG,
-                    "updateControllerOnHandler: Matching controller not found. uniqueSessionId="
-                            + sessionInfo.getId());
-            return;
-        }
-
-        RoutingSessionInfo oldInfo = matchingController.getRoutingSessionInfo();
-        if (!TextUtils.equals(oldInfo.getProviderId(), sessionInfo.getProviderId())) {
-            Log.w(
-                    TAG,
-                    "updateControllerOnHandler: Provider IDs are not matched. old="
-                            + oldInfo.getProviderId()
-                            + ", new="
-                            + sessionInfo.getProviderId());
-            return;
-        }
-
-        matchingController.setRoutingSessionInfo(sessionInfo);
-        notifyControllerUpdated(matchingController);
     }
 
     void releaseControllerOnHandler(RoutingSessionInfo sessionInfo) {
@@ -1371,34 +1349,47 @@
             return;
         }
 
-        RoutingController matchingController;
-        synchronized (mLock) {
-            matchingController = mNonSystemRoutingControllers.get(sessionInfo.getId());
-        }
+        RoutingController matchingController =
+                getMatchingController(sessionInfo, /* logPrefix */ "releaseControllerOnHandler");
 
-        if (matchingController == null) {
-            if (DEBUG) {
-                Log.d(
-                        TAG,
-                        "releaseControllerOnHandler: Matching controller not found. "
-                                + "uniqueSessionId="
-                                + sessionInfo.getId());
+        if (matchingController != null) {
+            matchingController.releaseInternal(/* shouldReleaseSession= */ false);
+        }
+    }
+
+    @Nullable
+    private RoutingController getMatchingController(
+            RoutingSessionInfo sessionInfo, String logPrefix) {
+        if (sessionInfo.isSystemSession()) {
+            return getSystemController();
+        } else {
+            RoutingController controller;
+            synchronized (mLock) {
+                controller = mNonSystemRoutingControllers.get(sessionInfo.getId());
             }
-            return;
-        }
 
-        RoutingSessionInfo oldInfo = matchingController.getRoutingSessionInfo();
-        if (!TextUtils.equals(oldInfo.getProviderId(), sessionInfo.getProviderId())) {
-            Log.w(
-                    TAG,
-                    "releaseControllerOnHandler: Provider IDs are not matched. old="
-                            + oldInfo.getProviderId()
-                            + ", new="
-                            + sessionInfo.getProviderId());
-            return;
-        }
+            if (controller == null) {
+                Log.w(
+                        TAG,
+                        logPrefix
+                                + ": Matching controller not found. uniqueSessionId="
+                                + sessionInfo.getId());
+                return null;
+            }
 
-        matchingController.releaseInternal(/* shouldReleaseSession= */ false);
+            RoutingSessionInfo oldInfo = controller.getRoutingSessionInfo();
+            if (!TextUtils.equals(oldInfo.getProviderId(), sessionInfo.getProviderId())) {
+                Log.w(
+                        TAG,
+                        logPrefix
+                                + ": Provider IDs are not matched. old="
+                                + oldInfo.getProviderId()
+                                + ", new="
+                                + sessionInfo.getProviderId());
+                return null;
+            }
+            return controller;
+        }
     }
 
     void onRequestCreateControllerByManagerOnHandler(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
index fcd7760..02a12e4 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
@@ -54,9 +54,9 @@
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.LowestZIndexScenePicker
 import com.android.compose.animation.scene.SceneScope
+import com.android.compose.animation.scene.TransitionState
 import com.android.compose.animation.scene.ValueKey
 import com.android.compose.animation.scene.animateElementFloatAsState
-import com.android.compose.animation.scene.animateSceneFloatAsState
 import com.android.compose.windowsizeclass.LocalWindowSizeClass
 import com.android.settingslib.Utils
 import com.android.systemui.battery.BatteryMeterView
@@ -87,10 +87,6 @@
         val ShadeCarrierGroup = ElementKey("ShadeCarrierGroup")
     }
 
-    object Keys {
-        val transitionProgress = ValueKey("ShadeHeaderTransitionProgress")
-    }
-
     object Values {
         val ClockScale = ValueKey("ShadeHeaderClockScale")
     }
@@ -119,19 +115,17 @@
         return
     }
 
-    val formatProgress =
-        animateSceneFloatAsState(0f, ShadeHeader.Keys.transitionProgress)
-            .unsafeCompositionState(initialValue = 0f)
-
     val cutoutWidth = LocalDisplayCutout.current.width()
     val cutoutLocation = LocalDisplayCutout.current.location
 
     val useExpandedFormat by
-        remember(formatProgress) {
+        remember(cutoutLocation) {
             derivedStateOf {
-                cutoutLocation != CutoutLocation.CENTER || formatProgress.value > 0.5f
+                cutoutLocation != CutoutLocation.CENTER ||
+                    shouldUseExpandedFormat(layoutState.transitionState)
             }
         }
+
     val isPrivacyChipVisible by viewModel.isPrivacyChipVisible.collectAsState()
 
     // This layout assumes it is globally positioned at (0, 0) and is the
@@ -207,7 +201,7 @@
 
         val screenWidth = constraints.maxWidth
         val cutoutWidthPx = cutoutWidth.roundToPx()
-        val height = ShadeHeader.Dimensions.CollapsedHeight.roundToPx()
+        val height = CollapsedHeight.roundToPx()
         val childConstraints = Constraints.fixed((screenWidth - cutoutWidthPx) / 2, height)
 
         val startMeasurable = measurables[0][0]
@@ -261,11 +255,10 @@
         return
     }
 
-    val formatProgress =
-        animateSceneFloatAsState(1f, ShadeHeader.Keys.transitionProgress)
-            .unsafeCompositionState(initialValue = 1f)
-    val useExpandedFormat by
-        remember(formatProgress) { derivedStateOf { formatProgress.value > 0.5f } }
+    val useExpandedFormat by remember {
+        derivedStateOf { shouldUseExpandedFormat(layoutState.transitionState) }
+    }
+
     val isPrivacyChipVisible by viewModel.isPrivacyChipVisible.collectAsState()
 
     Box(modifier = modifier) {
@@ -530,3 +523,15 @@
         modifier = modifier.element(ShadeHeader.Elements.PrivacyChip),
     )
 }
+
+private fun shouldUseExpandedFormat(state: TransitionState): Boolean {
+    return when (state) {
+        is TransitionState.Idle -> {
+            state.currentScene == Scenes.QuickSettings
+        }
+        is TransitionState.Transition -> {
+            (state.isTransitioning(Scenes.Shade, Scenes.QuickSettings) && state.progress >= 0.5) ||
+                (state.isTransitioning(Scenes.QuickSettings, Scenes.Shade) && state.progress < 0.5)
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
index 42bd4af..ce1aed0 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/preview/KeyguardPreviewRenderer.kt
@@ -621,7 +621,9 @@
     }
 
     private suspend fun updateClockAppearance(clock: ClockController) {
-        clockController.clock = clock
+        if (!MigrateClocksToBlueprint.isEnabled) {
+            clockController.clock = clock
+        }
         val colors = wallpaperColors
         if (clockRegistry.seedColor == null && colors != null) {
             // Seed color null means users do not override any color on the clock. The default
@@ -639,6 +641,11 @@
                 if (isWallpaperDark) lightClockColor else darkClockColor
             )
         }
+        // In clock preview, we should have a seed color for clock
+        // before setting clock to clockEventController to avoid updateColor with seedColor == null
+        if (MigrateClocksToBlueprint.isEnabled) {
+            clockController.clock = clock
+        }
     }
     private fun onClockChanged() {
         if (MigrateClocksToBlueprint.isEnabled) {
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
index d247122..f08bc17 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
@@ -25,8 +25,8 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.mediaprojection.appselector.data.ActivityTaskManagerLabelLoader
 import com.android.systemui.mediaprojection.appselector.data.ActivityTaskManagerThumbnailLoader
-import com.android.systemui.mediaprojection.appselector.data.AppIconLoader
-import com.android.systemui.mediaprojection.appselector.data.IconLoaderLibAppIconLoader
+import com.android.systemui.mediaprojection.appselector.data.BasicAppIconLoader
+import com.android.systemui.mediaprojection.appselector.data.BasicPackageManagerAppIconLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskLabelLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskListProvider
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskThumbnailLoader
@@ -102,7 +102,7 @@
 
     @Binds
     @MediaProjectionAppSelectorScope
-    fun bindAppIconLoader(impl: IconLoaderLibAppIconLoader): AppIconLoader
+    fun bindAppIconLoader(impl: BasicPackageManagerAppIconLoader): BasicAppIconLoader
 
     @Binds
     @IntoSet
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BadgedAppIconLoader.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BadgedAppIconLoader.kt
new file mode 100644
index 0000000..ca5b5f8
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BadgedAppIconLoader.kt
@@ -0,0 +1,68 @@
+/*
+ * 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.mediaprojection.appselector.data
+
+import android.content.ComponentName
+import android.content.Context
+import android.graphics.drawable.Drawable
+import android.os.UserHandle
+import com.android.launcher3.icons.BaseIconFactory
+import com.android.launcher3.icons.IconFactory
+import com.android.launcher3.util.UserIconInfo
+import com.android.systemui.dagger.qualifiers.Background
+import javax.inject.Inject
+import javax.inject.Provider
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+class BadgedAppIconLoader
+@Inject
+constructor(
+    private val basicAppIconLoader: BasicAppIconLoader,
+    @Background private val backgroundDispatcher: CoroutineDispatcher,
+    private val context: Context,
+    private val iconFactoryProvider: Provider<IconFactory>,
+) {
+
+    suspend fun loadIcon(
+        userId: Int,
+        userType: RecentTask.UserType,
+        componentName: ComponentName
+    ): Drawable? =
+        withContext(backgroundDispatcher) {
+            iconFactoryProvider.get().use<IconFactory, Drawable?> { iconFactory ->
+                val icon =
+                    basicAppIconLoader.loadIcon(userId, componentName) ?: return@withContext null
+                val userHandler = UserHandle.of(userId)
+                val iconType = getIconType(userType)
+                val options =
+                    BaseIconFactory.IconOptions().apply {
+                        setUser(UserIconInfo(userHandler, iconType))
+                    }
+                val badgedIcon = iconFactory.createBadgedIconBitmap(icon, options)
+                badgedIcon.newIcon(context)
+            }
+        }
+
+    private fun getIconType(userType: RecentTask.UserType): Int =
+        when (userType) {
+            RecentTask.UserType.CLONED -> UserIconInfo.TYPE_CLONED
+            RecentTask.UserType.WORK -> UserIconInfo.TYPE_WORK
+            RecentTask.UserType.PRIVATE -> UserIconInfo.TYPE_PRIVATE
+            RecentTask.UserType.STANDARD -> UserIconInfo.TYPE_MAIN
+        }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/AppIconLoader.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BasicAppIconLoader.kt
similarity index 61%
rename from packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/AppIconLoader.kt
rename to packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BasicAppIconLoader.kt
index b85d628..03f6f01 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/AppIconLoader.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/BasicAppIconLoader.kt
@@ -17,45 +17,29 @@
 package com.android.systemui.mediaprojection.appselector.data
 
 import android.content.ComponentName
-import android.content.Context
 import android.content.pm.PackageManager
 import android.graphics.drawable.Drawable
-import android.os.UserHandle
-import com.android.launcher3.icons.BaseIconFactory.IconOptions
-import com.android.launcher3.icons.IconFactory
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.shared.system.PackageManagerWrapper
 import javax.inject.Inject
-import javax.inject.Provider
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.withContext
 
-interface AppIconLoader {
+interface BasicAppIconLoader {
     suspend fun loadIcon(userId: Int, component: ComponentName): Drawable?
 }
 
-class IconLoaderLibAppIconLoader
+class BasicPackageManagerAppIconLoader
 @Inject
 constructor(
     @Background private val backgroundDispatcher: CoroutineDispatcher,
-    private val context: Context,
     // Use wrapper to access hidden API that allows to get ActivityInfo for any user id
     private val packageManagerWrapper: PackageManagerWrapper,
     private val packageManager: PackageManager,
-    private val iconFactoryProvider: Provider<IconFactory>
-) : AppIconLoader {
+) : BasicAppIconLoader {
 
     override suspend fun loadIcon(userId: Int, component: ComponentName): Drawable? =
         withContext(backgroundDispatcher) {
-            iconFactoryProvider.get().use<IconFactory, Drawable?> { iconFactory ->
-                val activityInfo =
-                    packageManagerWrapper.getActivityInfo(component, userId)
-                        ?: return@withContext null
-                val icon = activityInfo.loadIcon(packageManager) ?: return@withContext null
-                val userHandler = UserHandle.of(userId)
-                val options = IconOptions().apply { setUser(userHandler) }
-                val badgedIcon = iconFactory.createBadgedIconBitmap(icon, options)
-                badgedIcon.newIcon(context)
-            }
+            packageManagerWrapper.getActivityInfo(component, userId)?.loadIcon(packageManager)
         }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
index e9b4582..3e9b546 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTask.kt
@@ -28,4 +28,12 @@
     val baseIntentComponent: ComponentName?,
     @ColorInt val colorBackground: Int?,
     val isForegroundTask: Boolean,
-)
+    val userType: UserType,
+) {
+    enum class UserType {
+        STANDARD,
+        WORK,
+        PRIVATE,
+        CLONED
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
index 5dde14b..a6049c8 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
@@ -17,6 +17,8 @@
 package com.android.systemui.mediaprojection.appselector.data
 
 import android.app.ActivityManager.RECENT_IGNORE_UNAVAILABLE
+import android.content.pm.UserInfo
+import android.os.UserManager
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.kotlin.getOrNull
@@ -41,7 +43,8 @@
     @Background private val coroutineDispatcher: CoroutineDispatcher,
     @Background private val backgroundExecutor: Executor,
     private val recentTasks: Optional<RecentTasks>,
-    private val userTracker: UserTracker
+    private val userTracker: UserTracker,
+    private val userManager: UserManager,
 ) : RecentTaskListProvider {
 
     private val recents by lazy { recentTasks.getOrNull() }
@@ -65,7 +68,8 @@
                         it.topActivity,
                         it.baseIntent?.component,
                         it.taskDescription?.backgroundColor,
-                        isForegroundTask = it.taskId in foregroundTaskIds && it.isVisible
+                        isForegroundTask = it.taskId in foregroundTaskIds && it.isVisible,
+                        userType = userManager.getUserInfo(it.userId).toUserType(),
                     )
                 }
         }
@@ -81,4 +85,15 @@
                 continuation.resume(tasks)
             }
         }
+
+    private fun UserInfo.toUserType(): RecentTask.UserType =
+        if (isCloneProfile) {
+            RecentTask.UserType.CLONED
+        } else if (isManagedProfile) {
+            RecentTask.UserType.WORK
+        } else if (isPrivateProfile) {
+            RecentTask.UserType.PRIVATE
+        } else {
+            RecentTask.UserType.STANDARD
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt
index 3fe040a..3b84d2c 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/view/RecentTaskViewHolder.kt
@@ -21,12 +21,12 @@
 import android.view.ViewGroup
 import android.widget.ImageView
 import androidx.recyclerview.widget.RecyclerView.ViewHolder
-import com.android.systemui.res.R
 import com.android.systemui.mediaprojection.appselector.MediaProjectionAppSelector
-import com.android.systemui.mediaprojection.appselector.data.AppIconLoader
+import com.android.systemui.mediaprojection.appselector.data.BadgedAppIconLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTask
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskLabelLoader
 import com.android.systemui.mediaprojection.appselector.data.RecentTaskThumbnailLoader
+import com.android.systemui.res.R
 import com.android.systemui.statusbar.policy.ConfigurationController.ConfigurationListener
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
@@ -39,7 +39,7 @@
 @AssistedInject
 constructor(
     @Assisted private val root: ViewGroup,
-    private val iconLoader: AppIconLoader,
+    private val iconLoader: BadgedAppIconLoader,
     private val thumbnailLoader: RecentTaskThumbnailLoader,
     private val labelLoader: RecentTaskLabelLoader,
     private val taskViewSizeProvider: TaskPreviewSizeProvider,
@@ -63,7 +63,7 @@
             scope.launch {
                 task.baseIntentComponent?.let { component ->
                     launch {
-                        val icon = iconLoader.loadIcon(task.userId, component)
+                        val icon = iconLoader.loadIcon(task.userId, task.userType, component)
                         iconView.setImageDrawable(icon)
                     }
                     launch {
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt
index f69021f..0ccb19c 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotActionsProvider.kt
@@ -29,7 +29,7 @@
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_PREVIEW_TAPPED
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_SHARE_TAPPED
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_SMART_ACTION_TAPPED
-import com.android.systemui.screenshot.ui.viewmodel.ActionButtonViewModel
+import com.android.systemui.screenshot.ui.viewmodel.ActionButtonAppearance
 import com.android.systemui.screenshot.ui.viewmodel.ScreenshotViewModel
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
@@ -81,65 +81,67 @@
             }
         }
         viewModel.addAction(
-            ActionButtonViewModel(
+            ActionButtonAppearance(
                 AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_edit),
                 context.resources.getString(R.string.screenshot_edit_label),
                 context.resources.getString(R.string.screenshot_edit_description),
-            ) {
-                debugLog(LogConfig.DEBUG_ACTIONS) { "Edit tapped" }
-                uiEventLogger.log(SCREENSHOT_EDIT_TAPPED, 0, request.packageNameString)
-                onDeferrableActionTapped { result ->
-                    actionExecutor.startSharedTransition(
-                        createEdit(result.uri, context),
-                        result.user,
-                        true
-                    )
-                }
+            )
+        ) {
+            debugLog(LogConfig.DEBUG_ACTIONS) { "Edit tapped" }
+            uiEventLogger.log(SCREENSHOT_EDIT_TAPPED, 0, request.packageNameString)
+            onDeferrableActionTapped { result ->
+                actionExecutor.startSharedTransition(
+                    createEdit(result.uri, context),
+                    result.user,
+                    true
+                )
             }
-        )
+        }
+
         viewModel.addAction(
-            ActionButtonViewModel(
+            ActionButtonAppearance(
                 AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_share),
                 context.resources.getString(R.string.screenshot_share_label),
                 context.resources.getString(R.string.screenshot_share_description),
-            ) {
-                debugLog(LogConfig.DEBUG_ACTIONS) { "Share tapped" }
-                uiEventLogger.log(SCREENSHOT_SHARE_TAPPED, 0, request.packageNameString)
-                onDeferrableActionTapped { result ->
-                    actionExecutor.startSharedTransition(
-                        createShareWithSubject(result.uri, result.subject),
-                        result.user,
-                        false
-                    )
-                }
+            )
+        ) {
+            debugLog(LogConfig.DEBUG_ACTIONS) { "Share tapped" }
+            uiEventLogger.log(SCREENSHOT_SHARE_TAPPED, 0, request.packageNameString)
+            onDeferrableActionTapped { result ->
+                actionExecutor.startSharedTransition(
+                    createShareWithSubject(result.uri, result.subject),
+                    result.user,
+                    false
+                )
             }
-        )
+        }
+
         smartActionsProvider.requestQuickShare(request, requestId) { quickShare ->
             if (!quickShare.actionIntent.isImmutable) {
                 viewModel.addAction(
-                    ActionButtonViewModel(
+                    ActionButtonAppearance(
                         quickShare.getIcon().loadDrawable(context),
                         quickShare.title,
                         quickShare.title
-                    ) {
-                        debugLog(LogConfig.DEBUG_ACTIONS) { "Quickshare tapped" }
-                        onDeferrableActionTapped { result ->
-                            uiEventLogger.log(
-                                SCREENSHOT_SMART_ACTION_TAPPED,
-                                0,
-                                request.packageNameString
+                    )
+                ) {
+                    debugLog(LogConfig.DEBUG_ACTIONS) { "Quickshare tapped" }
+                    onDeferrableActionTapped { result ->
+                        uiEventLogger.log(
+                            SCREENSHOT_SMART_ACTION_TAPPED,
+                            0,
+                            request.packageNameString
+                        )
+                        val pendingIntentWithUri =
+                            smartActionsProvider.wrapIntent(
+                                quickShare,
+                                result.uri,
+                                result.subject,
+                                requestId
                             )
-                            val pendingIntentWithUri =
-                                smartActionsProvider.wrapIntent(
-                                    quickShare,
-                                    result.uri,
-                                    result.subject,
-                                    requestId
-                                )
-                            actionExecutor.sendPendingIntent(pendingIntentWithUri)
-                        }
+                        actionExecutor.sendPendingIntent(pendingIntentWithUri)
                     }
-                )
+                }
             } else {
                 Log.w(TAG, "Received immutable quick share pending intent; ignoring")
             }
@@ -148,14 +150,14 @@
 
     override fun onScrollChipReady(onClick: Runnable) {
         viewModel.addAction(
-            ActionButtonViewModel(
+            ActionButtonAppearance(
                 AppCompatResources.getDrawable(context, R.drawable.ic_screenshot_scroll),
                 context.resources.getString(R.string.screenshot_scroll_label),
                 context.resources.getString(R.string.screenshot_scroll_label),
-            ) {
-                onClick.run()
-            }
-        )
+            )
+        ) {
+            onClick.run()
+        }
     }
 
     override fun setCompletedScreenshot(result: ScreenshotSavedResult) {
@@ -166,13 +168,19 @@
         this.result = result
         pendingAction?.invoke(result)
         smartActionsProvider.requestSmartActions(request, requestId, result) { smartActions ->
-            viewModel.addActions(
-                smartActions.map {
-                    ActionButtonViewModel(it.getIcon().loadDrawable(context), it.title, it.title) {
-                        actionExecutor.sendPendingIntent(it.actionIntent)
+            smartActions.forEach {
+                smartActions.forEach { action ->
+                    viewModel.addAction(
+                        ActionButtonAppearance(
+                            action.getIcon().loadDrawable(context),
+                            action.title,
+                            action.title,
+                        )
+                    ) {
+                        actionExecutor.sendPendingIntent(action.actionIntent)
                     }
                 }
-            )
+            }
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt
index a6374ae..3c5a0ec 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt
@@ -28,16 +28,16 @@
     fun bind(view: View, viewModel: ActionButtonViewModel) {
         val iconView = view.requireViewById<ImageView>(R.id.overlay_action_chip_icon)
         val textView = view.requireViewById<TextView>(R.id.overlay_action_chip_text)
-        iconView.setImageDrawable(viewModel.icon)
-        textView.text = viewModel.name
-        setMargins(iconView, textView, viewModel.name?.isNotEmpty() ?: false)
+        iconView.setImageDrawable(viewModel.appearance.icon)
+        textView.text = viewModel.appearance.label
+        setMargins(iconView, textView, viewModel.appearance.label?.isNotEmpty() ?: false)
         if (viewModel.onClicked != null) {
             view.setOnClickListener { viewModel.onClicked.invoke() }
         } else {
             view.setOnClickListener(null)
         }
         view.tag = viewModel.id
-        view.contentDescription = viewModel.description
+        view.contentDescription = viewModel.appearance.description
         view.visibility = View.VISIBLE
         view.alpha = 1f
     }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ActionButtonAppearance.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ActionButtonAppearance.kt
new file mode 100644
index 0000000..55a2ad2
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ActionButtonAppearance.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.screenshot.ui.viewmodel
+
+import android.graphics.drawable.Drawable
+
+/** Data describing how an action should be shown to the user. */
+data class ActionButtonAppearance(
+    val icon: Drawable?,
+    val label: CharSequence?,
+    val description: CharSequence,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ActionButtonViewModel.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ActionButtonViewModel.kt
index 97b24c1..64b0105 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ActionButtonViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ActionButtonViewModel.kt
@@ -16,19 +16,19 @@
 
 package com.android.systemui.screenshot.ui.viewmodel
 
-import android.graphics.drawable.Drawable
-
 data class ActionButtonViewModel(
-    val icon: Drawable?,
-    val name: CharSequence?,
-    val description: CharSequence,
+    val appearance: ActionButtonAppearance,
+    val id: Int,
     val onClicked: (() -> Unit)?,
 ) {
-    val id: Int = getId()
-
     companion object {
         private var nextId = 0
 
         private fun getId() = nextId.also { nextId += 1 }
+
+        fun withNextId(
+            appearance: ActionButtonAppearance,
+            onClicked: (() -> Unit)?
+        ): ActionButtonViewModel = ActionButtonViewModel(appearance, getId(), onClicked)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt
index ddfa69b..fa34803 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.screenshot.ui.viewmodel
 
 import android.graphics.Bitmap
+import android.util.Log
 import android.view.accessibility.AccessibilityManager
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
@@ -39,16 +40,34 @@
         _previewAction.value = onClick
     }
 
-    fun addAction(action: ActionButtonViewModel) {
+    fun addAction(actionAppearance: ActionButtonAppearance, onClicked: (() -> Unit)): Int {
         val actionList = _actions.value.toMutableList()
+        val action = ActionButtonViewModel.withNextId(actionAppearance, onClicked)
         actionList.add(action)
         _actions.value = actionList
+        return action.id
     }
 
-    fun addActions(actions: List<ActionButtonViewModel>) {
+    fun updateActionAppearance(actionId: Int, appearance: ActionButtonAppearance) {
         val actionList = _actions.value.toMutableList()
-        actionList.addAll(actions)
-        _actions.value = actionList
+        val index = actionList.indexOfFirst { it.id == actionId }
+        if (index >= 0) {
+            actionList[index] =
+                ActionButtonViewModel(appearance, actionId, actionList[index].onClicked)
+            _actions.value = actionList
+        } else {
+            Log.w(TAG, "Attempted to update unknown action id $actionId")
+        }
+    }
+
+    fun removeAction(actionId: Int) {
+        val actionList = _actions.value.toMutableList()
+        if (actionList.removeIf { it.id == actionId }) {
+            // Update if something was removed.
+            _actions.value = actionList
+        } else {
+            Log.w(TAG, "Attempted to remove unknown action id $actionId")
+        }
     }
 
     fun reset() {
@@ -56,4 +75,8 @@
         _previewAction.value = null
         _actions.value = listOf()
     }
+
+    companion object {
+        const val TAG = "ScreenshotViewModel"
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/wallpapers/ImageWallpaper.java b/packages/SystemUI/src/com/android/systemui/wallpapers/ImageWallpaper.java
index d00081e..1568e8c0 100644
--- a/packages/SystemUI/src/com/android/systemui/wallpapers/ImageWallpaper.java
+++ b/packages/SystemUI/src/com/android/systemui/wallpapers/ImageWallpaper.java
@@ -20,6 +20,9 @@
 import static android.app.WallpaperManager.FLAG_SYSTEM;
 import static android.app.WallpaperManager.SetWallpaperFlags;
 
+import static com.android.window.flags.Flags.offloadColorExtraction;
+
+import android.annotation.Nullable;
 import android.app.WallpaperColors;
 import android.app.WallpaperManager;
 import android.content.Context;
@@ -135,6 +138,12 @@
                     mLongExecutor,
                     mLock,
                     new WallpaperLocalColorExtractor.WallpaperLocalColorExtractorCallback() {
+
+                        @Override
+                        public void onColorsProcessed() {
+                            CanvasEngine.this.notifyColorsChanged();
+                        }
+
                         @Override
                         public void onColorsProcessed(List<RectF> regions,
                                 List<WallpaperColors> colors) {
@@ -433,6 +442,12 @@
         }
 
         @Override
+        public @Nullable WallpaperColors onComputeColors() {
+            if (!offloadColorExtraction()) return null;
+            return mWallpaperLocalColorExtractor.onComputeColors();
+        }
+
+        @Override
         public boolean supportsLocalColorExtraction() {
             return true;
         }
@@ -469,6 +484,12 @@
         }
 
         @Override
+        public void onDimAmountChanged(float dimAmount) {
+            if (!offloadColorExtraction()) return;
+            mWallpaperLocalColorExtractor.onDimAmountChanged(dimAmount);
+        }
+
+        @Override
         public void onDisplayAdded(int displayId) {
 
         }
diff --git a/packages/SystemUI/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractor.java b/packages/SystemUI/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractor.java
index e2ec8dc..d37dfb4 100644
--- a/packages/SystemUI/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractor.java
+++ b/packages/SystemUI/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractor.java
@@ -17,6 +17,8 @@
 
 package com.android.systemui.wallpapers;
 
+import static com.android.window.flags.Flags.offloadColorExtraction;
+
 import android.app.WallpaperColors;
 import android.graphics.Bitmap;
 import android.graphics.Rect;
@@ -66,6 +68,12 @@
     private final List<RectF> mPendingRegions = new ArrayList<>();
     private final Set<RectF> mProcessedRegions = new ArraySet<>();
 
+    private float mWallpaperDimAmount = 0f;
+    private WallpaperColors mWallpaperColors;
+
+    // By default we assume that colors were loaded from disk and don't need to be recomputed
+    private boolean mRecomputeColors = false;
+
     @LongRunning
     private final Executor mLongExecutor;
 
@@ -75,6 +83,12 @@
      * Interface to handle the callbacks after the different steps of the color extraction
      */
     public interface WallpaperLocalColorExtractorCallback {
+
+        /**
+         * Callback after the wallpaper colors have been computed
+         */
+        void onColorsProcessed();
+
         /**
          * Callback after the colors of new regions have been extracted
          * @param regions the list of new regions that have been processed
@@ -129,7 +143,7 @@
             if (displayWidth == mDisplayWidth && displayHeight == mDisplayHeight) return;
             mDisplayWidth = displayWidth;
             mDisplayHeight = displayHeight;
-            processColorsInternal();
+            processLocalColorsInternal();
         }
     }
 
@@ -166,7 +180,8 @@
             mBitmapHeight = bitmap.getHeight();
             mMiniBitmap = createMiniBitmap(bitmap);
             mWallpaperLocalColorExtractorCallback.onMiniBitmapUpdated();
-            recomputeColors();
+            if (offloadColorExtraction() && mRecomputeColors) recomputeColorsInternal();
+            recomputeLocalColors();
         }
     }
 
@@ -184,16 +199,66 @@
             if (mPages == pages) return;
             mPages = pages;
             if (mMiniBitmap != null && !mMiniBitmap.isRecycled()) {
-                recomputeColors();
+                recomputeLocalColors();
             }
         }
     }
 
-    // helper to recompute colors, to be called in synchronized methods
-    private void recomputeColors() {
+    /**
+     * Should be called when the dim amount of the wallpaper changes, to recompute the colors
+     */
+    public void onDimAmountChanged(float dimAmount) {
+        mLongExecutor.execute(() -> onDimAmountChangedSynchronized(dimAmount));
+    }
+
+    private void onDimAmountChangedSynchronized(float dimAmount) {
+        synchronized (mLock) {
+            if (mWallpaperDimAmount == dimAmount) return;
+            mWallpaperDimAmount = dimAmount;
+            mRecomputeColors = true;
+            recomputeColorsInternal();
+        }
+    }
+
+    /**
+     * To be called by {@link ImageWallpaper.CanvasEngine#onComputeColors}. This will either
+     * return the current wallpaper colors, or if the bitmap is not yet loaded, return null and call
+     * {@link WallpaperLocalColorExtractorCallback#onColorsProcessed()} when the colors are ready.
+     */
+    public WallpaperColors onComputeColors() {
+        mLongExecutor.execute(this::onComputeColorsSynchronized);
+        return mWallpaperColors;
+    }
+
+    private void onComputeColorsSynchronized() {
+        synchronized (mLock) {
+            if (mRecomputeColors) return;
+            mRecomputeColors = true;
+            recomputeColorsInternal();
+        }
+    }
+
+    /**
+     * helper to recompute main colors, to be called in synchronized methods
+     */
+    private void recomputeColorsInternal() {
+        if (mMiniBitmap == null) return;
+        mWallpaperColors = getWallpaperColors(mMiniBitmap, mWallpaperDimAmount);
+        mWallpaperLocalColorExtractorCallback.onColorsProcessed();
+    }
+
+    @VisibleForTesting
+    WallpaperColors getWallpaperColors(@NonNull Bitmap bitmap, float dimAmount) {
+        return WallpaperColors.fromBitmap(bitmap, dimAmount);
+    }
+
+    /**
+     * helper to recompute local colors, to be called in synchronized methods
+     */
+    private void recomputeLocalColors() {
         mPendingRegions.addAll(mProcessedRegions);
         mProcessedRegions.clear();
-        processColorsInternal();
+        processLocalColorsInternal();
     }
 
     /**
@@ -216,7 +281,7 @@
             if (!wasActive && isActive()) {
                 mWallpaperLocalColorExtractorCallback.onActivated();
             }
-            processColorsInternal();
+            processLocalColorsInternal();
         }
     }
 
@@ -353,7 +418,7 @@
      * then notify the callback with the resulting colors for these regions
      * This method should only be called synchronously
      */
-    private void processColorsInternal() {
+    private void processLocalColorsInternal() {
         /*
          * if the miniBitmap is not yet loaded, that means the onBitmapChanged has not yet been
          * called, and thus the wallpaper is not yet loaded. In that case, exit, the function
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
index 44798ea..8b79fa4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorControllerTest.kt
@@ -257,6 +257,7 @@
             userId = userId,
             colorBackground = 0,
             isForegroundTask = isForegroundTask,
+            userType = RecentTask.UserType.STANDARD,
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/IconLoaderLibAppIconLoaderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/BasicPackageManagerAppIconLoaderTest.kt
similarity index 78%
rename from packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/IconLoaderLibAppIconLoaderTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/BasicPackageManagerAppIconLoaderTest.kt
index 9b346d0..fa1c8f8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/IconLoaderLibAppIconLoaderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/BasicPackageManagerAppIconLoaderTest.kt
@@ -20,15 +20,10 @@
 import android.content.pm.ActivityInfo
 import android.content.pm.PackageManager
 import android.graphics.Bitmap
-import android.graphics.drawable.Drawable
 import androidx.test.filters.SmallTest
-import com.android.launcher3.icons.BitmapInfo
 import com.android.launcher3.icons.FastBitmapDrawable
-import com.android.launcher3.icons.IconFactory
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.shared.system.PackageManagerWrapper
-import com.android.systemui.util.mockito.any
-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
@@ -40,20 +35,17 @@
 
 @SmallTest
 @RunWith(JUnit4::class)
-class IconLoaderLibAppIconLoaderTest : SysuiTestCase() {
+class BasicPackageManagerAppIconLoaderTest : SysuiTestCase() {
 
-    private val iconFactory: IconFactory = mock()
     private val packageManagerWrapper: PackageManagerWrapper = mock()
     private val packageManager: PackageManager = mock()
     private val dispatcher = Dispatchers.Unconfined
 
     private val appIconLoader =
-        IconLoaderLibAppIconLoader(
+        BasicPackageManagerAppIconLoader(
             backgroundDispatcher = dispatcher,
-            context = context,
             packageManagerWrapper = packageManagerWrapper,
             packageManager = packageManager,
-            iconFactoryProvider = { iconFactory }
         )
 
     @Test
@@ -70,12 +62,7 @@
     private fun givenIcon(component: ComponentName, userId: Int, icon: FastBitmapDrawable) {
         val activityInfo = mock<ActivityInfo>()
         whenever(packageManagerWrapper.getActivityInfo(component, userId)).thenReturn(activityInfo)
-        val rawIcon = mock<Drawable>()
-        whenever(activityInfo.loadIcon(packageManager)).thenReturn(rawIcon)
-
-        val bitmapInfo = mock<BitmapInfo>()
-        whenever(iconFactory.createBadgedIconBitmap(eq(rawIcon), any())).thenReturn(bitmapInfo)
-        whenever(bitmapInfo.newIcon(context)).thenReturn(icon)
+        whenever(activityInfo.loadIcon(packageManager)).thenReturn(icon)
     }
 
     private fun createIcon(): FastBitmapDrawable =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
index b593def..dd62112 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/data/ShellRecentTaskListProviderTest.kt
@@ -1,9 +1,15 @@
 package com.android.systemui.mediaprojection.appselector.data
 
 import android.app.ActivityManager.RecentTaskInfo
+import android.content.pm.UserInfo
+import android.os.UserManager
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.CLONED
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.PRIVATE
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.STANDARD
+import com.android.systemui.mediaprojection.appselector.data.RecentTask.UserType.WORK
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
@@ -17,6 +23,7 @@
 import kotlinx.coroutines.runBlocking
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
 
 @RunWith(AndroidTestingRunner::class)
 @SmallTest
@@ -25,12 +32,16 @@
     private val dispatcher = Dispatchers.Unconfined
     private val recentTasks: RecentTasks = mock()
     private val userTracker: UserTracker = mock()
+    private val userManager: UserManager = mock {
+        whenever(getUserInfo(anyInt())).thenReturn(mock())
+    }
     private val recentTaskListProvider =
         ShellRecentTaskListProvider(
             dispatcher,
             Runnable::run,
             Optional.of(recentTasks),
-            userTracker
+            userTracker,
+            userManager,
         )
 
     @Test
@@ -147,6 +158,22 @@
             .inOrder()
     }
 
+    @Test
+    fun loadRecentTasks_assignsCorrectUserType() {
+        givenRecentTasks(
+            createSingleTask(taskId = 1, userId = 10, userType = STANDARD),
+            createSingleTask(taskId = 2, userId = 20, userType = WORK),
+            createSingleTask(taskId = 3, userId = 30, userType = CLONED),
+            createSingleTask(taskId = 4, userId = 40, userType = PRIVATE),
+        )
+
+        val result = runBlocking { recentTaskListProvider.loadRecentTasks() }
+
+        assertThat(result.map { it.userType })
+            .containsExactly(STANDARD, WORK, CLONED, PRIVATE)
+            .inOrder()
+    }
+
     @Suppress("UNCHECKED_CAST")
     private fun givenRecentTasks(vararg tasks: GroupedRecentTaskInfo) {
         whenever(recentTasks.getRecentTasks(any(), any(), any(), any(), any())).thenAnswer {
@@ -155,7 +182,10 @@
         }
     }
 
-    private fun createRecentTask(taskId: Int): RecentTask =
+    private fun createRecentTask(
+        taskId: Int,
+        userType: RecentTask.UserType = STANDARD
+    ): RecentTask =
         RecentTask(
             taskId = taskId,
             displayId = 0,
@@ -164,25 +194,42 @@
             baseIntentComponent = null,
             colorBackground = null,
             isForegroundTask = false,
+            userType = userType,
         )
 
-    private fun createSingleTask(taskId: Int, isVisible: Boolean = false): GroupedRecentTaskInfo =
-        GroupedRecentTaskInfo.forSingleTask(createTaskInfo(taskId, isVisible))
+    private fun createSingleTask(
+        taskId: Int,
+        userId: Int = 0,
+        isVisible: Boolean = false,
+        userType: RecentTask.UserType = STANDARD,
+    ): GroupedRecentTaskInfo {
+        val userInfo =
+            mock<UserInfo> {
+                whenever(isCloneProfile).thenReturn(userType == CLONED)
+                whenever(isManagedProfile).thenReturn(userType == WORK)
+                whenever(isPrivateProfile).thenReturn(userType == PRIVATE)
+            }
+        whenever(userManager.getUserInfo(userId)).thenReturn(userInfo)
+        return GroupedRecentTaskInfo.forSingleTask(createTaskInfo(taskId, userId, isVisible))
+    }
 
     private fun createTaskPair(
         taskId1: Int,
+        userId1: Int = 0,
         taskId2: Int,
+        userId2: Int = 0,
         isVisible: Boolean = false
     ): GroupedRecentTaskInfo =
         GroupedRecentTaskInfo.forSplitTasks(
-            createTaskInfo(taskId1, isVisible),
-            createTaskInfo(taskId2, isVisible),
+            createTaskInfo(taskId1, userId1, isVisible),
+            createTaskInfo(taskId2, userId2, isVisible),
             null
         )
 
-    private fun createTaskInfo(taskId: Int, isVisible: Boolean = false) =
+    private fun createTaskInfo(taskId: Int, userId: Int, isVisible: Boolean = false) =
         RecentTaskInfo().apply {
             this.taskId = taskId
             this.isVisible = isVisible
+            this.userId = userId
         }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt
index ac41073..a84008b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediaprojection/appselector/view/MediaProjectionRecentsViewControllerTest.kt
@@ -56,7 +56,8 @@
             topActivityComponent = null,
             baseIntentComponent = null,
             colorBackground = null,
-            isForegroundTask = false
+            isForegroundTask = false,
+            userType = RecentTask.UserType.STANDARD,
         )
 
     private val taskView =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModelTest.kt
new file mode 100644
index 0000000..d44e26c
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModelTest.kt
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot.ui.viewmodel
+
+import android.view.accessibility.AccessibilityManager
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.mockito.Mockito.mock
+
+@SmallTest
+class ScreenshotViewModelTest {
+    private val accessibilityManager: AccessibilityManager = mock(AccessibilityManager::class.java)
+    private val appearance = ActionButtonAppearance(null, "Label", "Description")
+    private val onclick = {}
+
+    @Test
+    fun testAddAction() {
+        val viewModel = ScreenshotViewModel(accessibilityManager)
+
+        assertThat(viewModel.actions.value).isEmpty()
+
+        viewModel.addAction(appearance, onclick)
+
+        assertThat(viewModel.actions.value).hasSize(1)
+
+        val added = viewModel.actions.value[0]
+        assertThat(added.appearance).isEqualTo(appearance)
+        assertThat(added.onClicked).isEqualTo(onclick)
+    }
+
+    @Test
+    fun testRemoveAction() {
+        val viewModel = ScreenshotViewModel(accessibilityManager)
+        val firstId = viewModel.addAction(ActionButtonAppearance(null, "", ""), {})
+        val secondId = viewModel.addAction(appearance, onclick)
+
+        assertThat(viewModel.actions.value).hasSize(2)
+        assertThat(firstId).isNotEqualTo(secondId)
+
+        viewModel.removeAction(firstId)
+
+        assertThat(viewModel.actions.value).hasSize(1)
+
+        val remaining = viewModel.actions.value[0]
+        assertThat(remaining.appearance).isEqualTo(appearance)
+        assertThat(remaining.onClicked).isEqualTo(onclick)
+    }
+
+    @Test
+    fun testUpdateActionAppearance() {
+        val viewModel = ScreenshotViewModel(accessibilityManager)
+        val id = viewModel.addAction(appearance, onclick)
+        val otherAppearance = ActionButtonAppearance(null, "Other", "Other")
+
+        viewModel.updateActionAppearance(id, otherAppearance)
+
+        assertThat(viewModel.actions.value).hasSize(1)
+        val updated = viewModel.actions.value[0]
+        assertThat(updated.appearance).isEqualTo(otherAppearance)
+        assertThat(updated.onClicked).isEqualTo(onclick)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractorTest.java b/packages/SystemUI/tests/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractorTest.java
index 1125d41..0eba21a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wallpapers/WallpaperLocalColorExtractorTest.java
@@ -16,9 +16,12 @@
 
 package com.android.systemui.wallpapers;
 
+import static com.android.window.flags.Flags.FLAG_OFFLOAD_COLOR_EXTRACTION;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyFloat;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doNothing;
@@ -32,6 +35,7 @@
 import android.graphics.Color;
 import android.graphics.Rect;
 import android.graphics.RectF;
+import android.platform.test.annotations.EnableFlags;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 
@@ -77,6 +81,7 @@
     private Executor mBackgroundExecutor;
 
     private int mColorsProcessed;
+    private int mLocalColorsProcessed;
     private int mMiniBitmapUpdatedCount;
     private int mActivatedCount;
     private int mDeactivatedCount;
@@ -93,6 +98,7 @@
 
     private void resetCounters() {
         mColorsProcessed = 0;
+        mLocalColorsProcessed = 0;
         mMiniBitmapUpdatedCount = 0;
         mActivatedCount = 0;
         mDeactivatedCount = 0;
@@ -112,10 +118,14 @@
                 new Object(),
                 new WallpaperLocalColorExtractor.WallpaperLocalColorExtractorCallback() {
                     @Override
+                    public void onColorsProcessed() {
+                        mColorsProcessed++;
+                    }
+                    @Override
                     public void onColorsProcessed(List<RectF> regions,
                             List<WallpaperColors> colors) {
                         assertThat(regions.size()).isEqualTo(colors.size());
-                        mColorsProcessed += regions.size();
+                        mLocalColorsProcessed += regions.size();
                     }
 
                     @Override
@@ -148,8 +158,10 @@
                 .when(spyColorExtractor)
                 .createMiniBitmap(any(Bitmap.class), anyInt(), anyInt());
 
-        doReturn(new WallpaperColors(Color.valueOf(0), Color.valueOf(0), Color.valueOf(0)))
-                .when(spyColorExtractor).getLocalWallpaperColors(any(Rect.class));
+        WallpaperColors colors = new WallpaperColors(
+                Color.valueOf(0), Color.valueOf(0), Color.valueOf(0));
+        doReturn(colors).when(spyColorExtractor).getLocalWallpaperColors(any(Rect.class));
+        doReturn(colors).when(spyColorExtractor).getWallpaperColors(any(Bitmap.class), anyFloat());
 
         return spyColorExtractor;
     }
@@ -244,7 +256,7 @@
 
             assertThat(mActivatedCount).isEqualTo(1);
             assertThat(mMiniBitmapUpdatedCount).isEqualTo(1);
-            assertThat(mColorsProcessed).isEqualTo(regions.size());
+            assertThat(mLocalColorsProcessed).isEqualTo(regions.size());
 
             spyColorExtractor.removeLocalColorAreas(regions);
             assertThat(mDeactivatedCount).isEqualTo(1);
@@ -329,12 +341,69 @@
                 spyColorExtractor.onBitmapChanged(newBitmap);
                 assertThat(mMiniBitmapUpdatedCount).isEqualTo(1);
             }
-            assertThat(mColorsProcessed).isEqualTo(regions.size());
+            assertThat(mLocalColorsProcessed).isEqualTo(regions.size());
         }
         spyColorExtractor.removeLocalColorAreas(regions);
         assertThat(mDeactivatedCount).isEqualTo(1);
     }
 
+    /**
+     * Test that after the bitmap changes, the colors are computed only if asked via onComputeColors
+     */
+    @Test
+    @EnableFlags(FLAG_OFFLOAD_COLOR_EXTRACTION)
+    public void testRecomputeColors() {
+        resetCounters();
+        Bitmap bitmap = getMockBitmap(HIGH_BMP_WIDTH, HIGH_BMP_HEIGHT);
+        WallpaperLocalColorExtractor spyColorExtractor = getSpyWallpaperLocalColorExtractor();
+        spyColorExtractor.onBitmapChanged(bitmap);
+        assertThat(mColorsProcessed).isEqualTo(0);
+        spyColorExtractor.onComputeColors();
+        assertThat(mColorsProcessed).isEqualTo(1);
+    }
+
+    /**
+     * Test that after onComputeColors is called, the colors are computed once the bitmap is loaded
+     */
+    @Test
+    @EnableFlags(FLAG_OFFLOAD_COLOR_EXTRACTION)
+    public void testRecomputeColorsBeforeBitmapLoaded() {
+        resetCounters();
+        Bitmap bitmap = getMockBitmap(HIGH_BMP_WIDTH, HIGH_BMP_HEIGHT);
+        WallpaperLocalColorExtractor spyColorExtractor = getSpyWallpaperLocalColorExtractor();
+        spyColorExtractor.onComputeColors();
+        spyColorExtractor.onBitmapChanged(bitmap);
+        assertThat(mColorsProcessed).isEqualTo(1);
+    }
+
+    /**
+     * Test that after the dim changes, the colors are computed if the bitmap is already loaded
+     */
+    @Test
+    @EnableFlags(FLAG_OFFLOAD_COLOR_EXTRACTION)
+    public void testRecomputeColorsOnDimChanged() {
+        resetCounters();
+        Bitmap bitmap = getMockBitmap(HIGH_BMP_WIDTH, HIGH_BMP_HEIGHT);
+        WallpaperLocalColorExtractor spyColorExtractor = getSpyWallpaperLocalColorExtractor();
+        spyColorExtractor.onBitmapChanged(bitmap);
+        spyColorExtractor.onDimAmountChanged(0.5f);
+        assertThat(mColorsProcessed).isEqualTo(1);
+    }
+
+    /**
+     * Test that after the dim changes, the colors will be recomputed once the bitmap is loaded
+     */
+    @Test
+    @EnableFlags(FLAG_OFFLOAD_COLOR_EXTRACTION)
+    public void testRecomputeColorsOnDimChangedBeforeBitmapLoaded() {
+        resetCounters();
+        Bitmap bitmap = getMockBitmap(HIGH_BMP_WIDTH, HIGH_BMP_HEIGHT);
+        WallpaperLocalColorExtractor spyColorExtractor = getSpyWallpaperLocalColorExtractor();
+        spyColorExtractor.onDimAmountChanged(0.3f);
+        spyColorExtractor.onBitmapChanged(bitmap);
+        assertThat(mColorsProcessed).isEqualTo(1);
+    }
+
     @Test
     public void testCleanUp() {
         resetCounters();
@@ -346,6 +415,6 @@
         assertThat(mMiniBitmapUpdatedCount).isEqualTo(1);
         spyColorExtractor.cleanUp();
         spyColorExtractor.addLocalColorsAreas(listOfRandomAreas(MIN_AREAS, MAX_AREAS));
-        assertThat(mColorsProcessed).isEqualTo(0);
+        assertThat(mLocalColorsProcessed).isEqualTo(0);
     }
 }
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index 77119d5..f32c11d 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -1197,54 +1197,11 @@
     }
 
     @Override // Binder call
-    public KeyboardLayout[] getKeyboardLayoutsForInputDevice(
-            final InputDeviceIdentifier identifier) {
-        return mKeyboardLayoutManager.getKeyboardLayoutsForInputDevice(identifier);
-    }
-
-    @Override // Binder call
     public KeyboardLayout getKeyboardLayout(String keyboardLayoutDescriptor) {
         return mKeyboardLayoutManager.getKeyboardLayout(keyboardLayoutDescriptor);
     }
 
     @Override // Binder call
-    public String getCurrentKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier) {
-        return mKeyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(identifier);
-    }
-
-    @EnforcePermission(Manifest.permission.SET_KEYBOARD_LAYOUT)
-    @Override // Binder call
-    public void setCurrentKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor) {
-        super.setCurrentKeyboardLayoutForInputDevice_enforcePermission();
-        mKeyboardLayoutManager.setCurrentKeyboardLayoutForInputDevice(identifier,
-                keyboardLayoutDescriptor);
-    }
-
-    @Override // Binder call
-    public String[] getEnabledKeyboardLayoutsForInputDevice(InputDeviceIdentifier identifier) {
-        return mKeyboardLayoutManager.getEnabledKeyboardLayoutsForInputDevice(identifier);
-    }
-
-    @EnforcePermission(Manifest.permission.SET_KEYBOARD_LAYOUT)
-    @Override // Binder call
-    public void addKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor) {
-        super.addKeyboardLayoutForInputDevice_enforcePermission();
-        mKeyboardLayoutManager.addKeyboardLayoutForInputDevice(identifier,
-                keyboardLayoutDescriptor);
-    }
-
-    @EnforcePermission(Manifest.permission.SET_KEYBOARD_LAYOUT)
-    @Override // Binder call
-    public void removeKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor) {
-        super.removeKeyboardLayoutForInputDevice_enforcePermission();
-        mKeyboardLayoutManager.removeKeyboardLayoutForInputDevice(identifier,
-                keyboardLayoutDescriptor);
-    }
-
-    @Override // Binder call
     public KeyboardLayoutSelectionResult getKeyboardLayoutForInputDevice(
             InputDeviceIdentifier identifier, @UserIdInt int userId,
             @NonNull InputMethodInfo imeInfo, @Nullable InputMethodSubtype imeSubtype) {
@@ -1270,11 +1227,6 @@
                 imeInfo, imeSubtype);
     }
 
-
-    public void switchKeyboardLayout(int deviceId, int direction) {
-        mKeyboardLayoutManager.switchKeyboardLayout(deviceId, direction);
-    }
-
     public void setFocusedApplication(int displayId, InputApplicationHandle application) {
         mNative.setFocusedApplication(displayId, application);
     }
diff --git a/services/core/java/com/android/server/input/KeyboardLayoutManager.java b/services/core/java/com/android/server/input/KeyboardLayoutManager.java
index 6610081..9ba647f 100644
--- a/services/core/java/com/android/server/input/KeyboardLayoutManager.java
+++ b/services/core/java/com/android/server/input/KeyboardLayoutManager.java
@@ -60,7 +60,6 @@
 import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.ArrayMap;
-import android.util.FeatureFlagUtils;
 import android.util.Log;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -69,7 +68,6 @@
 import android.view.inputmethod.InputMethodInfo;
 import android.view.inputmethod.InputMethodManager;
 import android.view.inputmethod.InputMethodSubtype;
-import android.widget.Toast;
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
@@ -96,7 +94,6 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.stream.Stream;
 
 /**
  * A component of {@link InputManagerService} responsible for managing Physical Keyboard layouts.
@@ -112,9 +109,8 @@
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
 
     private static final int MSG_UPDATE_EXISTING_DEVICES = 1;
-    private static final int MSG_SWITCH_KEYBOARD_LAYOUT = 2;
-    private static final int MSG_RELOAD_KEYBOARD_LAYOUTS = 3;
-    private static final int MSG_UPDATE_KEYBOARD_LAYOUTS = 4;
+    private static final int MSG_RELOAD_KEYBOARD_LAYOUTS = 2;
+    private static final int MSG_UPDATE_KEYBOARD_LAYOUTS = 3;
 
     private final Context mContext;
     private final NativeInputManagerService mNative;
@@ -126,13 +122,14 @@
     // Connected keyboards with associated keyboard layouts (either auto-detected or manually
     // selected layout).
     private final SparseArray<KeyboardConfiguration> mConfiguredKeyboards = new SparseArray<>();
-    private Toast mSwitchedKeyboardLayoutToast;
 
     // This cache stores "best-matched" layouts so that we don't need to run the matching
     // algorithm repeatedly.
     @GuardedBy("mKeyboardLayoutCache")
     private final Map<String, KeyboardLayoutSelectionResult> mKeyboardLayoutCache =
             new ArrayMap<>();
+
+    private HashSet<String> mAvailableLayouts = new HashSet<>();
     private final Object mImeInfoLock = new Object();
     @Nullable
     @GuardedBy("mImeInfoLock")
@@ -206,68 +203,51 @@
         }
 
         boolean needToShowNotification = false;
-        if (!useNewSettingsUi()) {
-            synchronized (mDataStore) {
-                String layout = getCurrentKeyboardLayoutForInputDevice(inputDevice.getIdentifier());
-                if (layout == null) {
-                    layout = getDefaultKeyboardLayout(inputDevice);
-                    if (layout != null) {
-                        setCurrentKeyboardLayoutForInputDevice(inputDevice.getIdentifier(), layout);
-                    }
-                }
-                if (layout == null) {
-                    // In old settings show notification always until user manually selects a
-                    // layout in the settings.
+        Set<String> selectedLayouts = new HashSet<>();
+        List<ImeInfo> imeInfoList = getImeInfoListForLayoutMapping();
+        List<KeyboardLayoutSelectionResult> resultList = new ArrayList<>();
+        boolean hasMissingLayout = false;
+        for (ImeInfo imeInfo : imeInfoList) {
+            // Check if the layout has been previously configured
+            KeyboardLayoutSelectionResult result = getKeyboardLayoutForInputDeviceInternal(
+                    keyboardIdentifier, imeInfo);
+            if (result.getLayoutDescriptor() != null) {
+                selectedLayouts.add(result.getLayoutDescriptor());
+            } else {
+                hasMissingLayout = true;
+            }
+            resultList.add(result);
+        }
+
+        if (DEBUG) {
+            Slog.d(TAG,
+                    "Layouts selected for input device: " + keyboardIdentifier
+                            + " -> selectedLayouts: " + selectedLayouts);
+        }
+
+        // If even one layout not configured properly, we need to ask user to configure
+        // the keyboard properly from the Settings.
+        if (hasMissingLayout) {
+            selectedLayouts.clear();
+        }
+
+        config.setConfiguredLayouts(selectedLayouts);
+
+        synchronized (mDataStore) {
+            try {
+                final String key = keyboardIdentifier.toString();
+                if (mDataStore.setSelectedKeyboardLayouts(key, selectedLayouts)) {
+                    // Need to show the notification only if layout selection changed
+                    // from the previous configuration
                     needToShowNotification = true;
                 }
-            }
-        } else {
-            Set<String> selectedLayouts = new HashSet<>();
-            List<ImeInfo> imeInfoList = getImeInfoListForLayoutMapping();
-            List<KeyboardLayoutSelectionResult> resultList = new ArrayList<>();
-            boolean hasMissingLayout = false;
-            for (ImeInfo imeInfo : imeInfoList) {
-                // Check if the layout has been previously configured
-                KeyboardLayoutSelectionResult result = getKeyboardLayoutForInputDeviceInternal(
-                        keyboardIdentifier, imeInfo);
-                boolean noLayoutFound = result.getLayoutDescriptor() == null;
-                if (!noLayoutFound) {
-                    selectedLayouts.add(result.getLayoutDescriptor());
+
+                if (shouldLogConfiguration) {
+                    logKeyboardConfigurationEvent(inputDevice, imeInfoList, resultList,
+                            !mDataStore.hasInputDeviceEntry(key));
                 }
-                resultList.add(result);
-                hasMissingLayout |= noLayoutFound;
-            }
-
-            if (DEBUG) {
-                Slog.d(TAG,
-                        "Layouts selected for input device: " + keyboardIdentifier
-                                + " -> selectedLayouts: " + selectedLayouts);
-            }
-
-            // If even one layout not configured properly, we need to ask user to configure
-            // the keyboard properly from the Settings.
-            if (hasMissingLayout) {
-                selectedLayouts.clear();
-            }
-
-            config.setConfiguredLayouts(selectedLayouts);
-
-            synchronized (mDataStore) {
-                try {
-                    final String key = keyboardIdentifier.toString();
-                    if (mDataStore.setSelectedKeyboardLayouts(key, selectedLayouts)) {
-                        // Need to show the notification only if layout selection changed
-                        // from the previous configuration
-                        needToShowNotification = true;
-                    }
-
-                    if (shouldLogConfiguration) {
-                        logKeyboardConfigurationEvent(inputDevice, imeInfoList, resultList,
-                                !mDataStore.hasInputDeviceEntry(key));
-                    }
-                } finally {
-                    mDataStore.saveIfNeeded();
-                }
+            } finally {
+                mDataStore.saveIfNeeded();
             }
         }
         if (needToShowNotification) {
@@ -275,63 +255,6 @@
         }
     }
 
-    private String getDefaultKeyboardLayout(final InputDevice inputDevice) {
-        final Locale systemLocale = mContext.getResources().getConfiguration().locale;
-        // If our locale doesn't have a language for some reason, then we don't really have a
-        // reasonable default.
-        if (TextUtils.isEmpty(systemLocale.getLanguage())) {
-            return null;
-        }
-        final List<KeyboardLayout> layouts = new ArrayList<>();
-        visitAllKeyboardLayouts((resources, keyboardLayoutResId, layout) -> {
-            // Only select a default when we know the layout is appropriate. For now, this
-            // means it's a custom layout for a specific keyboard.
-            if (layout.getVendorId() != inputDevice.getVendorId()
-                    || layout.getProductId() != inputDevice.getProductId()) {
-                return;
-            }
-            final LocaleList locales = layout.getLocales();
-            for (int localeIndex = 0; localeIndex < locales.size(); ++localeIndex) {
-                final Locale locale = locales.get(localeIndex);
-                if (locale != null && isCompatibleLocale(systemLocale, locale)) {
-                    layouts.add(layout);
-                    break;
-                }
-            }
-        });
-
-        if (layouts.isEmpty()) {
-            return null;
-        }
-
-        // First sort so that ones with higher priority are listed at the top
-        Collections.sort(layouts);
-        // Next we want to try to find an exact match of language, country and variant.
-        for (KeyboardLayout layout : layouts) {
-            final LocaleList locales = layout.getLocales();
-            for (int localeIndex = 0; localeIndex < locales.size(); ++localeIndex) {
-                final Locale locale = locales.get(localeIndex);
-                if (locale != null && locale.getCountry().equals(systemLocale.getCountry())
-                        && locale.getVariant().equals(systemLocale.getVariant())) {
-                    return layout.getDescriptor();
-                }
-            }
-        }
-        // Then try an exact match of language and country
-        for (KeyboardLayout layout : layouts) {
-            final LocaleList locales = layout.getLocales();
-            for (int localeIndex = 0; localeIndex < locales.size(); ++localeIndex) {
-                final Locale locale = locales.get(localeIndex);
-                if (locale != null && locale.getCountry().equals(systemLocale.getCountry())) {
-                    return layout.getDescriptor();
-                }
-            }
-        }
-
-        // Give up and just use the highest priority layout with matching language
-        return layouts.get(0).getDescriptor();
-    }
-
     private static boolean isCompatibleLocale(Locale systemLocale, Locale keyboardLocale) {
         // Different languages are never compatible
         if (!systemLocale.getLanguage().equals(keyboardLocale.getLanguage())) {
@@ -343,11 +266,19 @@
                 || systemLocale.getCountry().equals(keyboardLocale.getCountry());
     }
 
+    @MainThread
     private void updateKeyboardLayouts() {
         // Scan all input devices state for keyboard layouts that have been uninstalled.
-        final HashSet<String> availableKeyboardLayouts = new HashSet<String>();
+        final HashSet<String> availableKeyboardLayouts = new HashSet<>();
         visitAllKeyboardLayouts((resources, keyboardLayoutResId, layout) ->
                 availableKeyboardLayouts.add(layout.getDescriptor()));
+
+        // If available layouts don't change, there is no need to reload layouts.
+        if (mAvailableLayouts.equals(availableKeyboardLayouts)) {
+            return;
+        }
+        mAvailableLayouts = availableKeyboardLayouts;
+
         synchronized (mDataStore) {
             try {
                 mDataStore.removeUninstalledKeyboardLayouts(availableKeyboardLayouts);
@@ -374,53 +305,6 @@
     }
 
     @AnyThread
-    public KeyboardLayout[] getKeyboardLayoutsForInputDevice(
-            final InputDeviceIdentifier identifier) {
-        if (useNewSettingsUi()) {
-            // Provide all supported keyboard layouts since Ime info is not provided
-            return getKeyboardLayouts();
-        }
-        final String[] enabledLayoutDescriptors =
-                getEnabledKeyboardLayoutsForInputDevice(identifier);
-        final ArrayList<KeyboardLayout> enabledLayouts =
-                new ArrayList<>(enabledLayoutDescriptors.length);
-        final ArrayList<KeyboardLayout> potentialLayouts = new ArrayList<>();
-        visitAllKeyboardLayouts(new KeyboardLayoutVisitor() {
-            boolean mHasSeenDeviceSpecificLayout;
-
-            @Override
-            public void visitKeyboardLayout(Resources resources,
-                    int keyboardLayoutResId, KeyboardLayout layout) {
-                // First check if it's enabled. If the keyboard layout is enabled then we always
-                // want to return it as a possible layout for the device.
-                for (String s : enabledLayoutDescriptors) {
-                    if (s != null && s.equals(layout.getDescriptor())) {
-                        enabledLayouts.add(layout);
-                        return;
-                    }
-                }
-                // Next find any potential layouts that aren't yet enabled for the device. For
-                // devices that have special layouts we assume there's a reason that the generic
-                // layouts don't work for them so we don't want to return them since it's likely
-                // to result in a poor user experience.
-                if (layout.getVendorId() == identifier.getVendorId()
-                        && layout.getProductId() == identifier.getProductId()) {
-                    if (!mHasSeenDeviceSpecificLayout) {
-                        mHasSeenDeviceSpecificLayout = true;
-                        potentialLayouts.clear();
-                    }
-                    potentialLayouts.add(layout);
-                } else if (layout.getVendorId() == -1 && layout.getProductId() == -1
-                        && !mHasSeenDeviceSpecificLayout) {
-                    potentialLayouts.add(layout);
-                }
-            }
-        });
-        return Stream.concat(enabledLayouts.stream(), potentialLayouts.stream()).toArray(
-                KeyboardLayout[]::new);
-    }
-
-    @AnyThread
     @Nullable
     public KeyboardLayout getKeyboardLayout(@NonNull String keyboardLayoutDescriptor) {
         Objects.requireNonNull(keyboardLayoutDescriptor,
@@ -580,195 +464,16 @@
         return LocaleList.forLanguageTags(languageTags.replace('|', ','));
     }
 
-    @AnyThread
-    @Nullable
-    public String getCurrentKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier) {
-        if (useNewSettingsUi()) {
-            Slog.e(TAG, "getCurrentKeyboardLayoutForInputDevice API not supported");
-            return null;
-        }
-        String key = new KeyboardIdentifier(identifier).toString();
-        synchronized (mDataStore) {
-            String layout;
-            // try loading it using the layout descriptor if we have it
-            layout = mDataStore.getCurrentKeyboardLayout(key);
-            if (layout == null && !key.equals(identifier.getDescriptor())) {
-                // if it doesn't exist fall back to the device descriptor
-                layout = mDataStore.getCurrentKeyboardLayout(identifier.getDescriptor());
-            }
-            if (DEBUG) {
-                Slog.d(TAG, "getCurrentKeyboardLayoutForInputDevice() "
-                        + identifier.toString() + ": " + layout);
-            }
-            return layout;
-        }
-    }
-
-    @AnyThread
-    public void setCurrentKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor) {
-        if (useNewSettingsUi()) {
-            Slog.e(TAG, "setCurrentKeyboardLayoutForInputDevice API not supported");
-            return;
-        }
-
-        Objects.requireNonNull(keyboardLayoutDescriptor,
-                "keyboardLayoutDescriptor must not be null");
-        String key = new KeyboardIdentifier(identifier).toString();
-        synchronized (mDataStore) {
-            try {
-                if (mDataStore.setCurrentKeyboardLayout(key, keyboardLayoutDescriptor)) {
-                    if (DEBUG) {
-                        Slog.d(TAG, "setCurrentKeyboardLayoutForInputDevice() " + identifier
-                                + " key: " + key
-                                + " keyboardLayoutDescriptor: " + keyboardLayoutDescriptor);
-                    }
-                    mHandler.sendEmptyMessage(MSG_RELOAD_KEYBOARD_LAYOUTS);
-                }
-            } finally {
-                mDataStore.saveIfNeeded();
-            }
-        }
-    }
-
-    @AnyThread
-    public String[] getEnabledKeyboardLayoutsForInputDevice(InputDeviceIdentifier identifier) {
-        if (useNewSettingsUi()) {
-            Slog.e(TAG, "getEnabledKeyboardLayoutsForInputDevice API not supported");
-            return new String[0];
-        }
-        String key = new KeyboardIdentifier(identifier).toString();
-        synchronized (mDataStore) {
-            String[] layouts = mDataStore.getKeyboardLayouts(key);
-            if ((layouts == null || layouts.length == 0)
-                    && !key.equals(identifier.getDescriptor())) {
-                layouts = mDataStore.getKeyboardLayouts(identifier.getDescriptor());
-            }
-            return layouts;
-        }
-    }
-
-    @AnyThread
-    public void addKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor) {
-        if (useNewSettingsUi()) {
-            Slog.e(TAG, "addKeyboardLayoutForInputDevice API not supported");
-            return;
-        }
-        Objects.requireNonNull(keyboardLayoutDescriptor,
-                "keyboardLayoutDescriptor must not be null");
-
-        String key = new KeyboardIdentifier(identifier).toString();
-        synchronized (mDataStore) {
-            try {
-                String oldLayout = mDataStore.getCurrentKeyboardLayout(key);
-                if (oldLayout == null && !key.equals(identifier.getDescriptor())) {
-                    oldLayout = mDataStore.getCurrentKeyboardLayout(identifier.getDescriptor());
-                }
-                if (mDataStore.addKeyboardLayout(key, keyboardLayoutDescriptor)
-                        && !Objects.equals(oldLayout,
-                        mDataStore.getCurrentKeyboardLayout(key))) {
-                    mHandler.sendEmptyMessage(MSG_RELOAD_KEYBOARD_LAYOUTS);
-                }
-            } finally {
-                mDataStore.saveIfNeeded();
-            }
-        }
-    }
-
-    @AnyThread
-    public void removeKeyboardLayoutForInputDevice(InputDeviceIdentifier identifier,
-            String keyboardLayoutDescriptor) {
-        if (useNewSettingsUi()) {
-            Slog.e(TAG, "removeKeyboardLayoutForInputDevice API not supported");
-            return;
-        }
-        Objects.requireNonNull(keyboardLayoutDescriptor,
-                "keyboardLayoutDescriptor must not be null");
-
-        String key = new KeyboardIdentifier(identifier).toString();
-        synchronized (mDataStore) {
-            try {
-                String oldLayout = mDataStore.getCurrentKeyboardLayout(key);
-                if (oldLayout == null && !key.equals(identifier.getDescriptor())) {
-                    oldLayout = mDataStore.getCurrentKeyboardLayout(identifier.getDescriptor());
-                }
-                boolean removed = mDataStore.removeKeyboardLayout(key, keyboardLayoutDescriptor);
-                if (!key.equals(identifier.getDescriptor())) {
-                    // We need to remove from both places to ensure it is gone
-                    removed |= mDataStore.removeKeyboardLayout(identifier.getDescriptor(),
-                            keyboardLayoutDescriptor);
-                }
-                if (removed && !Objects.equals(oldLayout,
-                        mDataStore.getCurrentKeyboardLayout(key))) {
-                    mHandler.sendEmptyMessage(MSG_RELOAD_KEYBOARD_LAYOUTS);
-                }
-            } finally {
-                mDataStore.saveIfNeeded();
-            }
-        }
-    }
-
-    @AnyThread
-    public void switchKeyboardLayout(int deviceId, int direction) {
-        if (useNewSettingsUi()) {
-            Slog.e(TAG, "switchKeyboardLayout API not supported");
-            return;
-        }
-        mHandler.obtainMessage(MSG_SWITCH_KEYBOARD_LAYOUT, deviceId, direction).sendToTarget();
-    }
-
-    @MainThread
-    private void handleSwitchKeyboardLayout(int deviceId, int direction) {
-        final InputDevice device = getInputDevice(deviceId);
-        if (device != null) {
-            final boolean changed;
-            final String keyboardLayoutDescriptor;
-
-            String key = new KeyboardIdentifier(device.getIdentifier()).toString();
-            synchronized (mDataStore) {
-                try {
-                    changed = mDataStore.switchKeyboardLayout(key, direction);
-                    keyboardLayoutDescriptor = mDataStore.getCurrentKeyboardLayout(
-                            key);
-                } finally {
-                    mDataStore.saveIfNeeded();
-                }
-            }
-
-            if (changed) {
-                if (mSwitchedKeyboardLayoutToast != null) {
-                    mSwitchedKeyboardLayoutToast.cancel();
-                    mSwitchedKeyboardLayoutToast = null;
-                }
-                if (keyboardLayoutDescriptor != null) {
-                    KeyboardLayout keyboardLayout = getKeyboardLayout(keyboardLayoutDescriptor);
-                    if (keyboardLayout != null) {
-                        mSwitchedKeyboardLayoutToast = Toast.makeText(
-                                mContext, keyboardLayout.getLabel(), Toast.LENGTH_SHORT);
-                        mSwitchedKeyboardLayoutToast.show();
-                    }
-                }
-
-                reloadKeyboardLayouts();
-            }
-        }
-    }
-
     @Nullable
     @AnyThread
     public String[] getKeyboardLayoutOverlay(InputDeviceIdentifier identifier, String languageTag,
             String layoutType) {
         String keyboardLayoutDescriptor;
-        if (useNewSettingsUi()) {
-            synchronized (mImeInfoLock) {
-                KeyboardLayoutSelectionResult result = getKeyboardLayoutForInputDeviceInternal(
-                        new KeyboardIdentifier(identifier, languageTag, layoutType),
-                        mCurrentImeInfo);
-                keyboardLayoutDescriptor = result.getLayoutDescriptor();
-            }
-        } else {
-            keyboardLayoutDescriptor = getCurrentKeyboardLayoutForInputDevice(identifier);
+        synchronized (mImeInfoLock) {
+            KeyboardLayoutSelectionResult result = getKeyboardLayoutForInputDeviceInternal(
+                    new KeyboardIdentifier(identifier, languageTag, layoutType),
+                    mCurrentImeInfo);
+            keyboardLayoutDescriptor = result.getLayoutDescriptor();
         }
         if (keyboardLayoutDescriptor == null) {
             return null;
@@ -797,10 +502,6 @@
     public KeyboardLayoutSelectionResult getKeyboardLayoutForInputDevice(
             InputDeviceIdentifier identifier, @UserIdInt int userId,
             @NonNull InputMethodInfo imeInfo, @Nullable InputMethodSubtype imeSubtype) {
-        if (!useNewSettingsUi()) {
-            Slog.e(TAG, "getKeyboardLayoutForInputDevice() API not supported");
-            return FAILED;
-        }
         InputDevice inputDevice = getInputDevice(identifier);
         if (inputDevice == null || inputDevice.isVirtual() || !inputDevice.isFullKeyboard()) {
             return FAILED;
@@ -820,10 +521,6 @@
             @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
             @Nullable InputMethodSubtype imeSubtype,
             String keyboardLayoutDescriptor) {
-        if (!useNewSettingsUi()) {
-            Slog.e(TAG, "setKeyboardLayoutForInputDevice() API not supported");
-            return;
-        }
         Objects.requireNonNull(keyboardLayoutDescriptor,
                 "keyboardLayoutDescriptor must not be null");
         InputDevice inputDevice = getInputDevice(identifier);
@@ -854,10 +551,6 @@
     public KeyboardLayout[] getKeyboardLayoutListForInputDevice(InputDeviceIdentifier identifier,
             @UserIdInt int userId, @NonNull InputMethodInfo imeInfo,
             @Nullable InputMethodSubtype imeSubtype) {
-        if (!useNewSettingsUi()) {
-            Slog.e(TAG, "getKeyboardLayoutListForInputDevice() API not supported");
-            return new KeyboardLayout[0];
-        }
         InputDevice inputDevice = getInputDevice(identifier);
         if (inputDevice == null || inputDevice.isVirtual() || !inputDevice.isFullKeyboard()) {
             return new KeyboardLayout[0];
@@ -923,10 +616,6 @@
     public void onInputMethodSubtypeChanged(@UserIdInt int userId,
             @Nullable InputMethodSubtypeHandle subtypeHandle,
             @Nullable InputMethodSubtype subtype) {
-        if (!useNewSettingsUi()) {
-            Slog.e(TAG, "onInputMethodSubtypeChanged() API not supported");
-            return;
-        }
         if (subtypeHandle == null) {
             if (DEBUG) {
                 Slog.d(TAG, "No InputMethod is running, ignoring change");
@@ -1289,9 +978,6 @@
                     onInputDeviceAdded(deviceId);
                 }
                 return true;
-            case MSG_SWITCH_KEYBOARD_LAYOUT:
-                handleSwitchKeyboardLayout(msg.arg1, msg.arg2);
-                return true;
             case MSG_RELOAD_KEYBOARD_LAYOUTS:
                 reloadKeyboardLayouts();
                 return true;
@@ -1303,10 +989,6 @@
         }
     }
 
-    private boolean useNewSettingsUi() {
-        return FeatureFlagUtils.isEnabled(mContext, FeatureFlagUtils.SETTINGS_NEW_KEYBOARD_UI);
-    }
-
     @Nullable
     private InputDevice getInputDevice(int deviceId) {
         InputManager inputManager = mContext.getSystemService(InputManager.class);
diff --git a/services/core/java/com/android/server/input/PersistentDataStore.java b/services/core/java/com/android/server/input/PersistentDataStore.java
index 31083fd..7859253 100644
--- a/services/core/java/com/android/server/input/PersistentDataStore.java
+++ b/services/core/java/com/android/server/input/PersistentDataStore.java
@@ -27,7 +27,6 @@
 import android.view.Surface;
 
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
@@ -42,7 +41,6 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -74,7 +72,7 @@
             new HashMap<String, InputDeviceState>();
 
     // The interface for methods which should be replaced by the test harness.
-    private Injector mInjector;
+    private final Injector mInjector;
 
     // True if the data has been loaded.
     private boolean mLoaded;
@@ -83,7 +81,7 @@
     private boolean mDirty;
 
     // Storing key remapping
-    private Map<Integer, Integer> mKeyRemapping = new HashMap<>();
+    private final Map<Integer, Integer> mKeyRemapping = new HashMap<>();
 
     public PersistentDataStore() {
         this(new Injector());
@@ -130,22 +128,6 @@
     }
 
     @Nullable
-    public String getCurrentKeyboardLayout(String inputDeviceDescriptor) {
-        InputDeviceState state = getInputDeviceState(inputDeviceDescriptor);
-        return state != null ? state.getCurrentKeyboardLayout() : null;
-    }
-
-    public boolean setCurrentKeyboardLayout(String inputDeviceDescriptor,
-            String keyboardLayoutDescriptor) {
-        InputDeviceState state = getOrCreateInputDeviceState(inputDeviceDescriptor);
-        if (state.setCurrentKeyboardLayout(keyboardLayoutDescriptor)) {
-            setDirty();
-            return true;
-        }
-        return false;
-    }
-
-    @Nullable
     public String getKeyboardLayout(String inputDeviceDescriptor, String key) {
         InputDeviceState state = getInputDeviceState(inputDeviceDescriptor);
         return state != null ? state.getKeyboardLayout(key) : null;
@@ -171,43 +153,6 @@
         return false;
     }
 
-    public String[] getKeyboardLayouts(String inputDeviceDescriptor) {
-        InputDeviceState state = getInputDeviceState(inputDeviceDescriptor);
-        if (state == null) {
-            return (String[])ArrayUtils.emptyArray(String.class);
-        }
-        return state.getKeyboardLayouts();
-    }
-
-    public boolean addKeyboardLayout(String inputDeviceDescriptor,
-            String keyboardLayoutDescriptor) {
-        InputDeviceState state = getOrCreateInputDeviceState(inputDeviceDescriptor);
-        if (state.addKeyboardLayout(keyboardLayoutDescriptor)) {
-            setDirty();
-            return true;
-        }
-        return false;
-    }
-
-    public boolean removeKeyboardLayout(String inputDeviceDescriptor,
-            String keyboardLayoutDescriptor) {
-        InputDeviceState state = getOrCreateInputDeviceState(inputDeviceDescriptor);
-        if (state.removeKeyboardLayout(keyboardLayoutDescriptor)) {
-            setDirty();
-            return true;
-        }
-        return false;
-    }
-
-    public boolean switchKeyboardLayout(String inputDeviceDescriptor, int direction) {
-        InputDeviceState state = getInputDeviceState(inputDeviceDescriptor);
-        if (state != null && state.switchKeyboardLayout(direction)) {
-            setDirty();
-            return true;
-        }
-        return false;
-    }
-
     public boolean setKeyboardBacklightBrightness(String inputDeviceDescriptor, int lightId,
             int brightness) {
         InputDeviceState state = getOrCreateInputDeviceState(inputDeviceDescriptor);
@@ -417,9 +362,6 @@
                 "x_ymix", "x_offset", "y_xmix", "y_scale", "y_offset" };
 
         private final TouchCalibration[] mTouchCalibration = new TouchCalibration[4];
-        @Nullable
-        private String mCurrentKeyboardLayout;
-        private final ArrayList<String> mKeyboardLayouts = new ArrayList<String>();
         private final SparseIntArray mKeyboardBacklightBrightnessMap = new SparseIntArray();
 
         private final Map<String, String> mKeyboardLayoutMap = new ArrayMap<>();
@@ -465,49 +407,6 @@
             return true;
         }
 
-        @Nullable
-        public String getCurrentKeyboardLayout() {
-            return mCurrentKeyboardLayout;
-        }
-
-        public boolean setCurrentKeyboardLayout(String keyboardLayout) {
-            if (Objects.equals(mCurrentKeyboardLayout, keyboardLayout)) {
-                return false;
-            }
-            addKeyboardLayout(keyboardLayout);
-            mCurrentKeyboardLayout = keyboardLayout;
-            return true;
-        }
-
-        public String[] getKeyboardLayouts() {
-            if (mKeyboardLayouts.isEmpty()) {
-                return (String[])ArrayUtils.emptyArray(String.class);
-            }
-            return mKeyboardLayouts.toArray(new String[mKeyboardLayouts.size()]);
-        }
-
-        public boolean addKeyboardLayout(String keyboardLayout) {
-            int index = Collections.binarySearch(mKeyboardLayouts, keyboardLayout);
-            if (index >= 0) {
-                return false;
-            }
-            mKeyboardLayouts.add(-index - 1, keyboardLayout);
-            if (mCurrentKeyboardLayout == null) {
-                mCurrentKeyboardLayout = keyboardLayout;
-            }
-            return true;
-        }
-
-        public boolean removeKeyboardLayout(String keyboardLayout) {
-            int index = Collections.binarySearch(mKeyboardLayouts, keyboardLayout);
-            if (index < 0) {
-                return false;
-            }
-            mKeyboardLayouts.remove(index);
-            updateCurrentKeyboardLayoutIfRemoved(keyboardLayout, index);
-            return true;
-        }
-
         public boolean setKeyboardBacklightBrightness(int lightId, int brightness) {
             if (mKeyboardBacklightBrightnessMap.get(lightId, INVALID_VALUE) == brightness) {
                 return false;
@@ -521,48 +420,8 @@
             return brightness == INVALID_VALUE ? OptionalInt.empty() : OptionalInt.of(brightness);
         }
 
-        private void updateCurrentKeyboardLayoutIfRemoved(
-                String removedKeyboardLayout, int removedIndex) {
-            if (Objects.equals(mCurrentKeyboardLayout, removedKeyboardLayout)) {
-                if (!mKeyboardLayouts.isEmpty()) {
-                    int index = removedIndex;
-                    if (index == mKeyboardLayouts.size()) {
-                        index = 0;
-                    }
-                    mCurrentKeyboardLayout = mKeyboardLayouts.get(index);
-                } else {
-                    mCurrentKeyboardLayout = null;
-                }
-            }
-        }
-
-        public boolean switchKeyboardLayout(int direction) {
-            final int size = mKeyboardLayouts.size();
-            if (size < 2) {
-                return false;
-            }
-            int index = Collections.binarySearch(mKeyboardLayouts, mCurrentKeyboardLayout);
-            assert index >= 0;
-            if (direction > 0) {
-                index = (index + 1) % size;
-            } else {
-                index = (index + size - 1) % size;
-            }
-            mCurrentKeyboardLayout = mKeyboardLayouts.get(index);
-            return true;
-        }
-
         public boolean removeUninstalledKeyboardLayouts(Set<String> availableKeyboardLayouts) {
             boolean changed = false;
-            for (int i = mKeyboardLayouts.size(); i-- > 0; ) {
-                String keyboardLayout = mKeyboardLayouts.get(i);
-                if (!availableKeyboardLayouts.contains(keyboardLayout)) {
-                    Slog.i(TAG, "Removing uninstalled keyboard layout " + keyboardLayout);
-                    mKeyboardLayouts.remove(i);
-                    updateCurrentKeyboardLayoutIfRemoved(keyboardLayout, i);
-                    changed = true;
-                }
-            }
             List<String> removedEntries = new ArrayList<>();
             for (String key : mKeyboardLayoutMap.keySet()) {
                 if (!availableKeyboardLayouts.contains(mKeyboardLayoutMap.get(key))) {
@@ -582,27 +441,7 @@
                 throws IOException, XmlPullParserException {
             final int outerDepth = parser.getDepth();
             while (XmlUtils.nextElementWithin(parser, outerDepth)) {
-                if (parser.getName().equals("keyboard-layout")) {
-                    String descriptor = parser.getAttributeValue(null, "descriptor");
-                    if (descriptor == null) {
-                        throw new XmlPullParserException(
-                                "Missing descriptor attribute on keyboard-layout.");
-                    }
-                    String current = parser.getAttributeValue(null, "current");
-                    if (mKeyboardLayouts.contains(descriptor)) {
-                        throw new XmlPullParserException(
-                                "Found duplicate keyboard layout.");
-                    }
-
-                    mKeyboardLayouts.add(descriptor);
-                    if (current != null && current.equals("true")) {
-                        if (mCurrentKeyboardLayout != null) {
-                            throw new XmlPullParserException(
-                                    "Found multiple current keyboard layouts.");
-                        }
-                        mCurrentKeyboardLayout = descriptor;
-                    }
-                } else if (parser.getName().equals("keyed-keyboard-layout")) {
+                if (parser.getName().equals("keyed-keyboard-layout")) {
                     String key = parser.getAttributeValue(null, "key");
                     if (key == null) {
                         throw new XmlPullParserException(
@@ -676,27 +515,9 @@
                     }
                 }
             }
-
-            // Maintain invariant that layouts are sorted.
-            Collections.sort(mKeyboardLayouts);
-
-            // Maintain invariant that there is always a current keyboard layout unless
-            // there are none installed.
-            if (mCurrentKeyboardLayout == null && !mKeyboardLayouts.isEmpty()) {
-                mCurrentKeyboardLayout = mKeyboardLayouts.get(0);
-            }
         }
 
         public void saveToXml(TypedXmlSerializer serializer) throws IOException {
-            for (String layout : mKeyboardLayouts) {
-                serializer.startTag(null, "keyboard-layout");
-                serializer.attribute(null, "descriptor", layout);
-                if (layout.equals(mCurrentKeyboardLayout)) {
-                    serializer.attributeBoolean(null, "current", true);
-                }
-                serializer.endTag(null, "keyboard-layout");
-            }
-
             for (String key : mKeyboardLayoutMap.keySet()) {
                 serializer.startTag(null, "keyed-keyboard-layout");
                 serializer.attribute(null, "key", key);
diff --git a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
index 31ce630..ffc2319 100644
--- a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
+++ b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
@@ -196,7 +196,10 @@
         return true;
     }
 
-    private void sendResultReceiverFailure(ResultReceiver resultReceiver) {
+    private void sendResultReceiverFailure(@Nullable ResultReceiver resultReceiver) {
+        if (resultReceiver == null) {
+            return;
+        }
         resultReceiver.send(
                 mIsInputShown.getAsBoolean()
                         ? InputMethodManager.RESULT_UNCHANGED_SHOWN
diff --git a/services/core/java/com/android/server/media/MediaSessionRecord.java b/services/core/java/com/android/server/media/MediaSessionRecord.java
index eb4e6e4..62a9471 100644
--- a/services/core/java/com/android/server/media/MediaSessionRecord.java
+++ b/services/core/java/com/android/server/media/MediaSessionRecord.java
@@ -849,7 +849,7 @@
             }
         }
         if (deadCallbackHolders != null) {
-            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+            removeControllerHoldersSafely(deadCallbackHolders);
         }
     }
 
@@ -876,7 +876,7 @@
             }
         }
         if (deadCallbackHolders != null) {
-            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+            removeControllerHoldersSafely(deadCallbackHolders);
         }
     }
 
@@ -911,7 +911,7 @@
             }
         }
         if (deadCallbackHolders != null) {
-            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+            removeControllerHoldersSafely(deadCallbackHolders);
         }
     }
 
@@ -938,7 +938,7 @@
             }
         }
         if (deadCallbackHolders != null) {
-            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+            removeControllerHoldersSafely(deadCallbackHolders);
         }
     }
 
@@ -965,7 +965,7 @@
             }
         }
         if (deadCallbackHolders != null) {
-            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+            removeControllerHoldersSafely(deadCallbackHolders);
         }
     }
 
@@ -992,7 +992,7 @@
             }
         }
         if (deadCallbackHolders != null) {
-            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+            removeControllerHoldersSafely(deadCallbackHolders);
         }
     }
 
@@ -1017,7 +1017,7 @@
             }
         }
         if (deadCallbackHolders != null) {
-            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+            removeControllerHoldersSafely(deadCallbackHolders);
         }
     }
 
@@ -1042,7 +1042,7 @@
             }
         }
         // After notifying clear all listeners
-        mControllerCallbackHolders.clear();
+        removeControllerHoldersSafely(null);
     }
 
     private PlaybackState getStateWithUpdatedPosition() {
@@ -1090,6 +1090,17 @@
         return -1;
     }
 
+    private void removeControllerHoldersSafely(
+            Collection<ISessionControllerCallbackHolder> holders) {
+        synchronized (mLock) {
+            if (holders == null) {
+                mControllerCallbackHolders.clear();
+            } else {
+                mControllerCallbackHolders.removeAll(holders);
+            }
+        }
+    }
+
     private PlaybackInfo getVolumeAttributes() {
         int volumeType;
         AudioAttributes attributes;
diff --git a/services/core/java/com/android/server/os/BugreportManagerServiceImpl.java b/services/core/java/com/android/server/os/BugreportManagerServiceImpl.java
index 4eb8b2b..c8fd7e4 100644
--- a/services/core/java/com/android/server/os/BugreportManagerServiceImpl.java
+++ b/services/core/java/com/android/server/os/BugreportManagerServiceImpl.java
@@ -184,6 +184,16 @@
                         throwInvalidBugreportFileForCallerException(
                                 bugreportFile, callingInfo.second);
                     }
+
+                    boolean keepBugreportOnRetrieval = false;
+                    if (onboardingBugreportV2Enabled()) {
+                        keepBugreportOnRetrieval = mBugreportFilesToPersist.contains(
+                                bugreportFile);
+                    }
+
+                    if (!keepBugreportOnRetrieval) {
+                        bugreportFilesForUid.remove(bugreportFile);
+                    }
                 } else {
                     ArraySet<String> bugreportFilesForCaller = mBugreportFiles.get(callingInfo);
                     if (bugreportFilesForCaller != null
diff --git a/services/core/java/com/android/server/pm/PackageInstallerSession.java b/services/core/java/com/android/server/pm/PackageInstallerSession.java
index 3eeeae7..80a5f3a 100644
--- a/services/core/java/com/android/server/pm/PackageInstallerSession.java
+++ b/services/core/java/com/android/server/pm/PackageInstallerSession.java
@@ -1085,11 +1085,17 @@
         final boolean isUpdateOwnershipEnforcementEnabled =
                 mPm.isUpdateOwnershipEnforcementAvailable()
                         && existingUpdateOwnerPackageName != null;
+        // For an installation that un-archives an app, if the installer doesn't have the
+        // INSTALL_PACKAGES permission, the user should have already been prompted to confirm the
+        // un-archive request. There's no need for another confirmation during the installation.
+        final boolean isInstallUnarchive =
+                (params.installFlags & PackageManager.INSTALL_UNARCHIVE) != 0;
 
         // Device owners and affiliated profile owners are allowed to silently install packages, so
         // the permission check is waived if the installer is the device owner.
         final boolean noUserActionNecessary = isInstallerRoot || isInstallerSystem
-                || isInstallerDeviceOwnerOrAffiliatedProfileOwner() || isEmergencyInstall;
+                || isInstallerDeviceOwnerOrAffiliatedProfileOwner() || isEmergencyInstall
+                || isInstallUnarchive;
 
         if (noUserActionNecessary) {
             return userActionNotTypicallyNeededResponse;
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 7db83d7..b4919a4 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -174,7 +174,6 @@
 import android.service.vr.IPersistentVrStateCallbacks;
 import android.speech.RecognizerIntent;
 import android.telecom.TelecomManager;
-import android.util.FeatureFlagUtils;
 import android.util.Log;
 import android.util.MathUtils;
 import android.util.MutableBoolean;
@@ -4121,14 +4120,10 @@
 
     private void handleSwitchKeyboardLayout(@NonNull KeyEvent event, int direction,
             IBinder focusedToken) {
-        if (FeatureFlagUtils.isEnabled(mContext, FeatureFlagUtils.SETTINGS_NEW_KEYBOARD_UI)) {
-            IBinder targetWindowToken =
-                    mWindowManagerInternal.getTargetWindowTokenFromInputToken(focusedToken);
-            InputMethodManagerInternal.get().onSwitchKeyboardLayoutShortcut(direction,
-                    event.getDisplayId(), targetWindowToken);
-        } else {
-            mWindowManagerFuncs.switchKeyboardLayout(event.getDeviceId(), direction);
-        }
+        IBinder targetWindowToken =
+                mWindowManagerInternal.getTargetWindowTokenFromInputToken(focusedToken);
+        InputMethodManagerInternal.get().onSwitchKeyboardLayoutShortcut(direction,
+                event.getDisplayId(), targetWindowToken);
     }
 
     private boolean interceptFallback(IBinder focusedToken, KeyEvent fallbackEvent,
diff --git a/services/core/java/com/android/server/policy/WindowManagerPolicy.java b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
index 2623025..9ca4e27 100644
--- a/services/core/java/com/android/server/policy/WindowManagerPolicy.java
+++ b/services/core/java/com/android/server/policy/WindowManagerPolicy.java
@@ -251,12 +251,6 @@
          */
         public int getCameraLensCoverState();
 
-        /**
-         * Switch the keyboard layout for the given device.
-         * Direction should be +1 or -1 to go to the next or previous keyboard layout.
-         */
-        public void switchKeyboardLayout(int deviceId, int direction);
-
         public void shutdown(boolean confirm);
         public void reboot(boolean confirm);
         public void rebootSafeMode(boolean confirm);
diff --git a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
index 8c4c0de..802f196 100644
--- a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
+++ b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
@@ -42,6 +42,7 @@
 import static com.android.server.wallpaper.WallpaperUtils.getWallpaperDir;
 import static com.android.server.wallpaper.WallpaperUtils.makeWallpaperIdLocked;
 import static com.android.window.flags.Flags.multiCrop;
+import static com.android.window.flags.Flags.offloadColorExtraction;
 
 import android.annotation.NonNull;
 import android.app.ActivityManager;
@@ -380,7 +381,7 @@
             }
 
             // Outside of the lock since it will synchronize itself
-            notifyWallpaperColorsChanged(wallpaper);
+            if (!offloadColorExtraction()) notifyWallpaperColorsChanged(wallpaper);
         }
 
         @Override
@@ -403,12 +404,16 @@
     }
 
     void notifyWallpaperColorsChanged(@NonNull WallpaperData wallpaper) {
+        notifyWallpaperColorsChanged(wallpaper, wallpaper.mWhich);
+    }
+
+    private void notifyWallpaperColorsChanged(@NonNull WallpaperData wallpaper, int which) {
         if (DEBUG) {
             Slog.i(TAG, "Notifying wallpaper colors changed");
         }
         if (wallpaper.connection != null) {
             wallpaper.connection.forEachDisplayConnector(connector ->
-                    notifyWallpaperColorsChangedOnDisplay(wallpaper, connector.mDisplayId));
+                    notifyWallpaperColorsChangedOnDisplay(wallpaper, connector.mDisplayId, which));
         }
     }
 
@@ -425,6 +430,11 @@
 
     private void notifyWallpaperColorsChangedOnDisplay(@NonNull WallpaperData wallpaper,
             int displayId) {
+        notifyWallpaperColorsChangedOnDisplay(wallpaper, displayId, wallpaper.mWhich);
+    }
+
+    private void notifyWallpaperColorsChangedOnDisplay(@NonNull WallpaperData wallpaper,
+            int displayId, int which) {
         boolean needsExtraction;
         synchronized (mLock) {
             final RemoteCallbackList<IWallpaperManagerCallback> currentUserColorListeners =
@@ -449,8 +459,8 @@
             notify = extractColors(wallpaper);
         }
         if (notify) {
-            notifyColorListeners(getAdjustedWallpaperColorsOnDimming(wallpaper),
-                    wallpaper.mWhich, wallpaper.userId, displayId);
+            notifyColorListeners(getAdjustedWallpaperColorsOnDimming(wallpaper), which,
+                    wallpaper.userId, displayId);
         }
     }
 
@@ -504,6 +514,7 @@
      * @return true unless the wallpaper changed during the color computation
      */
     private boolean extractColors(WallpaperData wallpaper) {
+        if (offloadColorExtraction()) return !mImageWallpaper.equals(wallpaper.wallpaperComponent);
         String cropFile = null;
         boolean defaultImageWallpaper = false;
         int wallpaperId;
@@ -1148,10 +1159,16 @@
             synchronized (mLock) {
                 // Do not broadcast changes on ImageWallpaper since it's handled
                 // internally by this class.
-                if (mImageWallpaper.equals(mWallpaper.wallpaperComponent)) {
+                boolean isImageWallpaper = mImageWallpaper.equals(mWallpaper.wallpaperComponent);
+                if (isImageWallpaper && (!offloadColorExtraction() || primaryColors == null)) {
                     return;
                 }
                 mWallpaper.primaryColors = primaryColors;
+                // only save the colors for ImageWallpaper - for live wallpapers, the colors
+                // are always recomputed after a reboot.
+                if (offloadColorExtraction() && isImageWallpaper) {
+                    saveSettingsLocked(mWallpaper.userId);
+                }
             }
             notifyWallpaperColorsChangedOnDisplay(mWallpaper, displayId);
         }
@@ -1177,7 +1194,9 @@
                 try {
                     // This will trigger onComputeColors in the wallpaper engine.
                     // It's fine to be locked in here since the binder is oneway.
-                    connector.mEngine.requestWallpaperColors();
+                    if (!offloadColorExtraction() || mWallpaper.primaryColors == null) {
+                        connector.mEngine.requestWallpaperColors();
+                    }
                 } catch (RemoteException e) {
                     Slog.w(TAG, "Failed to request wallpaper colors", e);
                 }
@@ -1811,6 +1830,7 @@
             // Offload color extraction to another thread since switchUser will be called
             // from the main thread.
             FgThread.getHandler().post(() -> {
+                if (offloadColorExtraction()) return;
                 notifyWallpaperColorsChanged(systemWallpaper);
                 if (lockWallpaper != systemWallpaper) notifyWallpaperColorsChanged(lockWallpaper);
                 notifyWallpaperColorsChanged(mFallbackWallpaper);
@@ -2722,8 +2742,10 @@
                         });
                         // Need to extract colors again to re-calculate dark hints after
                         // applying dimming.
-                        wp.mIsColorExtractedFromDim = true;
-                        pendingColorExtraction.add(wp);
+                        if (!offloadColorExtraction()) {
+                            wp.mIsColorExtractedFromDim = true;
+                            pendingColorExtraction.add(wp);
+                        }
                         changed = true;
                     }
                 }
@@ -2732,7 +2754,7 @@
                 }
             }
             for (WallpaperData wp: pendingColorExtraction) {
-                notifyWallpaperColorsChanged(wp);
+                if (!offloadColorExtraction()) notifyWallpaperColorsChanged(wp);
             }
         } finally {
             Binder.restoreCallingIdentity(ident);
@@ -2927,6 +2949,7 @@
                     }
                     wallpaper.allowBackup = allowBackup;
                     wallpaper.mWallpaperDimAmount = getWallpaperDimAmount();
+                    if (offloadColorExtraction()) wallpaper.primaryColors = null;
                 }
                 return pfd;
             } finally {
@@ -3069,6 +3092,10 @@
         checkPermission(android.Manifest.permission.SET_WALLPAPER_COMPONENT);
 
         boolean shouldNotifyColors = false;
+
+        // If the lockscreen wallpaper is set to the same as the home screen, notify that the
+        // lockscreen wallpaper colors changed, even if we don't bind a new wallpaper engine.
+        boolean shouldNotifyLockscreenColors = false;
         boolean bindSuccess;
         final WallpaperData newWallpaper;
 
@@ -3114,7 +3141,7 @@
                 bindSuccess = bindWallpaperComponentLocked(name, /* force */
                         forceRebind, /* fromUser */ true, newWallpaper, reply);
                 if (bindSuccess) {
-                    if (!same) {
+                    if (!same || (offloadColorExtraction() && forceRebind)) {
                         newWallpaper.primaryColors = null;
                     } else {
                         if (newWallpaper.connection != null) {
@@ -3138,6 +3165,11 @@
                     newWallpaper.wallpaperId = makeWallpaperIdLocked();
                     notifyCallbacksLocked(newWallpaper);
                     shouldNotifyColors = true;
+                    if (offloadColorExtraction()) {
+                        shouldNotifyColors = false;
+                        shouldNotifyLockscreenColors = !force && same && !systemIsBoth
+                                && which == (FLAG_SYSTEM | FLAG_LOCK);
+                    }
 
                     if (which == (FLAG_SYSTEM | FLAG_LOCK)) {
                         if (DEBUG) {
@@ -3166,6 +3198,10 @@
         if (shouldNotifyColors) {
             notifyWallpaperColorsChanged(newWallpaper);
         }
+        if (shouldNotifyLockscreenColors) {
+            notifyWallpaperColorsChanged(newWallpaper, FLAG_LOCK);
+        }
+
         return bindSuccess;
     }
 
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index b8bb258..0ad601d 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -61,6 +61,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;
 
@@ -225,13 +226,16 @@
             if (keyguardShowing) {
                 state.mDismissalRequested = false;
             }
-            if (goingAwayRemoved) {
-                // Keyguard dismiss is canceled. Send a transition to undo the changes and clean up
-                // before holding the sleep token again.
+            if (goingAwayRemoved || (keyguardShowing && Flags.keyguardAppearTransition())) {
+                // Keyguard decided to show or stopped going away. Send a transition to animate back
+                // to the locked state before holding the sleep token again
                 final DisplayContent dc = mRootWindowContainer.getDefaultDisplay();
                 dc.requestTransitionAndLegacyPrepare(
                         TRANSIT_TO_FRONT, TRANSIT_FLAG_KEYGUARD_APPEARING);
-                mWindowManager.executeAppTransition();
+                if (Flags.keyguardAppearTransition()) {
+                    dc.mWallpaperController.adjustWallpaperWindows();
+                }
+                dc.executeAppTransition();
             }
         }
 
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index c8cb934..59bda54 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -857,10 +857,6 @@
     }
 
     public void updateWallpaperTokens(boolean keyguardLocked) {
-        if (DEBUG_WALLPAPER) {
-            Slog.v(TAG, "Wallpaper vis: target " + mWallpaperTarget + " prev="
-                    + mPrevWallpaperTarget);
-        }
         updateWallpaperTokens(mWallpaperTarget != null || mPrevWallpaperTarget != null,
                 keyguardLocked);
     }
@@ -869,6 +865,8 @@
      * Change the visibility of the top wallpaper to {@param visibility} and hide all the others.
      */
     private void updateWallpaperTokens(boolean visibility, boolean keyguardLocked) {
+        ProtoLog.v(WM_DEBUG_WALLPAPER, "updateWallpaperTokens requestedVisibility=%b on"
+                + " keyguardLocked=%b", visibility, keyguardLocked);
         WindowState topWallpaper = mFindResults.getTopWallpaper(keyguardLocked);
         WallpaperWindowToken topWallpaperToken =
                 topWallpaper == null ? null : topWallpaper.mToken.asWallpaperToken();
diff --git a/services/core/java/com/android/server/wm/WallpaperWindowToken.java b/services/core/java/com/android/server/wm/WallpaperWindowToken.java
index 55eeaf2..5c24eee 100644
--- a/services/core/java/com/android/server/wm/WallpaperWindowToken.java
+++ b/services/core/java/com/android/server/wm/WallpaperWindowToken.java
@@ -274,6 +274,12 @@
     }
 
     @Override
+    boolean isSyncFinished(BLASTSyncEngine.SyncGroup group) {
+        // TODO(b/233286785): Support sync state for wallpaper. See WindowState#prepareSync.
+        return !mVisibleRequested || !hasVisibleNotDrawnWallpaper();
+    }
+
+    @Override
     public String toString() {
         if (stringName == null) {
             StringBuilder sb = new StringBuilder();
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index aeefc34..8a68afb 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -3685,12 +3685,6 @@
 
     // Called by window manager policy.  Not exposed externally.
     @Override
-    public void switchKeyboardLayout(int deviceId, int direction) {
-        mInputManager.switchKeyboardLayout(deviceId, direction);
-    }
-
-    // Called by window manager policy.  Not exposed externally.
-    @Override
     public void shutdown(boolean confirm) {
         // Pass in the UI context, since ShutdownThread requires it (to show UI).
         ShutdownThread.shutdown(ActivityThread.currentActivityThread().getSystemUiContext(),
diff --git a/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java b/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java
index 52df010..eac9929 100644
--- a/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java
+++ b/services/tests/wmtests/src/com/android/server/policy/TestPhoneWindowManager.java
@@ -85,7 +85,6 @@
 import android.os.test.TestLooper;
 import android.service.dreams.DreamManagerInternal;
 import android.telecom.TelecomManager;
-import android.util.FeatureFlagUtils;
 import android.view.Display;
 import android.view.InputDevice;
 import android.view.KeyEvent;
@@ -743,15 +742,8 @@
 
     void assertSwitchKeyboardLayout(int direction, int displayId) {
         mTestLooper.dispatchAll();
-        if (FeatureFlagUtils.isEnabled(mContext, FeatureFlagUtils.SETTINGS_NEW_KEYBOARD_UI)) {
-            verify(mInputMethodManagerInternal).onSwitchKeyboardLayoutShortcut(eq(direction),
-                    eq(displayId), eq(mImeTargetWindowToken));
-            verify(mWindowManagerFuncsImpl, never()).switchKeyboardLayout(anyInt(), anyInt());
-        } else {
-            verify(mWindowManagerFuncsImpl).switchKeyboardLayout(anyInt(), eq(direction));
-            verify(mInputMethodManagerInternal, never())
-                    .onSwitchKeyboardLayoutShortcut(anyInt(), anyInt(), any());
-        }
+        verify(mInputMethodManagerInternal).onSwitchKeyboardLayoutShortcut(eq(direction),
+                eq(displayId), eq(mImeTargetWindowToken));
     }
 
     void assertTakeBugreport() {
diff --git a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt
index 511c948..1390218 100644
--- a/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt
+++ b/tests/FlickerTests/ActivityEmbedding/src/com/android/server/wm/flicker/activityembedding/rotation/RotationTransition.kt
@@ -17,8 +17,12 @@
 package com.android.server.wm.flicker.activityembedding.rotation
 
 import android.platform.test.annotations.Presubmit
+import android.tools.Position
+import android.tools.datatypes.Rect
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.traces.Condition
+import android.tools.traces.DeviceStateDump
 import android.tools.traces.component.ComponentNameMatcher
 import com.android.server.wm.flicker.activityembedding.ActivityEmbeddingTestBase
 import com.android.server.wm.flicker.helpers.setRotation
@@ -30,7 +34,14 @@
     override val transition: FlickerBuilder.() -> Unit = {
         setup { this.setRotation(flicker.scenario.startRotation) }
         teardown { testApp.exit(wmHelper) }
-        transitions { this.setRotation(flicker.scenario.endRotation) }
+        transitions {
+            this.setRotation(flicker.scenario.endRotation)
+            if (!flicker.scenario.isTablet) {
+                wmHelper.StateSyncBuilder()
+                    .add(navBarInPosition(flicker.scenario.isGesturalNavigation))
+                    .waitForAndVerify()
+            }
+        }
     }
 
     /** {@inheritDoc} */
@@ -76,4 +87,37 @@
         appLayerRotates_StartingPos()
         appLayerRotates_EndingPos()
     }
+
+    private fun navBarInPosition(isGesturalNavigation: Boolean): Condition<DeviceStateDump> {
+        return Condition("navBarPosition") { dump ->
+            val display =
+                dump.layerState.displays.filterNot { it.isOff }.minByOrNull { it.id }
+                    ?: error("There is no display!")
+            val displayArea = display.layerStackSpace
+            val navBarPosition = display.navBarPosition(isGesturalNavigation)
+            val navBarRegion = dump.layerState
+                .getLayerWithBuffer(ComponentNameMatcher.NAV_BAR)
+                ?.visibleRegion?.bounds ?: Rect.EMPTY
+
+            when (navBarPosition) {
+                Position.TOP ->
+                    navBarRegion.top == displayArea.top &&
+                        navBarRegion.left == displayArea.left &&
+                        navBarRegion.right == displayArea.right
+                Position.BOTTOM ->
+                    navBarRegion.bottom == displayArea.bottom &&
+                        navBarRegion.left == displayArea.left &&
+                        navBarRegion.right == displayArea.right
+                Position.LEFT ->
+                    navBarRegion.left == displayArea.left &&
+                        navBarRegion.top == displayArea.top &&
+                        navBarRegion.bottom == displayArea.bottom
+                Position.RIGHT ->
+                    navBarRegion.right == displayArea.right &&
+                        navBarRegion.top == displayArea.top &&
+                        navBarRegion.bottom == displayArea.bottom
+                else -> error("Unknown position $navBarPosition")
+            }
+        }
+    }
 }
diff --git a/tests/FlickerTests/README.md b/tests/FlickerTests/README.md
index 6b28fdf..7429250 100644
--- a/tests/FlickerTests/README.md
+++ b/tests/FlickerTests/README.md
@@ -7,82 +7,17 @@
 
 ## Adding a Test
 
-By default tests should inherit from `RotationTestBase` or `NonRotationTestBase` and must override the variable `transitionToRun` (Kotlin) or the function `getTransitionToRun()` (Java).
-Only tests that are not supported by these classes should inherit directly from the `FlickerTestBase` class.
+By default, tests should inherit from `TestBase` and override the variable `transition` (Kotlin) or the function `getTransition()` (Java).
 
-### Rotation animations and transitions
+Inheriting from this class ensures the common assertions will be executed, namely:
 
-Tests that rotate the device should inherit from `RotationTestBase`.
-Tests that inherit from the class automatically receive start and end rotation values.  
-Moreover, these tests inherit the following checks:
 * all regions on the screen are covered
 * status bar is always visible
-* status bar rotates
+* status bar is at the correct position at the start and end of the transition
 * nav bar is always visible
-* nav bar is rotates
+* nav bar is at the correct position at the start and end of the transition
 
 The default tests can be disabled by overriding the respective methods and including an `@Ignore` annotation.
 
-### Non-Rotation animations and transitions
+For more examples of how a test looks like check `ChangeAppRotationTest` within the `Rotation` subdirectory.
 
-`NonRotationTestBase` was created to make it easier to write tests that do not involve rotation (e.g., `Pip`, `split screen` or `IME`).
-Tests that inherit from the class are automatically executed twice: once in portrait and once in landscape mode and the assertions are checked independently.
-Moreover, these tests inherit the following checks:
-* all regions on the screen are covered
-* status bar is always visible
-* nav bar is always visible
-
-The default tests can be disabled by overriding the respective methods and including an `@Ignore` annotation.
-
-### Exceptional cases
-
-Tests that rotate the device should inherit from `RotationTestBase`.
-This class allows the test to be freely configured and does not provide any assertions.  
-
-
-### Example
-
-Start by defining common or error prone transitions using `TransitionRunner`.
-```kotlin
-@LargeTest
-@RunWith(Parameterized::class)
-@FixMethodOrder(MethodSorters.NAME_ASCENDING)
-class MyTest(
-    beginRotationName: String,
-    beginRotation: Int
-) : NonRotationTestBase(beginRotationName, beginRotation) {
-    init {
-        mTestApp = MyAppHelper(InstrumentationRegistry.getInstrumentation())
-    }
-
-    override val transitionToRun: TransitionRunner
-        get() = TransitionRunner.newBuilder()
-            .withTag("myTest")
-            .recordAllRuns()
-            .runBefore { device.pressHome() }
-            .runBefore { device.waitForIdle() }
-            .run { testApp.open() }
-            .runAfter{ testApp.exit() }
-            .repeat(2)
-            .includeJankyRuns()
-            .build()
-
-    @Test
-    fun myWMTest() {
-        checkResults {
-            WmTraceSubject.assertThat(it)
-                    .showsAppWindow(MyTestApp)
-                    .forAllEntries()
-        }
-    }
-
-    @Test
-    fun mySFTest() {
-        checkResults {
-            LayersTraceSubject.assertThat(it)
-                    .showsLayer(MyTestApp)
-                    .forAllEntries()
-        }
-    }
-}
-```
diff --git a/tests/Input/src/com/android/server/input/KeyboardLayoutManagerTests.kt b/tests/Input/src/com/android/server/input/KeyboardLayoutManagerTests.kt
index e60764f..80282c3 100644
--- a/tests/Input/src/com/android/server/input/KeyboardLayoutManagerTests.kt
+++ b/tests/Input/src/com/android/server/input/KeyboardLayoutManagerTests.kt
@@ -33,7 +33,6 @@
 import android.os.Bundle
 import android.os.test.TestLooper
 import android.platform.test.annotations.Presubmit
-import android.provider.Settings
 import android.util.proto.ProtoOutputStream
 import android.view.InputDevice
 import android.view.inputmethod.InputMethodInfo
@@ -47,9 +46,7 @@
 import org.junit.After
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotEquals
-import org.junit.Assert.assertNull
 import org.junit.Assert.assertTrue
-import org.junit.Assert.assertThrows
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -234,631 +231,326 @@
     }
 
     @Test
-    fun testDefaultUi_getKeyboardLayouts() {
-        NewSettingsApiFlag(false).use {
-            val keyboardLayouts = keyboardLayoutManager.keyboardLayouts
-            assertNotEquals(
-                "Default UI: Keyboard layout API should not return empty array",
-                0,
-                keyboardLayouts.size
-            )
-            assertTrue(
-                "Default UI: Keyboard layout API should provide English(US) layout",
-                hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
-            )
-        }
+    fun testGetKeyboardLayouts() {
+        val keyboardLayouts = keyboardLayoutManager.keyboardLayouts
+        assertNotEquals(
+            "Keyboard layout API should not return empty array",
+            0,
+            keyboardLayouts.size
+        )
+        assertTrue(
+            "Keyboard layout API should provide English(US) layout",
+            hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
+        )
     }
 
     @Test
-    fun testNewUi_getKeyboardLayouts() {
-        NewSettingsApiFlag(true).use {
-            val keyboardLayouts = keyboardLayoutManager.keyboardLayouts
-            assertNotEquals(
-                "New UI: Keyboard layout API should not return empty array",
-                0,
-                keyboardLayouts.size
-            )
-            assertTrue(
-                "New UI: Keyboard layout API should provide English(US) layout",
-                hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
-            )
-        }
+    fun testGetKeyboardLayout() {
+        val keyboardLayout =
+            keyboardLayoutManager.getKeyboardLayout(ENGLISH_US_LAYOUT_DESCRIPTOR)
+        assertEquals("getKeyboardLayout API should return correct Layout from " +
+                "available layouts",
+            ENGLISH_US_LAYOUT_DESCRIPTOR,
+            keyboardLayout!!.descriptor
+        )
     }
 
     @Test
-    fun testDefaultUi_getKeyboardLayoutsForInputDevice() {
-        NewSettingsApiFlag(false).use {
-            val keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutsForInputDevice(keyboardDevice.identifier)
-            assertNotEquals(
-                "Default UI: getKeyboardLayoutsForInputDevice API should not return empty array",
-                0,
-                keyboardLayouts.size
-            )
-            assertTrue(
-                "Default UI: getKeyboardLayoutsForInputDevice API should provide English(US) " +
-                        "layout",
-                hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
-            )
+    fun testGetSetKeyboardLayoutForInputDevice_withImeInfo() {
+        val imeSubtype = createImeSubtype()
 
-            val vendorSpecificKeyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutsForInputDevice(
-                    vendorSpecificKeyboardDevice.identifier
-                )
-            assertEquals(
-                "Default UI: getKeyboardLayoutsForInputDevice API should return only vendor " +
-                        "specific layout",
+        keyboardLayoutManager.setKeyboardLayoutForInputDevice(
+            keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
+            ENGLISH_UK_LAYOUT_DESCRIPTOR
+        )
+        var result =
+            keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
+            )
+        assertEquals(
+            "getKeyboardLayoutForInputDevice API should return the set layout",
+            ENGLISH_UK_LAYOUT_DESCRIPTOR,
+            result.layoutDescriptor
+        )
+
+        // This should replace previously set layout
+        keyboardLayoutManager.setKeyboardLayoutForInputDevice(
+            keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
+            ENGLISH_US_LAYOUT_DESCRIPTOR
+        )
+        result =
+            keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
+            )
+        assertEquals(
+            "getKeyboardLayoutForInputDevice API should return the last set layout",
+            ENGLISH_US_LAYOUT_DESCRIPTOR,
+            result.layoutDescriptor
+        )
+    }
+
+    @Test
+    fun testGetKeyboardLayoutListForInputDevice() {
+        // Check Layouts for "hi-Latn". It should return all 'Latn' keyboard layouts
+        var keyboardLayouts =
+            keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo,
+                createImeSubtypeForLanguageTag("hi-Latn")
+            )
+        assertNotEquals(
+            "getKeyboardLayoutListForInputDevice API should return the list of " +
+                    "supported layouts with matching script code",
+            0,
+            keyboardLayouts.size
+        )
+        assertTrue("getKeyboardLayoutListForInputDevice API should return a list " +
+                "containing English(US) layout for hi-Latn",
+            containsLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
+        )
+        assertTrue("getKeyboardLayoutListForInputDevice API should return a list " +
+                "containing English(No script code) layout for hi-Latn",
+            containsLayout(
+                keyboardLayouts,
+                createLayoutDescriptor("keyboard_layout_english_without_script_code")
+            )
+        )
+
+        // Check Layouts for "hi" which by default uses 'Deva' script.
+        keyboardLayouts =
+            keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo,
+                createImeSubtypeForLanguageTag("hi")
+            )
+        assertEquals("getKeyboardLayoutListForInputDevice API should return empty " +
+                "list if no supported layouts available",
+            0,
+            keyboardLayouts.size
+        )
+
+        // If user manually selected some layout, always provide it in the layout list
+        val imeSubtype = createImeSubtypeForLanguageTag("hi")
+        keyboardLayoutManager.setKeyboardLayoutForInputDevice(
+            keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
+            ENGLISH_US_LAYOUT_DESCRIPTOR
+        )
+        keyboardLayouts =
+            keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo,
+                imeSubtype
+            )
+        assertEquals("getKeyboardLayoutListForInputDevice API should return user " +
+                "selected layout even if the script is incompatible with IME",
                 1,
-                vendorSpecificKeyboardLayouts.size
-            )
-            assertEquals(
-                "Default UI: getKeyboardLayoutsForInputDevice API should return vendor specific " +
-                        "layout",
-                VENDOR_SPECIFIC_LAYOUT_DESCRIPTOR,
-                vendorSpecificKeyboardLayouts[0].descriptor
-            )
-        }
-    }
+            keyboardLayouts.size
+        )
 
-    @Test
-    fun testNewUi_getKeyboardLayoutsForInputDevice() {
-        NewSettingsApiFlag(true).use {
-            val keyboardLayouts = keyboardLayoutManager.keyboardLayouts
-            assertNotEquals(
-                    "New UI: getKeyboardLayoutsForInputDevice API should not return empty array",
-                    0,
-                    keyboardLayouts.size
-            )
-            assertTrue(
-                    "New UI: getKeyboardLayoutsForInputDevice API should provide English(US) " +
-                            "layout",
-                    hasLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
-            )
-        }
-    }
-
-    @Test
-    fun testDefaultUi_getSetCurrentKeyboardLayoutForInputDevice() {
-        NewSettingsApiFlag(false).use {
-            assertNull(
-                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return null if " +
-                        "nothing was set",
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-
-            keyboardLayoutManager.setCurrentKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier,
-                ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            val keyboardLayout =
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            assertEquals(
-                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return the set " +
-                        "layout",
-                ENGLISH_US_LAYOUT_DESCRIPTOR,
-                keyboardLayout
-            )
-        }
-    }
-
-    @Test
-    fun testNewUi_getSetCurrentKeyboardLayoutForInputDevice() {
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.setCurrentKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier,
-                ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            assertNull(
-                "New UI: getCurrentKeyboardLayoutForInputDevice API should always return null " +
-                        "even after setCurrentKeyboardLayoutForInputDevice",
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-        }
-    }
-
-    @Test
-    fun testDefaultUi_getEnabledKeyboardLayoutsForInputDevice() {
-        NewSettingsApiFlag(false).use {
-            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-
-            val keyboardLayouts =
-                keyboardLayoutManager.getEnabledKeyboardLayoutsForInputDevice(
-                    keyboardDevice.identifier
-                )
-            assertEquals(
-                "Default UI: getEnabledKeyboardLayoutsForInputDevice API should return added " +
-                        "layout",
-                1,
-                keyboardLayouts.size
-            )
-            assertEquals(
-                "Default UI: getEnabledKeyboardLayoutsForInputDevice API should return " +
-                        "English(US) layout",
-                ENGLISH_US_LAYOUT_DESCRIPTOR,
-                keyboardLayouts[0]
-            )
-            assertEquals(
-                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return " +
-                        "English(US) layout (Auto select the first enabled layout)",
-                ENGLISH_US_LAYOUT_DESCRIPTOR,
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-
-            keyboardLayoutManager.removeKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            assertEquals(
-                "Default UI: getKeyboardLayoutsForInputDevice API should return 0 layouts",
-                0,
-                keyboardLayoutManager.getEnabledKeyboardLayoutsForInputDevice(
-                    keyboardDevice.identifier
-                ).size
-            )
-            assertNull(
-                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return null after " +
-                        "the enabled layout is removed",
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-        }
-    }
-
-    @Test
-    fun testNewUi_getEnabledKeyboardLayoutsForInputDevice() {
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-
-            assertEquals(
-                "New UI: getEnabledKeyboardLayoutsForInputDevice API should return always return " +
-                        "an empty array",
-                0,
-                keyboardLayoutManager.getEnabledKeyboardLayoutsForInputDevice(
-                    keyboardDevice.identifier
-                ).size
-            )
-            assertNull(
-                "New UI: getCurrentKeyboardLayoutForInputDevice API should always return null",
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-        }
-    }
-
-    @Test
-    fun testDefaultUi_switchKeyboardLayout() {
-        NewSettingsApiFlag(false).use {
-            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, ENGLISH_UK_LAYOUT_DESCRIPTOR
-            )
-            assertEquals(
-                "Default UI: getCurrentKeyboardLayoutForInputDevice API should return " +
-                        "English(US) layout",
-                ENGLISH_US_LAYOUT_DESCRIPTOR,
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-
-            keyboardLayoutManager.switchKeyboardLayout(DEVICE_ID, 1)
-
-            // Throws null pointer because trying to show toast using TestLooper
-            assertThrows(NullPointerException::class.java) { testLooper.dispatchAll() }
-            assertEquals("Default UI: getCurrentKeyboardLayoutForInputDevice API should return " +
-                    "English(UK) layout",
-                ENGLISH_UK_LAYOUT_DESCRIPTOR,
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-        }
-    }
-
-    @Test
-    fun testNewUi_switchKeyboardLayout() {
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            keyboardLayoutManager.addKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, ENGLISH_UK_LAYOUT_DESCRIPTOR
-            )
-
-            keyboardLayoutManager.switchKeyboardLayout(DEVICE_ID, 1)
-            testLooper.dispatchAll()
-
-            assertNull("New UI: getCurrentKeyboardLayoutForInputDevice API should always return " +
-                    "null",
-                keyboardLayoutManager.getCurrentKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier
-                )
-            )
-        }
-    }
-
-    @Test
-    fun testDefaultUi_getKeyboardLayout() {
-        NewSettingsApiFlag(false).use {
-            val keyboardLayout =
-                keyboardLayoutManager.getKeyboardLayout(ENGLISH_US_LAYOUT_DESCRIPTOR)
-            assertEquals("Default UI: getKeyboardLayout API should return correct Layout from " +
-                    "available layouts",
-                ENGLISH_US_LAYOUT_DESCRIPTOR,
-                keyboardLayout!!.descriptor
-            )
-        }
-    }
-
-    @Test
-    fun testNewUi_getKeyboardLayout() {
-        NewSettingsApiFlag(true).use {
-            val keyboardLayout =
-                keyboardLayoutManager.getKeyboardLayout(ENGLISH_US_LAYOUT_DESCRIPTOR)
-            assertEquals("New UI: getKeyboardLayout API should return correct Layout from " +
-                    "available layouts",
-                ENGLISH_US_LAYOUT_DESCRIPTOR,
-                keyboardLayout!!.descriptor
-            )
-        }
-    }
-
-    @Test
-    fun testDefaultUi_getSetKeyboardLayoutForInputDevice_WithImeInfo() {
-        NewSettingsApiFlag(false).use {
-            val imeSubtype = createImeSubtype()
-            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
-                ENGLISH_UK_LAYOUT_DESCRIPTOR
-            )
-            assertEquals(
-                "Default UI: getKeyboardLayoutForInputDevice API should always return " +
-                        "KeyboardLayoutSelectionResult.FAILED",
-                KeyboardLayoutSelectionResult.FAILED,
-                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
-                )
-            )
-        }
-    }
-
-    @Test
-    fun testNewUi_getSetKeyboardLayoutForInputDevice_withImeInfo() {
-        NewSettingsApiFlag(true).use {
-            val imeSubtype = createImeSubtype()
-
-            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
-                ENGLISH_UK_LAYOUT_DESCRIPTOR
-            )
-            var result =
-                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
-                )
-            assertEquals(
-                "New UI: getKeyboardLayoutForInputDevice API should return the set layout",
-                ENGLISH_UK_LAYOUT_DESCRIPTOR,
-                result.layoutDescriptor
-            )
-
-            // This should replace previously set layout
-            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
-                ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            result =
-                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype
-                )
-            assertEquals(
-                "New UI: getKeyboardLayoutForInputDevice API should return the last set layout",
-                ENGLISH_US_LAYOUT_DESCRIPTOR,
-                result.layoutDescriptor
-            )
-        }
-    }
-
-    @Test
-    fun testDefaultUi_getKeyboardLayoutListForInputDevice() {
-        NewSettingsApiFlag(false).use {
-            val keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+        // Special case Japanese: UScript ignores provided script code for certain language tags
+        // Should manually match provided script codes and then rely on Uscript to derive
+        // script from language tags and match those.
+        keyboardLayouts =
+            keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
                     keyboardDevice.identifier, USER_ID, imeInfo,
-                    createImeSubtype()
-                )
-            assertEquals("Default UI: getKeyboardLayoutListForInputDevice API should always " +
-                    "return empty array",
-                0,
-                keyboardLayouts.size
+                    createImeSubtypeForLanguageTag("ja-Latn-JP")
             )
-        }
-    }
+        assertNotEquals(
+            "getKeyboardLayoutListForInputDevice API should return the list of " +
+                    "supported layouts with matching script code for ja-Latn-JP",
+            0,
+            keyboardLayouts.size
+        )
+        assertTrue("getKeyboardLayoutListForInputDevice API should return a list " +
+                "containing English(US) layout for ja-Latn-JP",
+            containsLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
+        )
+        assertTrue("getKeyboardLayoutListForInputDevice API should return a list " +
+                "containing English(No script code) layout for ja-Latn-JP",
+            containsLayout(
+                keyboardLayouts,
+                createLayoutDescriptor("keyboard_layout_english_without_script_code")
+            )
+        )
 
-    @Test
-    fun testNewUi_getKeyboardLayoutListForInputDevice() {
-        NewSettingsApiFlag(true).use {
-            // Check Layouts for "hi-Latn". It should return all 'Latn' keyboard layouts
-            var keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+        // If script code not explicitly provided for Japanese should rely on Uscript to find
+        // derived script code and hence no suitable layout will be found.
+        keyboardLayouts =
+            keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
                     keyboardDevice.identifier, USER_ID, imeInfo,
-                    createImeSubtypeForLanguageTag("hi-Latn")
-                )
-            assertNotEquals(
-                "New UI: getKeyboardLayoutListForInputDevice API should return the list of " +
-                        "supported layouts with matching script code",
-                0,
-                keyboardLayouts.size
+                    createImeSubtypeForLanguageTag("ja-JP")
             )
-            assertTrue("New UI: getKeyboardLayoutListForInputDevice API should return a list " +
-                    "containing English(US) layout for hi-Latn",
-                containsLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
-            )
-            assertTrue("New UI: getKeyboardLayoutListForInputDevice API should return a list " +
-                    "containing English(No script code) layout for hi-Latn",
-                containsLayout(
-                    keyboardLayouts,
-                    createLayoutDescriptor("keyboard_layout_english_without_script_code")
-                )
-            )
+        assertEquals(
+            "getKeyboardLayoutListForInputDevice API should return empty list of " +
+                    "supported layouts with matching script code for ja-JP",
+            0,
+            keyboardLayouts.size
+        )
 
-            // Check Layouts for "hi" which by default uses 'Deva' script.
-            keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo,
-                    createImeSubtypeForLanguageTag("hi")
-                )
-            assertEquals("New UI: getKeyboardLayoutListForInputDevice API should return empty " +
-                    "list if no supported layouts available",
-                0,
-                keyboardLayouts.size
+        // If IME doesn't have a corresponding language tag, then should show all available
+        // layouts no matter the script code.
+        keyboardLayouts =
+            keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo, null
             )
-
-            // If user manually selected some layout, always provide it in the layout list
-            val imeSubtype = createImeSubtypeForLanguageTag("hi")
-            keyboardLayoutManager.setKeyboardLayoutForInputDevice(
-                keyboardDevice.identifier, USER_ID, imeInfo, imeSubtype,
-                ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo,
-                    imeSubtype
-                )
-            assertEquals("New UI: getKeyboardLayoutListForInputDevice API should return user " +
-                    "selected layout even if the script is incompatible with IME",
-                    1,
-                keyboardLayouts.size
-            )
-
-            // Special case Japanese: UScript ignores provided script code for certain language tags
-            // Should manually match provided script codes and then rely on Uscript to derive
-            // script from language tags and match those.
-            keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
-                        keyboardDevice.identifier, USER_ID, imeInfo,
-                        createImeSubtypeForLanguageTag("ja-Latn-JP")
-                )
-            assertNotEquals(
-                "New UI: getKeyboardLayoutListForInputDevice API should return the list of " +
-                        "supported layouts with matching script code for ja-Latn-JP",
-                0,
-                keyboardLayouts.size
-            )
-            assertTrue("New UI: getKeyboardLayoutListForInputDevice API should return a list " +
-                    "containing English(US) layout for ja-Latn-JP",
-                containsLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
-            )
-            assertTrue("New UI: getKeyboardLayoutListForInputDevice API should return a list " +
-                    "containing English(No script code) layout for ja-Latn-JP",
-                containsLayout(
-                    keyboardLayouts,
-                    createLayoutDescriptor("keyboard_layout_english_without_script_code")
-                )
-            )
-
-            // If script code not explicitly provided for Japanese should rely on Uscript to find
-            // derived script code and hence no suitable layout will be found.
-            keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
-                        keyboardDevice.identifier, USER_ID, imeInfo,
-                        createImeSubtypeForLanguageTag("ja-JP")
-                )
-            assertEquals(
-                "New UI: getKeyboardLayoutListForInputDevice API should return empty list of " +
-                        "supported layouts with matching script code for ja-JP",
-                0,
-                keyboardLayouts.size
-            )
-
-            // If IME doesn't have a corresponding language tag, then should show all available
-            // layouts no matter the script code.
-            keyboardLayouts =
-                keyboardLayoutManager.getKeyboardLayoutListForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo, null
-                )
-            assertNotEquals(
-                "New UI: getKeyboardLayoutListForInputDevice API should return all layouts if" +
-                    "language tag or subtype not provided",
-                0,
-                keyboardLayouts.size
-            )
-            assertTrue("New UI: getKeyboardLayoutListForInputDevice API should contain Latin " +
-                "layouts if language tag or subtype not provided",
-                containsLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
-            )
-            assertTrue("New UI: getKeyboardLayoutListForInputDevice API should contain Cyrillic " +
-                "layouts if language tag or subtype not provided",
-                containsLayout(
-                    keyboardLayouts,
-                    createLayoutDescriptor("keyboard_layout_russian")
-                )
-            )
-        }
-    }
-
-    @Test
-    fun testNewUi_getDefaultKeyboardLayoutForInputDevice_withImeLanguageTag() {
-        NewSettingsApiFlag(true).use {
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTag("en-US"),
-                ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTag("en-GB"),
-                ENGLISH_UK_LAYOUT_DESCRIPTOR
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTag("de"),
-                GERMAN_LAYOUT_DESCRIPTOR
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTag("fr-FR"),
-                createLayoutDescriptor("keyboard_layout_french")
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTag("ru"),
+        assertNotEquals(
+            "getKeyboardLayoutListForInputDevice API should return all layouts if" +
+                "language tag or subtype not provided",
+            0,
+            keyboardLayouts.size
+        )
+        assertTrue("getKeyboardLayoutListForInputDevice API should contain Latin " +
+            "layouts if language tag or subtype not provided",
+            containsLayout(keyboardLayouts, ENGLISH_US_LAYOUT_DESCRIPTOR)
+        )
+        assertTrue("getKeyboardLayoutListForInputDevice API should contain Cyrillic " +
+            "layouts if language tag or subtype not provided",
+            containsLayout(
+                keyboardLayouts,
                 createLayoutDescriptor("keyboard_layout_russian")
             )
-            assertEquals(
-                "New UI: getDefaultKeyboardLayoutForInputDevice should return " +
-                        "KeyboardLayoutSelectionResult.FAILED when no layout available",
-                KeyboardLayoutSelectionResult.FAILED,
-                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo,
-                    createImeSubtypeForLanguageTag("it")
-                )
-            )
-            assertEquals(
-                "New UI: getDefaultKeyboardLayoutForInputDevice should return " +
-                        "KeyboardLayoutSelectionResult.FAILED when no layout for script code is" +
-                        "available",
-                KeyboardLayoutSelectionResult.FAILED,
-                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo,
-                    createImeSubtypeForLanguageTag("en-Deva")
-                )
-            )
-        }
+        )
     }
 
     @Test
-    fun testNewUi_getDefaultKeyboardLayoutForInputDevice_withImeLanguageTagAndLayoutType() {
-        NewSettingsApiFlag(true).use {
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("en-US", "qwerty"),
-                ENGLISH_US_LAYOUT_DESCRIPTOR
+    fun testGetDefaultKeyboardLayoutForInputDevice_withImeLanguageTag() {
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTag("en-US"),
+            ENGLISH_US_LAYOUT_DESCRIPTOR
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTag("en-GB"),
+            ENGLISH_UK_LAYOUT_DESCRIPTOR
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTag("de"),
+            GERMAN_LAYOUT_DESCRIPTOR
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTag("fr-FR"),
+            createLayoutDescriptor("keyboard_layout_french")
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTag("ru"),
+            createLayoutDescriptor("keyboard_layout_russian")
+        )
+        assertEquals(
+            "getDefaultKeyboardLayoutForInputDevice should return " +
+                    "KeyboardLayoutSelectionResult.FAILED when no layout available",
+            KeyboardLayoutSelectionResult.FAILED,
+            keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo,
+                createImeSubtypeForLanguageTag("it")
             )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("en-US", "dvorak"),
-                createLayoutDescriptor("keyboard_layout_english_us_dvorak")
-            )
-            // Try to match layout type even if country doesn't match
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("en-GB", "dvorak"),
-                createLayoutDescriptor("keyboard_layout_english_us_dvorak")
-            )
-            // Choose layout based on layout type priority, if layout type is not provided by IME
-            // (Qwerty > Dvorak > Extended)
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("en-US", ""),
-                ENGLISH_US_LAYOUT_DESCRIPTOR
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("en-GB", "qwerty"),
-                ENGLISH_UK_LAYOUT_DESCRIPTOR
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("de", "qwertz"),
-                GERMAN_LAYOUT_DESCRIPTOR
-            )
-            // Wrong layout type should match with language if provided layout type not available
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("de", "qwerty"),
-                GERMAN_LAYOUT_DESCRIPTOR
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("fr-FR", "azerty"),
-                createLayoutDescriptor("keyboard_layout_french")
-            )
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("ru", "qwerty"),
-                createLayoutDescriptor("keyboard_layout_russian_qwerty")
-            )
-            // If layout type is empty then prioritize KCM with empty layout type
-            assertCorrectLayout(
-                keyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("ru", ""),
-                createLayoutDescriptor("keyboard_layout_russian")
-            )
-            assertEquals("New UI: getDefaultKeyboardLayoutForInputDevice should return " +
+        )
+        assertEquals(
+            "getDefaultKeyboardLayoutForInputDevice should return " +
                     "KeyboardLayoutSelectionResult.FAILED when no layout for script code is" +
                     "available",
-                KeyboardLayoutSelectionResult.FAILED,
-                keyboardLayoutManager.getKeyboardLayoutForInputDevice(
-                    keyboardDevice.identifier, USER_ID, imeInfo,
-                    createImeSubtypeForLanguageTagAndLayoutType("en-Deva-US", "")
-                )
+            KeyboardLayoutSelectionResult.FAILED,
+            keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo,
+                createImeSubtypeForLanguageTag("en-Deva")
             )
-        }
+        )
     }
 
     @Test
-    fun testNewUi_getDefaultKeyboardLayoutForInputDevice_withHwLanguageTagAndLayoutType() {
-        NewSettingsApiFlag(true).use {
-            val frenchSubtype = createImeSubtypeForLanguageTagAndLayoutType("fr", "azerty")
-            // Should return English dvorak even if IME current layout is French, since HW says the
-            // keyboard is a Dvorak keyboard
-            assertCorrectLayout(
-                englishDvorakKeyboardDevice,
-                frenchSubtype,
-                createLayoutDescriptor("keyboard_layout_english_us_dvorak")
+    fun testGetDefaultKeyboardLayoutForInputDevice_withImeLanguageTagAndLayoutType() {
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("en-US", "qwerty"),
+            ENGLISH_US_LAYOUT_DESCRIPTOR
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("en-US", "dvorak"),
+            createLayoutDescriptor("keyboard_layout_english_us_dvorak")
+        )
+        // Try to match layout type even if country doesn't match
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("en-GB", "dvorak"),
+            createLayoutDescriptor("keyboard_layout_english_us_dvorak")
+        )
+        // Choose layout based on layout type priority, if layout type is not provided by IME
+        // (Qwerty > Dvorak > Extended)
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("en-US", ""),
+            ENGLISH_US_LAYOUT_DESCRIPTOR
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("en-GB", "qwerty"),
+            ENGLISH_UK_LAYOUT_DESCRIPTOR
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("de", "qwertz"),
+            GERMAN_LAYOUT_DESCRIPTOR
+        )
+        // Wrong layout type should match with language if provided layout type not available
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("de", "qwerty"),
+            GERMAN_LAYOUT_DESCRIPTOR
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("fr-FR", "azerty"),
+            createLayoutDescriptor("keyboard_layout_french")
+        )
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("ru", "qwerty"),
+            createLayoutDescriptor("keyboard_layout_russian_qwerty")
+        )
+        // If layout type is empty then prioritize KCM with empty layout type
+        assertCorrectLayout(
+            keyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("ru", ""),
+            createLayoutDescriptor("keyboard_layout_russian")
+        )
+        assertEquals("getDefaultKeyboardLayoutForInputDevice should return " +
+                "KeyboardLayoutSelectionResult.FAILED when no layout for script code is" +
+                "available",
+            KeyboardLayoutSelectionResult.FAILED,
+            keyboardLayoutManager.getKeyboardLayoutForInputDevice(
+                keyboardDevice.identifier, USER_ID, imeInfo,
+                createImeSubtypeForLanguageTagAndLayoutType("en-Deva-US", "")
             )
+        )
+    }
 
-            // Back to back changing HW keyboards with same product and vendor ID but different
-            // language and layout type should configure the layouts correctly.
-            assertCorrectLayout(
-                englishQwertyKeyboardDevice,
-                frenchSubtype,
-                createLayoutDescriptor("keyboard_layout_english_us")
-            )
+    @Test
+    fun testGetDefaultKeyboardLayoutForInputDevice_withHwLanguageTagAndLayoutType() {
+        val frenchSubtype = createImeSubtypeForLanguageTagAndLayoutType("fr", "azerty")
+        // Should return English dvorak even if IME current layout is French, since HW says the
+        // keyboard is a Dvorak keyboard
+        assertCorrectLayout(
+            englishDvorakKeyboardDevice,
+            frenchSubtype,
+            createLayoutDescriptor("keyboard_layout_english_us_dvorak")
+        )
 
-            // Fallback to IME information if the HW provided layout script is incompatible with the
-            // provided IME subtype
-            assertCorrectLayout(
-                englishDvorakKeyboardDevice,
-                createImeSubtypeForLanguageTagAndLayoutType("ru", ""),
-                createLayoutDescriptor("keyboard_layout_russian")
-            )
-        }
+        // Back to back changing HW keyboards with same product and vendor ID but different
+        // language and layout type should configure the layouts correctly.
+        assertCorrectLayout(
+            englishQwertyKeyboardDevice,
+            frenchSubtype,
+            createLayoutDescriptor("keyboard_layout_english_us")
+        )
+
+        // Fallback to IME information if the HW provided layout script is incompatible with the
+        // provided IME subtype
+        assertCorrectLayout(
+            englishDvorakKeyboardDevice,
+            createImeSubtypeForLanguageTagAndLayoutType("ru", ""),
+            createLayoutDescriptor("keyboard_layout_russian")
+        )
     }
 
     @Test
@@ -867,27 +559,25 @@
                 KeyboardLayoutManager.ImeInfo(0, imeInfo,
                         createImeSubtypeForLanguageTagAndLayoutType("de-Latn", "qwertz")))
         Mockito.doReturn(imeInfos).`when`(keyboardLayoutManager).imeInfoListForLayoutMapping
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.onInputDeviceAdded(keyboardDevice.id)
-            ExtendedMockito.verify {
-                FrameworkStatsLog.write(
-                        ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
-                        ArgumentMatchers.anyBoolean(),
-                        ArgumentMatchers.eq(keyboardDevice.vendorId),
-                        ArgumentMatchers.eq(keyboardDevice.productId),
-                        ArgumentMatchers.eq(
-                            createByteArray(
-                                KeyboardMetricsCollector.DEFAULT_LANGUAGE_TAG,
-                                LAYOUT_TYPE_DEFAULT,
-                                GERMAN_LAYOUT_NAME,
-                                KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_VIRTUAL_KEYBOARD,
-                                "de-Latn",
-                                LAYOUT_TYPE_QWERTZ
-                            ),
+        keyboardLayoutManager.onInputDeviceAdded(keyboardDevice.id)
+        ExtendedMockito.verify {
+            FrameworkStatsLog.write(
+                    ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
+                    ArgumentMatchers.anyBoolean(),
+                    ArgumentMatchers.eq(keyboardDevice.vendorId),
+                    ArgumentMatchers.eq(keyboardDevice.productId),
+                    ArgumentMatchers.eq(
+                        createByteArray(
+                            KeyboardMetricsCollector.DEFAULT_LANGUAGE_TAG,
+                            LAYOUT_TYPE_DEFAULT,
+                            GERMAN_LAYOUT_NAME,
+                            KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_VIRTUAL_KEYBOARD,
+                            "de-Latn",
+                            LAYOUT_TYPE_QWERTZ
                         ),
-                        ArgumentMatchers.eq(keyboardDevice.deviceBus),
-                )
-            }
+                    ),
+                    ArgumentMatchers.eq(keyboardDevice.deviceBus),
+            )
         }
     }
 
@@ -897,27 +587,25 @@
                 KeyboardLayoutManager.ImeInfo(0, imeInfo,
                         createImeSubtypeForLanguageTagAndLayoutType("de-Latn", "qwertz")))
         Mockito.doReturn(imeInfos).`when`(keyboardLayoutManager).imeInfoListForLayoutMapping
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.onInputDeviceAdded(englishQwertyKeyboardDevice.id)
-            ExtendedMockito.verify {
-                FrameworkStatsLog.write(
-                        ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
-                        ArgumentMatchers.anyBoolean(),
-                        ArgumentMatchers.eq(englishQwertyKeyboardDevice.vendorId),
-                        ArgumentMatchers.eq(englishQwertyKeyboardDevice.productId),
-                        ArgumentMatchers.eq(
-                            createByteArray(
-                                "en",
-                                LAYOUT_TYPE_QWERTY,
-                                ENGLISH_US_LAYOUT_NAME,
-                                KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_DEVICE,
-                                "de-Latn",
-                                LAYOUT_TYPE_QWERTZ
-                            )
-                        ),
-                        ArgumentMatchers.eq(keyboardDevice.deviceBus),
-                )
-            }
+        keyboardLayoutManager.onInputDeviceAdded(englishQwertyKeyboardDevice.id)
+        ExtendedMockito.verify {
+            FrameworkStatsLog.write(
+                    ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
+                    ArgumentMatchers.anyBoolean(),
+                    ArgumentMatchers.eq(englishQwertyKeyboardDevice.vendorId),
+                    ArgumentMatchers.eq(englishQwertyKeyboardDevice.productId),
+                    ArgumentMatchers.eq(
+                        createByteArray(
+                            "en",
+                            LAYOUT_TYPE_QWERTY,
+                            ENGLISH_US_LAYOUT_NAME,
+                            KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_DEVICE,
+                            "de-Latn",
+                            LAYOUT_TYPE_QWERTZ
+                        )
+                    ),
+                    ArgumentMatchers.eq(keyboardDevice.deviceBus),
+            )
         }
     }
 
@@ -925,27 +613,25 @@
     fun testConfigurationLogged_onInputDeviceAdded_DefaultSelection() {
         val imeInfos = listOf(KeyboardLayoutManager.ImeInfo(0, imeInfo, createImeSubtype()))
         Mockito.doReturn(imeInfos).`when`(keyboardLayoutManager).imeInfoListForLayoutMapping
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.onInputDeviceAdded(keyboardDevice.id)
-            ExtendedMockito.verify {
-                FrameworkStatsLog.write(
-                        ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
-                        ArgumentMatchers.anyBoolean(),
-                        ArgumentMatchers.eq(keyboardDevice.vendorId),
-                        ArgumentMatchers.eq(keyboardDevice.productId),
-                        ArgumentMatchers.eq(
-                            createByteArray(
-                                KeyboardMetricsCollector.DEFAULT_LANGUAGE_TAG,
-                                LAYOUT_TYPE_DEFAULT,
-                                "Default",
-                                KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_DEFAULT,
-                                KeyboardMetricsCollector.DEFAULT_LANGUAGE_TAG,
-                                LAYOUT_TYPE_DEFAULT
-                            ),
+        keyboardLayoutManager.onInputDeviceAdded(keyboardDevice.id)
+        ExtendedMockito.verify {
+            FrameworkStatsLog.write(
+                    ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
+                    ArgumentMatchers.anyBoolean(),
+                    ArgumentMatchers.eq(keyboardDevice.vendorId),
+                    ArgumentMatchers.eq(keyboardDevice.productId),
+                    ArgumentMatchers.eq(
+                        createByteArray(
+                            KeyboardMetricsCollector.DEFAULT_LANGUAGE_TAG,
+                            LAYOUT_TYPE_DEFAULT,
+                            "Default",
+                            KeyboardLayoutSelectionResult.LAYOUT_SELECTION_CRITERIA_DEFAULT,
+                            KeyboardMetricsCollector.DEFAULT_LANGUAGE_TAG,
+                            LAYOUT_TYPE_DEFAULT
                         ),
-                        ArgumentMatchers.eq(keyboardDevice.deviceBus),
-                )
-            }
+                    ),
+                    ArgumentMatchers.eq(keyboardDevice.deviceBus),
+            )
         }
     }
 
@@ -953,19 +639,17 @@
     fun testConfigurationNotLogged_onInputDeviceChanged() {
         val imeInfos = listOf(KeyboardLayoutManager.ImeInfo(0, imeInfo, createImeSubtype()))
         Mockito.doReturn(imeInfos).`when`(keyboardLayoutManager).imeInfoListForLayoutMapping
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.onInputDeviceChanged(keyboardDevice.id)
-            ExtendedMockito.verify({
-                FrameworkStatsLog.write(
-                        ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
-                        ArgumentMatchers.anyBoolean(),
-                        ArgumentMatchers.anyInt(),
-                        ArgumentMatchers.anyInt(),
-                        ArgumentMatchers.any(ByteArray::class.java),
-                        ArgumentMatchers.anyInt(),
-                )
-            }, Mockito.times(0))
-        }
+        keyboardLayoutManager.onInputDeviceChanged(keyboardDevice.id)
+        ExtendedMockito.verify({
+            FrameworkStatsLog.write(
+                    ArgumentMatchers.eq(FrameworkStatsLog.KEYBOARD_CONFIGURED),
+                    ArgumentMatchers.anyBoolean(),
+                    ArgumentMatchers.anyInt(),
+                    ArgumentMatchers.anyInt(),
+                    ArgumentMatchers.any(ByteArray::class.java),
+                    ArgumentMatchers.anyInt(),
+            )
+        }, Mockito.times(0))
     }
 
     @Test
@@ -975,18 +659,16 @@
         Mockito.doReturn(false).`when`(keyboardLayoutManager).isVirtualDevice(
             ArgumentMatchers.eq(keyboardDevice.id)
         )
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.onInputDeviceChanged(keyboardDevice.id)
-            ExtendedMockito.verify(
-                notificationManager,
-                Mockito.times(1)
-            ).notifyAsUser(
-                ArgumentMatchers.isNull(),
-                ArgumentMatchers.anyInt(),
-                ArgumentMatchers.any(),
-                ArgumentMatchers.any()
-            )
-        }
+        keyboardLayoutManager.onInputDeviceChanged(keyboardDevice.id)
+        ExtendedMockito.verify(
+            notificationManager,
+            Mockito.times(1)
+        ).notifyAsUser(
+            ArgumentMatchers.isNull(),
+            ArgumentMatchers.anyInt(),
+            ArgumentMatchers.any(),
+            ArgumentMatchers.any()
+        )
     }
 
     @Test
@@ -996,18 +678,16 @@
         Mockito.doReturn(true).`when`(keyboardLayoutManager).isVirtualDevice(
             ArgumentMatchers.eq(keyboardDevice.id)
         )
-        NewSettingsApiFlag(true).use {
-            keyboardLayoutManager.onInputDeviceChanged(keyboardDevice.id)
-            ExtendedMockito.verify(
-                notificationManager,
-                Mockito.never()
-            ).notifyAsUser(
-                ArgumentMatchers.isNull(),
-                ArgumentMatchers.anyInt(),
-                ArgumentMatchers.any(),
-                ArgumentMatchers.any()
-            )
-        }
+        keyboardLayoutManager.onInputDeviceChanged(keyboardDevice.id)
+        ExtendedMockito.verify(
+            notificationManager,
+            Mockito.never()
+        ).notifyAsUser(
+            ArgumentMatchers.isNull(),
+            ArgumentMatchers.anyInt(),
+            ArgumentMatchers.any(),
+            ArgumentMatchers.any()
+        )
     }
 
     private fun assertCorrectLayout(
@@ -1019,7 +699,7 @@
             device.identifier, USER_ID, imeInfo, imeSubtype
         )
         assertEquals(
-            "New UI: getDefaultKeyboardLayoutForInputDevice should return $expectedLayout",
+            "getDefaultKeyboardLayoutForInputDevice should return $expectedLayout",
             expectedLayout,
             result.layoutDescriptor
         )
@@ -1123,21 +803,4 @@
         info.serviceInfo.name = RECEIVER_NAME
         return info
     }
-
-    private inner class NewSettingsApiFlag constructor(enabled: Boolean) : AutoCloseable {
-        init {
-            Settings.Global.putString(
-                context.contentResolver,
-                "settings_new_keyboard_ui", enabled.toString()
-            )
-        }
-
-        override fun close() {
-            Settings.Global.putString(
-                context.contentResolver,
-                "settings_new_keyboard_ui",
-                ""
-            )
-        }
-    }
 }