Merge "Set exclusion region based on drag resizability not task resizability" into main
diff --git a/core/java/android/app/appfunctions/AppFunctionRuntimeMetadata.java b/core/java/android/app/appfunctions/AppFunctionRuntimeMetadata.java
index 36daaab..83b5aa0 100644
--- a/core/java/android/app/appfunctions/AppFunctionRuntimeMetadata.java
+++ b/core/java/android/app/appfunctions/AppFunctionRuntimeMetadata.java
@@ -164,7 +164,13 @@
      */
     @Nullable
     public Boolean getEnabled() {
-        return (Boolean) getProperty(PROPERTY_ENABLED);
+        // We can't use getPropertyBoolean here. getPropertyBoolean returns false instead of null
+        // if the value is missing.
+        boolean[] enabled = getPropertyBooleanArray(PROPERTY_ENABLED);
+        if (enabled == null || enabled.length == 0) {
+            return null;
+        }
+        return enabled[0];
     }
 
     /** Returns the qualified id linking to the static metadata of the app function. */
@@ -201,11 +207,16 @@
         /**
          * Sets an indicator specifying if the function is enabled or not. This would override the
          * default enabled state in the static metadata ({@link
-         * AppFunctionStaticMetadataHelper#STATIC_PROPERTY_ENABLED_BY_DEFAULT}).
+         * AppFunctionStaticMetadataHelper#STATIC_PROPERTY_ENABLED_BY_DEFAULT}). Sets this to
+         * null to clear the override.
          */
         @NonNull
-        public Builder setEnabled(boolean enabled) {
-            setPropertyBoolean(PROPERTY_ENABLED, enabled);
+        public Builder setEnabled(@Nullable Boolean enabled) {
+            if (enabled == null) {
+                setPropertyBoolean(PROPERTY_ENABLED);
+            } else {
+                setPropertyBoolean(PROPERTY_ENABLED, enabled);
+            }
             return this;
         }
 
diff --git a/core/java/android/hardware/input/KeyGestureEvent.java b/core/java/android/hardware/input/KeyGestureEvent.java
index 0cabc4c..bdbec55 100644
--- a/core/java/android/hardware/input/KeyGestureEvent.java
+++ b/core/java/android/hardware/input/KeyGestureEvent.java
@@ -22,8 +22,6 @@
 import android.view.Display;
 import android.view.KeyCharacterMap;
 
-import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.util.AnnotationValidations;
 import com.android.internal.util.FrameworkStatsLog;
 
 import java.lang.annotation.Retention;
@@ -171,6 +169,14 @@
     }
 
     /**
+     * Tests whether this keyboard shortcut event has the given modifiers (i.e. all of the given
+     * modifiers were pressed when this shortcut was triggered).
+     */
+    public boolean hasModifiers(int modifiers) {
+        return (getModifierState() & modifiers) == modifiers;
+    }
+
+    /**
      * Key gesture event builder used to create a KeyGestureEvent for tests in Java.
      *
      * @hide
diff --git a/core/java/android/os/IUserManager.aidl b/core/java/android/os/IUserManager.aidl
index 00ba3bf..18f9b2b 100644
--- a/core/java/android/os/IUserManager.aidl
+++ b/core/java/android/os/IUserManager.aidl
@@ -139,7 +139,7 @@
     boolean isUserForeground(int userId);
     boolean isUserVisible(int userId);
     int[] getVisibleUsers();
-    int getMainDisplayIdAssignedToUser();
+    int getMainDisplayIdAssignedToUser(int userId);
     boolean isForegroundUserAdmin();
     boolean isUserNameSet(int userId);
     boolean hasRestrictedProfiles(int userId);
diff --git a/core/java/android/os/UserManager.java b/core/java/android/os/UserManager.java
index 1ca4574..461f1e0 100644
--- a/core/java/android/os/UserManager.java
+++ b/core/java/android/os/UserManager.java
@@ -3706,9 +3706,13 @@
      * @hide
      */
     @TestApi
+    @UserHandleAware(
+            requiresAnyOfPermissionsIfNotCaller = {
+                    android.Manifest.permission.MANAGE_USERS,
+                    android.Manifest.permission.INTERACT_ACROSS_USERS})
     public int getMainDisplayIdAssignedToUser() {
         try {
-            return mService.getMainDisplayIdAssignedToUser();
+            return mService.getMainDisplayIdAssignedToUser(mUserId);
         } catch (RemoteException re) {
             throw re.rethrowFromSystemServer();
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
index dfc5ab3..2bc01b2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandler.kt
@@ -304,14 +304,14 @@
         val leafTaskFilter = TransitionUtil.LeafTaskFilter()
         info.changes.withIndex().forEach { (i, change) ->
             if (TransitionUtil.isWallpaper(change)) {
-                val layer = layers.wallpaperLayers - i
+                val layer = layers.topWallpaperLayer - i
                 startTransaction.apply {
                     setLayer(change.leash, layer)
                     show(change.leash)
                 }
             } else if (isHomeChange(change)) {
                 state.homeChange = change
-                val layer = layers.homeLayers - i
+                val layer = layers.topHomeLayer - i
                 startTransaction.apply {
                     setLayer(change.leash, layer)
                     show(change.leash)
@@ -325,7 +325,7 @@
                             if (state.cancelState == CancelState.NO_CANCEL) {
                                 // Normal case, split root goes to the bottom behind everything
                                 // else.
-                                layers.appLayers - i
+                                layers.topAppLayer - i
                             } else {
                                 // Cancel-early case, pretend nothing happened so split root stays
                                 // top.
@@ -357,7 +357,7 @@
                             state.otherRootChanges.add(change)
                             val bounds = change.endAbsBounds
                             startTransaction.apply {
-                                setLayer(change.leash, layers.appLayers - i)
+                                setLayer(change.leash, layers.topAppLayer - i)
                                 setWindowCrop(change.leash, bounds.width(), bounds.height())
                                 show(change.leash)
                             }
@@ -398,6 +398,7 @@
                 }
             }
         }
+        state.surfaceLayers = layers
         state.startTransitionFinishCb = finishCallback
         state.startTransitionFinishTransaction = finishTransaction
         startTransaction.apply()
@@ -522,6 +523,10 @@
                     startTransaction.show(change.leash)
                     finishTransaction.show(change.leash)
                     state.draggedTaskChange = change
+                    // Restoring the dragged leash layer as it gets reset in the merge transition
+                    state.surfaceLayers?.let {
+                        startTransaction.setLayer(change.leash, it.dragLayer)
+                    }
                 }
                 change.taskInfo?.windowingMode == WINDOWING_MODE_FREEFORM -> {
                     // Other freeform tasks that are being restored go behind the dragged task.
@@ -647,8 +652,15 @@
         }
     }
 
-    private fun isHomeChange(change: Change): Boolean {
-        return change.taskInfo?.activityType == ACTIVITY_TYPE_HOME
+    /** Checks if the change is a home task change */
+    @VisibleForTesting
+    fun isHomeChange(change: Change): Boolean {
+        return change.taskInfo?.let {
+            it.activityType == ACTIVITY_TYPE_HOME &&
+                // Skip translucent wizard task with type home
+                // TODO(b/368334295): Remove when the multiple home changes issue is resolved
+                !(it.isTopActivityTransparent && it.numActivities == 1)
+        } ?: false
     }
 
     private fun startCancelAnimation() {
@@ -765,12 +777,18 @@
 
     /**
      * Represents the layering (Z order) that will be given to any window based on its type during
-     * the "start" transition of the drag-to-desktop transition
+     * the "start" transition of the drag-to-desktop transition.
+     *
+     * @param topAppLayer Used to calculate the app layer z-order = `topAppLayer - changeIndex`.
+     * @param topHomeLayer Used to calculate the home layer z-order = `topHomeLayer - changeIndex`.
+     * @param topWallpaperLayer Used to calculate the wallpaper layer z-order = `topWallpaperLayer -
+     *   changeIndex`
+     * @param dragLayer Defines the drag layer z-order
      */
-    protected data class DragToDesktopLayers(
-        val appLayers: Int,
-        val homeLayers: Int,
-        val wallpaperLayers: Int,
+    data class DragToDesktopLayers(
+        val topAppLayer: Int,
+        val topHomeLayer: Int,
+        val topWallpaperLayer: Int,
         val dragLayer: Int,
     )
 
@@ -790,6 +808,7 @@
         abstract var homeChange: Change?
         abstract var draggedTaskChange: Change?
         abstract var freeformTaskChanges: List<Change>
+        abstract var surfaceLayers: DragToDesktopLayers?
         abstract var cancelState: CancelState
         abstract var startAborted: Boolean
 
@@ -803,6 +822,7 @@
             override var homeChange: Change? = null,
             override var draggedTaskChange: Change? = null,
             override var freeformTaskChanges: List<Change> = emptyList(),
+            override var surfaceLayers: DragToDesktopLayers? = null,
             override var cancelState: CancelState = CancelState.NO_CANCEL,
             override var startAborted: Boolean = false,
             var otherRootChanges: MutableList<Change> = mutableListOf()
@@ -818,6 +838,7 @@
             override var homeChange: Change? = null,
             override var draggedTaskChange: Change? = null,
             override var freeformTaskChanges: List<Change> = emptyList(),
+            override var surfaceLayers: DragToDesktopLayers? = null,
             override var cancelState: CancelState = CancelState.NO_CANCEL,
             override var startAborted: Boolean = false,
             var splitRootChange: Change? = null,
@@ -872,9 +893,9 @@
      */
     override fun calculateStartDragToDesktopLayers(info: TransitionInfo): DragToDesktopLayers =
         DragToDesktopLayers(
-            appLayers = info.changes.size,
-            homeLayers = info.changes.size * 2,
-            wallpaperLayers = info.changes.size * 3,
+            topAppLayer = info.changes.size,
+            topHomeLayer = info.changes.size * 2,
+            topWallpaperLayer = info.changes.size * 3,
             dragLayer = info.changes.size * 3
         )
 }
@@ -914,9 +935,9 @@
      */
     override fun calculateStartDragToDesktopLayers(info: TransitionInfo): DragToDesktopLayers =
         DragToDesktopLayers(
-            appLayers = -1,
-            homeLayers = info.changes.size - 1,
-            wallpaperLayers = info.changes.size * 2 - 1,
+            topAppLayer = -1,
+            topHomeLayer = info.changes.size - 1,
+            topWallpaperLayer = info.changes.size * 2 - 1,
             dragLayer = info.changes.size * 2
         )
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index e47ebea..b24c17c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -576,10 +576,6 @@
         mWindowDecorCaptionHandleRepository.notifyCaptionChanged(captionState);
     }
 
-    private static boolean isDragResizable(ActivityManager.RunningTaskInfo taskInfo) {
-        return taskInfo.isFreeform() && taskInfo.isResizeable;
-    }
-
     private void updateMaximizeMenu(SurfaceControl.Transaction startT) {
         if (!isDragResizable(mTaskInfo) || !isMaximizeMenuActive()) {
             return;
diff --git a/libs/WindowManager/Shell/tests/e2e/splitscreen/scenarios/src/com/android/wm/shell/scenarios/SwitchBackToSplitFromRecent.kt b/libs/WindowManager/Shell/tests/e2e/splitscreen/scenarios/src/com/android/wm/shell/scenarios/SwitchBackToSplitFromRecent.kt
index c7cbc3e4..22adf6c 100644
--- a/libs/WindowManager/Shell/tests/e2e/splitscreen/scenarios/src/com/android/wm/shell/scenarios/SwitchBackToSplitFromRecent.kt
+++ b/libs/WindowManager/Shell/tests/e2e/splitscreen/scenarios/src/com/android/wm/shell/scenarios/SwitchBackToSplitFromRecent.kt
@@ -48,6 +48,7 @@
     fun setup() {
         tapl.workspace.switchToOverview().dismissAllTasks()
 
+        tapl.setExpectedRotationCheckEnabled(false)
         tapl.setEnableRotation(true)
         tapl.setExpectedRotation(rotation.value)
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java
index 24f4d92..e6bd05b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestRunningTaskInfoBuilder.java
@@ -47,6 +47,8 @@
     private ActivityManager.TaskDescription.Builder mTaskDescriptionBuilder = null;
     private final Point mPositionInParent = new Point();
     private boolean mIsVisible = false;
+    private boolean mIsTopActivityTransparent = false;
+    private int mNumActivities = 1;
     private long mLastActiveTime;
 
     public static WindowContainerToken createMockWCToken() {
@@ -113,6 +115,16 @@
         return this;
     }
 
+    public TestRunningTaskInfoBuilder setTopActivityTransparent(boolean isTopActivityTransparent) {
+        mIsTopActivityTransparent = isTopActivityTransparent;
+        return this;
+    }
+
+    public TestRunningTaskInfoBuilder setNumActivities(int numActivities) {
+        mNumActivities = numActivities;
+        return this;
+    }
+
     public TestRunningTaskInfoBuilder setLastActiveTime(long lastActiveTime) {
         mLastActiveTime = lastActiveTime;
         return this;
@@ -134,6 +146,8 @@
                 mTaskDescriptionBuilder != null ? mTaskDescriptionBuilder.build() : null;
         info.positionInParent = mPositionInParent;
         info.isVisible = mIsVisible;
+        info.isTopActivityTransparent = mIsTopActivityTransparent;
+        info.numActivities = mNumActivities;
         info.lastActiveTime = mLastActiveTime;
         return info;
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
index 497d0e5..d9387d2 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DragToDesktopTransitionHandlerTest.kt
@@ -35,6 +35,7 @@
 import java.util.function.Supplier
 import junit.framework.Assert.assertEquals
 import junit.framework.Assert.assertFalse
+import junit.framework.Assert.assertTrue
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
@@ -212,6 +213,60 @@
     }
 
     @Test
+    fun isHomeChange_withoutTaskInfo_returnsFalse() {
+        val change =
+            TransitionInfo.Change(mock(), homeTaskLeash).apply {
+                parent = null
+                taskInfo = null
+            }
+
+        assertFalse(defaultHandler.isHomeChange(change))
+        assertFalse(springHandler.isHomeChange(change))
+    }
+
+    @Test
+    fun isHomeChange_withStandardActivityTaskInfo_returnsFalse() {
+        val change =
+            TransitionInfo.Change(mock(), homeTaskLeash).apply {
+                parent = null
+                taskInfo =
+                    TestRunningTaskInfoBuilder().setActivityType(ACTIVITY_TYPE_STANDARD).build()
+            }
+
+        assertFalse(defaultHandler.isHomeChange(change))
+        assertFalse(springHandler.isHomeChange(change))
+    }
+
+    @Test
+    fun isHomeChange_withHomeActivityTaskInfo_returnsTrue() {
+        val change =
+            TransitionInfo.Change(mock(), homeTaskLeash).apply {
+                parent = null
+                taskInfo = TestRunningTaskInfoBuilder().setActivityType(ACTIVITY_TYPE_HOME).build()
+            }
+
+        assertTrue(defaultHandler.isHomeChange(change))
+        assertTrue(springHandler.isHomeChange(change))
+    }
+
+    @Test
+    fun isHomeChange_withSingleTranslucentHomeActivityTaskInfo_returnsFalse() {
+        val change =
+            TransitionInfo.Change(mock(), homeTaskLeash).apply {
+                parent = null
+                taskInfo =
+                    TestRunningTaskInfoBuilder()
+                        .setActivityType(ACTIVITY_TYPE_HOME)
+                        .setTopActivityTransparent(true)
+                        .setNumActivities(1)
+                        .build()
+            }
+
+        assertFalse(defaultHandler.isHomeChange(change))
+        assertFalse(springHandler.isHomeChange(change))
+    }
+
+    @Test
     fun cancelDragToDesktop_startWasReady_cancel() {
         startDrag(defaultHandler)
 
@@ -343,6 +398,8 @@
         // Should show dragged task layer in start and finish transaction
         verify(mergedStartTransaction).show(draggedTaskLeash)
         verify(playingFinishTransaction).show(draggedTaskLeash)
+        // Should update the dragged task layer
+        verify(mergedStartTransaction).setLayer(eq(draggedTaskLeash), anyInt())
         // Should merge animation
         verify(finishCallback).onTransitionFinished(null)
     }
@@ -373,6 +430,8 @@
         // Should show dragged task layer in start and finish transaction
         verify(mergedStartTransaction).show(draggedTaskLeash)
         verify(playingFinishTransaction).show(draggedTaskLeash)
+        // Should update the dragged task layer
+        verify(mergedStartTransaction).setLayer(eq(draggedTaskLeash), anyInt())
         // Should hide home task leash in finish transaction
         verify(playingFinishTransaction).hide(homeTaskLeash)
         // Should merge animation
diff --git a/libs/appfunctions/Android.bp b/libs/appfunctions/Android.bp
index 09e2f42..c6cee07 100644
--- a/libs/appfunctions/Android.bp
+++ b/libs/appfunctions/Android.bp
@@ -29,3 +29,11 @@
     no_dist: true,
     unsafe_ignore_missing_latest_api: true,
 }
+
+prebuilt_etc {
+    name: "appfunctions.sidecar.xml",
+    system_ext_specific: true,
+    sub_dir: "permissions",
+    src: "appfunctions.sidecar.xml",
+    filename_from_src: true,
+}
diff --git a/libs/appfunctions/appfunctions.sidecar.xml b/libs/appfunctions/appfunctions.sidecar.xml
new file mode 100644
index 0000000..bef8b6e
--- /dev/null
+++ b/libs/appfunctions/appfunctions.sidecar.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<permissions>
+    <library
+        name="com.google.android.appfunctions.sidecar"
+        file="/system_ext/framework/com.google.android.appfunctions.sidecar.jar"/>
+</permissions>
\ No newline at end of file
diff --git a/media/java/android/media/IAudioService.aidl b/media/java/android/media/IAudioService.aidl
index 9af6b28..8394daf 100644
--- a/media/java/android/media/IAudioService.aidl
+++ b/media/java/android/media/IAudioService.aidl
@@ -709,6 +709,10 @@
     @EnforcePermission("MODIFY_AUDIO_ROUTING")
     List<AudioFocusInfo> getFocusStack();
 
+    @EnforcePermission("MODIFY_AUDIO_ROUTING")
+    oneway void sendFocusLossAndUpdate(in AudioFocusInfo focusLoser, in IAudioPolicyCallback apcb);
+
+    @EnforcePermission("MODIFY_AUDIO_ROUTING")
     boolean sendFocusLoss(in AudioFocusInfo focusLoser, in IAudioPolicyCallback apcb);
 
     @EnforcePermission("MODIFY_AUDIO_ROUTING")
diff --git a/media/java/android/media/audiopolicy/AudioPolicy.java b/media/java/android/media/audiopolicy/AudioPolicy.java
index 293a8f8..2c8e352 100644
--- a/media/java/android/media/audiopolicy/AudioPolicy.java
+++ b/media/java/android/media/audiopolicy/AudioPolicy.java
@@ -927,6 +927,29 @@
     }
 
     /**
+     * @hide
+     * Causes the given audio focus owner to lose audio focus with
+     * {@link android.media.AudioManager#AUDIOFOCUS_LOSS}, and be removed from the focus stack.
+     * Unlike {@link #sendFocusLoss(AudioFocusInfo)}, the method causes the focus stack
+     * to be reevaluated as the discarded focus owner may have been at the top of stack,
+     * and now the new owner needs to be notified of the gain.
+     * @param focusLoser identifies the focus owner to discard from the focus stack
+     * @throws IllegalStateException if used on an unregistered policy, or a registered policy
+     * with no {@link AudioPolicyFocusListener} set
+     * @see #getFocusStack()
+     * @see #sendFocusLoss(AudioFocusInfo)
+     */
+    @RequiresPermission(android.Manifest.permission.MODIFY_AUDIO_ROUTING)
+    public void sendFocusLossAndUpdate(@NonNull AudioFocusInfo focusLoser)
+            throws IllegalStateException {
+        try {
+            getService().sendFocusLossAndUpdate(Objects.requireNonNull(focusLoser), cb());
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Send AUDIOFOCUS_LOSS to a specific stack entry, causing it to be notified of the focus
      * loss, and for it to exit the focus stack (its focus listener will not be invoked after that).
      * This operation is only valid for a registered policy (with
diff --git a/media/java/android/media/midi/MidiManager.java b/media/java/android/media/midi/MidiManager.java
index 244292d..bfc1324 100644
--- a/media/java/android/media/midi/MidiManager.java
+++ b/media/java/android/media/midi/MidiManager.java
@@ -393,6 +393,16 @@
 
     /**
      * Opens a Bluetooth MIDI device for reading and writing.
+     * Bluetooth MIDI devices are only available after openBluetoothDevice() is called.
+     * Once that happens anywhere in the system, then the BLE-MIDI device will appear as just
+     * another MidiDevice to other apps.
+     *
+     * If the device opened using openBluetoothDevice()  is closed, then it will no longer be
+     * available. To other apps, it will appear as if the BLE MidiDevice had been unplugged.
+     * If a MidiDevice is garbage collected then it will be closed automatically.
+     * If you want the BLE-MIDI device to remain available you should keep the object alive.
+     *
+     * You may close the device with MidiDevice.close().
      *
      * @param bluetoothDevice a {@link android.bluetooth.BluetoothDevice} to open as a MIDI device
      * @param listener a {@link MidiManager.OnDeviceOpenedListener} to be called to receive the
diff --git a/packages/SettingsLib/SettingsTheme/Android.bp b/packages/SettingsLib/SettingsTheme/Android.bp
index baeff7e..1661dfb 100644
--- a/packages/SettingsLib/SettingsTheme/Android.bp
+++ b/packages/SettingsLib/SettingsTheme/Android.bp
@@ -15,7 +15,10 @@
         "src/**/*.kt",
     ],
     resource_dirs: ["res"],
-    static_libs: ["androidx.preference_preference"],
+    static_libs: [
+        "androidx.preference_preference",
+        "com.google.android.material_material",
+    ],
     sdk_version: "system_current",
     min_sdk_version: "21",
     apex_available: [
diff --git a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_icon_check.xml b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_icon_check.xml
new file mode 100644
index 0000000..309dbdf
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_icon_check.xml
@@ -0,0 +1,27 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960"
+    android:tint="?attr/colorControlNormal">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M382,720L154,492L211,435L382,606L749,239L806,296L382,720Z"/>
+</vector>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_icon_close.xml b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_icon_close.xml
new file mode 100644
index 0000000..e6df8a4
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_icon_close.xml
@@ -0,0 +1,27 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960"
+    android:tint="?attr/colorControlNormal">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M256,760L200,704L424,480L200,256L256,200L480,424L704,200L760,256L536,480L760,704L704,760L480,536L256,760Z"/>
+</vector>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_switch_thumb_icon.xml b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_switch_thumb_icon.xml
new file mode 100644
index 0000000..342729d
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_expressive_switch_thumb_icon.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:state_checked="true" android:drawable="@drawable/settingslib_expressive_icon_check"/>
+    <item android:state_checked="false" android:drawable="@drawable/settingslib_expressive_icon_close"/>
+</selector>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_bottom_selected.xml b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_bottom_selected.xml
index f4766ee..543b237 100644
--- a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_bottom_selected.xml
+++ b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_bottom_selected.xml
@@ -26,7 +26,10 @@
             <solid
                 android:color="@color/settingslib_materialColorSurfaceContainer" />
             <corners
-                android:radius="@dimen/settingslib_preference_corner_radius_selected" />
+                android:topLeftRadius="4dp"
+                android:bottomLeftRadius="@dimen/settingslib_preference_corner_radius"
+                android:topRightRadius="4dp"
+                android:bottomRightRadius="@dimen/settingslib_preference_corner_radius" />
             <padding
                 android:bottom="16dp"/>
         </shape>
diff --git a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_center_selected.xml b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_center_selected.xml
index 40eafc2..6d2cd1a 100644
--- a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_center_selected.xml
+++ b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_center_selected.xml
@@ -25,7 +25,7 @@
             <solid
                 android:color="@color/settingslib_materialColorSurfaceContainer" />
             <corners
-                android:radius="@dimen/settingslib_preference_corner_radius_selected" />
+                android:radius="4dp" />
         </shape>
     </item>
 </ripple>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_selected.xml b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_selected.xml
index f4766ee..bcdbf1d 100644
--- a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_selected.xml
+++ b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_selected.xml
@@ -26,7 +26,7 @@
             <solid
                 android:color="@color/settingslib_materialColorSurfaceContainer" />
             <corners
-                android:radius="@dimen/settingslib_preference_corner_radius_selected" />
+                android:radius="@dimen/settingslib_preference_corner_radius" />
             <padding
                 android:bottom="16dp"/>
         </shape>
diff --git a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_top_selected.xml b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_top_selected.xml
index 40eafc2..d4b658c 100644
--- a/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_top_selected.xml
+++ b/packages/SettingsLib/SettingsTheme/res/drawable-v35/settingslib_round_background_top_selected.xml
@@ -25,7 +25,10 @@
             <solid
                 android:color="@color/settingslib_materialColorSurfaceContainer" />
             <corners
-                android:radius="@dimen/settingslib_preference_corner_radius_selected" />
+                android:topLeftRadius="@dimen/settingslib_preference_corner_radius"
+                android:bottomLeftRadius="4dp"
+                android:topRightRadius="@dimen/settingslib_preference_corner_radius"
+                android:bottomRightRadius="4dp" />
         </shape>
     </item>
 </ripple>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference.xml b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference.xml
new file mode 100644
index 0000000..2475dfd
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference.xml
@@ -0,0 +1,44 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:minHeight="?android:attr/listPreferredItemHeightSmall"
+    android:gravity="center_vertical"
+    android:paddingStart="?android:attr/listPreferredItemPaddingStart"
+    android:paddingEnd="?android:attr/listPreferredItemPaddingEnd"
+    android:background="?android:attr/selectableItemBackground"
+    android:clipToPadding="false"
+    android:baselineAligned="false">
+
+    <include layout="@layout/settingslib_expressive_preference_icon_frame"/>
+
+    <include layout="@layout/settingslib_expressive_preference_text_frame"/>
+
+    <!-- Preference should place its actual preference widget here. -->
+    <LinearLayout
+        android:id="@android:id/widget_frame"
+        android:layout_width="wrap_content"
+        android:layout_height="match_parent"
+        android:gravity="end|center_vertical"
+        android:paddingStart="@dimen/settingslib_expressive_space_small1"
+        android:paddingEnd="0dp"
+        android:orientation="vertical"/>
+
+</LinearLayout>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_icon_frame.xml b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_icon_frame.xml
new file mode 100644
index 0000000..f5017a5
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_icon_frame.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/icon_frame"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:minWidth="@dimen/settingslib_expressive_space_medium3"
+    android:minHeight="@dimen/settingslib_expressive_space_medium3"
+    android:gravity="center"
+    android:layout_marginEnd="-8dp">
+
+    <androidx.preference.internal.PreferenceImageView
+        android:id="@android:id/icon"
+        android:layout_width="@dimen/settingslib_expressive_space_medium3"
+        android:layout_height="@dimen/settingslib_expressive_space_medium3"
+        android:scaleType="centerInside"/>
+
+</LinearLayout>
diff --git a/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_switch.xml b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_switch.xml
new file mode 100644
index 0000000..4cbdfd5
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_switch.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<com.google.android.material.materialswitch.MaterialSwitch
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:theme="@style/Theme.Material3.DynamicColors.DayNight"
+    android:id="@+id/switchWidget"
+    style="@style/SettingslibSwitchStyle.Expressive"/>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_text_frame.xml b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_text_frame.xml
new file mode 100644
index 0000000..e3e689b
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/layout-v35/settingslib_expressive_preference_text_frame.xml
@@ -0,0 +1,47 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<RelativeLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="@dimen/settingslib_expressive_space_none"
+    android:layout_height="wrap_content"
+    android:layout_weight="1"
+    android:padding="@dimen/settingslib_expressive_space_small1">
+
+    <TextView
+        android:id="@android:id/title"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_gravity="start"
+        android:textAlignment="viewStart"
+        android:textAppearance="?android:attr/textAppearanceListItem"
+        android:maxLines="2"
+        android:ellipsize="marquee"/>
+
+    <TextView
+        android:id="@android:id/summary"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_below="@android:id/title"
+        android:layout_alignLeft="@android:id/title"
+        android:layout_alignStart="@android:id/title"
+        android:layout_gravity="start"
+        android:textAlignment="viewStart"
+        android:textAppearance="?android:attr/textAppearanceListItemSecondary"
+        android:textColor="?android:attr/textColorSecondary"
+        android:maxLines="10"/>
+</RelativeLayout>
diff --git a/packages/SettingsLib/SettingsTheme/res/values-v35/dimens_expressive.xml b/packages/SettingsLib/SettingsTheme/res/values-v35/dimens_expressive.xml
new file mode 100644
index 0000000..2320aab
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/values-v35/dimens_expressive.xml
@@ -0,0 +1,57 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  Copyright (C) 2024 The Android Open Source Project
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+  -->
+
+<resources>
+    <!-- Expressive design start -->
+    <!-- corner radius token -->
+    <dimen name="settingslib_expressive_radius_none">0dp</dimen>
+    <dimen name="settingslib_expressive_radius_full">360dp</dimen>
+    <dimen name="settingslib_expressive_radius_extrasmall1">2dp</dimen>
+    <dimen name="settingslib_expressive_radius_extrasmall2">4dp</dimen>
+    <dimen name="settingslib_expressive_radius_small">8dp</dimen>
+    <dimen name="settingslib_expressive_radius_medium">12dp</dimen>
+    <dimen name="settingslib_expressive_radius_large1">16dp</dimen>
+    <dimen name="settingslib_expressive_radius_large2">20dp</dimen>
+    <dimen name="settingslib_expressive_radius_large3">24dp</dimen>
+    <dimen name="settingslib_expressive_radius_extralarge1">28dp</dimen>
+    <dimen name="settingslib_expressive_radius_extralarge2">32dp</dimen>
+    <dimen name="settingslib_expressive_radius_extralarge3">42dp</dimen>
+
+    <!-- space token -->
+    <dimen name="settingslib_expressive_space_none">0dp</dimen>
+    <dimen name="settingslib_expressive_space_extrasmall1">2dp</dimen>
+    <dimen name="settingslib_expressive_space_extrasmall2">4dp</dimen>
+    <dimen name="settingslib_expressive_space_extrasmall3">6dp</dimen>
+    <dimen name="settingslib_expressive_space_extrasmall4">8dp</dimen>
+    <dimen name="settingslib_expressive_space_extrasmall5">10dp</dimen>
+    <dimen name="settingslib_expressive_space_extrasmall6">12dp</dimen>
+    <dimen name="settingslib_expressive_space_extrasmall7">14dp</dimen>
+    <dimen name="settingslib_expressive_space_small1">16dp</dimen>
+    <dimen name="settingslib_expressive_space_small2">18dp</dimen>
+    <dimen name="settingslib_expressive_space_small3">20dp</dimen>
+    <dimen name="settingslib_expressive_space_small4">24dp</dimen>
+    <dimen name="settingslib_expressive_space_medium1">32dp</dimen>
+    <dimen name="settingslib_expressive_space_medium2">36dp</dimen>
+    <dimen name="settingslib_expressive_space_medium3">40dp</dimen>
+    <dimen name="settingslib_expressive_space_medium4">48dp</dimen>
+    <dimen name="settingslib_expressive_space_large1">60dp</dimen>
+    <dimen name="settingslib_expressive_space_large2">64dp</dimen>
+    <dimen name="settingslib_expressive_space_large3">72dp</dimen>
+    <dimen name="settingslib_expressive_space_large4">80dp</dimen>
+    <dimen name="settingslib_expressive_space_large5">96dp</dimen>
+    <!-- Expressive theme end -->
+</resources>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/values-v35/styles_expressive.xml b/packages/SettingsLib/SettingsTheme/res/values-v35/styles_expressive.xml
new file mode 100644
index 0000000..04ae80e
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/values-v35/styles_expressive.xml
@@ -0,0 +1,172 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  Copyright (C) 2024 The Android Open Source Project
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+  -->
+
+<resources>
+    <style name="SettingsLibTextAppearance" parent="@android:style/TextAppearance.DeviceDefault">
+        <!--item name="android:fontFamily"></item-->
+        <item name="android:hyphenationFrequency">normalFast</item>
+        <item name="android:lineBreakWordStyle">phrase</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Primary">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+
+    <style name="SettingsLibTextAppearance.Primary.Display">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Display.Large">
+        <item name="android:textSize">57sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Display.Medium">
+        <item name="android:textSize">45sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Display.Small">
+        <item name="android:textSize">36sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Primary.Headline">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Headline.Large">
+        <item name="android:textSize">32sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Headline.Medium">
+        <item name="android:textSize">28sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Headline.Small">
+        <item name="android:textSize">24sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Primary.Title">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Title.Large">
+        <item name="android:textSize">22sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Title.Medium">
+        <item name="android:textSize">16sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Title.Small">
+        <item name="android:textSize">14sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Primary.Label">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Label.Large">
+        <item name="android:textSize">14sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Label.Medium">
+        <item name="android:textSize">12sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Label.Small">
+        <item name="android:textSize">11sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Primary.Body">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Body.Large">
+        <item name="android:textSize">16sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Body.Medium">
+        <item name="android:textSize">14sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Primary.Body.Small">
+        <item name="android:textSize">12sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Emphasized">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+
+    <style name="SettingsLibTextAppearance.Emphasized.Display">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Display.Large">
+        <item name="android:textSize">57sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Display.Medium">
+        <item name="android:textSize">45sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Display.Small">
+        <item name="android:textSize">36sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Emphasized.Headline">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Headline.Large">
+        <item name="android:textSize">32sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Headline.Medium">
+        <item name="android:textSize">28sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Headline.Small">
+        <item name="android:textSize">24sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Emphasized.Title">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Title.Large">
+        <item name="android:textSize">22sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Title.Medium">
+        <item name="android:textSize">16sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Title.Small">
+        <item name="android:textSize">14sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Emphasized.Label">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Label.Large">
+        <item name="android:textSize">14sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Label.Medium">
+        <item name="android:textSize">12sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Label.Small">
+        <item name="android:textSize">11sp</item>
+    </style>
+
+    <style name="SettingsLibTextAppearance.Emphasized.Body">
+        <!--item name="android:fontFamily"></item-->
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Body.Large">
+        <item name="android:textSize">16sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Body.Medium">
+        <item name="android:textSize">14sp</item>
+    </style>
+    <style name="SettingsLibTextAppearance.Emphasized.Body.Small">
+        <item name="android:textSize">12sp</item>
+    </style>
+
+    <style name="SettingslibSwitchStyle.Expressive" parent="">
+        <item name="android:layout_width">wrap_content</item>
+        <item name="android:layout_height">wrap_content</item>
+        <item name="android:minWidth">@dimen/settingslib_expressive_space_medium4</item>
+        <item name="android:background">@null</item>
+        <item name="android:clickable">false</item>
+        <item name="android:focusable">false</item>
+        <item name="thumbIcon">@drawable/settingslib_expressive_switch_thumb_icon</item>
+    </style>
+</resources>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/values-v35/styles_preference_expressive.xml b/packages/SettingsLib/SettingsTheme/res/values-v35/styles_preference_expressive.xml
new file mode 100644
index 0000000..3c69027
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/values-v35/styles_preference_expressive.xml
@@ -0,0 +1,72 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+<resources>
+    <style name="SettingsLibPreference" parent="SettingsPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.Category" parent="SettingsCategoryPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.CheckBoxPreference" parent="SettingsCheckBoxPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.SwitchPreferenceCompat" parent="SettingsSwitchPreferenceCompat.SettingsLib"/>
+
+    <style name="SettingsLibPreference.SeekBarPreference" parent="SettingsSeekbarPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.PreferenceScreen" parent="SettingsPreferenceScreen.SettingsLib"/>
+
+    <style name="SettingsLibPreference.DialogPreference" parent="SettingsPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.DialogPreference.EditTextPreference" parent="SettingsEditTextPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.DropDown" parent="SettingsDropdownPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.SwitchPreference" parent="SettingsSwitchPreference.SettingsLib"/>
+
+    <style name="SettingsLibPreference.Expressive">
+        <item name="android:layout">@layout/settingslib_expressive_preference</item>
+    </style>
+
+    <style name="SettingsLibPreference.Category.Expressive">
+    </style>
+
+    <style name="SettingsLibPreference.CheckBoxPreference.Expressive">
+        <item name="android:layout">@layout/settingslib_expressive_preference</item>
+    </style>
+
+    <style name="SettingsLibPreference.SwitchPreferenceCompat.Expressive">
+        <item name="android:layout">@layout/settingslib_expressive_preference</item>
+        <item name="android:widgetLayout">@layout/settingslib_expressive_preference_switch</item>
+    </style>
+
+    <style name="SettingsLibPreference.SeekBarPreference.Expressive"/>
+
+    <style name="SettingsLibPreference.PreferenceScreen.Expressive">
+        <item name="android:layout">@layout/settingslib_expressive_preference</item>
+    </style>
+
+    <style name="SettingsLibPreference.DialogPreference.Expressive">
+    </style>
+
+    <style name="SettingsLibPreference.DialogPreference.EditTextPreference.Expressive">
+        <item name="android:layout">@layout/settingslib_expressive_preference</item>
+        <item name="android:dialogLayout">@layout/settingslib_preference_dialog_edittext</item>
+    </style>
+
+    <style name="SettingsLibPreference.DropDown.Expressive">
+    </style>
+
+    <style name="SettingsLibPreference.SwitchPreference.Expressive"/>
+</resources>
diff --git a/packages/SettingsLib/SettingsTheme/res/values-v35/themes_expressive.xml b/packages/SettingsLib/SettingsTheme/res/values-v35/themes_expressive.xml
new file mode 100644
index 0000000..fea8739
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/values-v35/themes_expressive.xml
@@ -0,0 +1,58 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<resources>
+    <style name="Theme.SettingsBase.Expressive">
+        <!-- Set up Preference title text style -->
+        <!--item name="android:textAppearanceListItem">@style/TextAppearance.PreferenceTitle.SettingsLib</item-->
+        <!--item name="android:textAppearanceListItemSecondary">@style/textAppearanceListItemSecondary</item-->
+
+        <!-- Set up  list item padding -->
+        <item name="android:listPreferredItemPaddingStart">@dimen/settingslib_expressive_space_small1</item>
+        <item name="android:listPreferredItemPaddingLeft">@dimen/settingslib_expressive_space_small1</item>
+        <item name="android:listPreferredItemPaddingEnd">@dimen/settingslib_expressive_space_small1</item>
+        <item name="android:listPreferredItemPaddingRight">@dimen/settingslib_expressive_space_small1</item>
+        <item name="android:listPreferredItemHeightSmall">@dimen/settingslib_expressive_space_large3</item>
+
+        <!-- Set up preference theme -->
+        <item name="preferenceTheme">@style/PreferenceTheme.SettingsLib.Expressive</item>
+
+        <!-- Set up Spinner style -->
+        <!--item name="android:spinnerStyle"></item>
+        <item name="android:spinnerItemStyle"></item>
+        <item name="android:spinnerDropDownItemStyle"></item-->
+
+        <!-- Set up edge-to-edge configuration for top app bar -->
+        <item name="android:clipToPadding">false</item>
+        <item name="android:clipChildren">false</item>
+    </style>
+
+    <!-- Using in SubSettings page including injected settings page -->
+    <style name="Theme.SubSettingsBase.Expressive" parent="Theme.SettingsBase.Expressive">
+        <!-- Suppress the built-in action bar -->
+        <item name="android:windowActionBar">false</item>
+        <item name="android:windowNoTitle">true</item>
+
+        <!-- Set up edge-to-edge configuration for top app bar -->
+        <item name="android:navigationBarColor">@android:color/transparent</item>
+        <item name="android:statusBarColor">@android:color/transparent</item>
+        <item name="colorControlNormal">?android:attr/colorControlNormal</item>
+
+        <!-- For AndroidX AlertDialog -->
+        <!--item name="alertDialogTheme">@style/Theme.AlertDialog.SettingsLib</item-->
+    </style>
+</resources>
\ No newline at end of file
diff --git a/packages/SettingsLib/SettingsTheme/res/values-v35/themes_preference_expressive.xml b/packages/SettingsLib/SettingsTheme/res/values-v35/themes_preference_expressive.xml
new file mode 100644
index 0000000..41fe225
--- /dev/null
+++ b/packages/SettingsLib/SettingsTheme/res/values-v35/themes_preference_expressive.xml
@@ -0,0 +1,32 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+  -->
+
+<resources>
+    <style name="PreferenceTheme.SettingsLib.Expressive">
+        <item name="checkBoxPreferenceStyle">@style/SettingsLibPreference.CheckBoxPreference.Expressive</item>
+        <item name="dialogPreferenceStyle">@style/SettingsLibPreference.DialogPreference.Expressive</item>
+        <item name="dropdownPreferenceStyle">@style/SettingsLibPreference.DropDown.Expressive</item>
+        <item name="editTextPreferenceStyle">@style/SettingsLibPreference.DialogPreference.EditTextPreference.Expressive</item>
+        <item name="seekBarPreferenceStyle">@style/SettingsLibPreference.SeekBarPreference.Expressive</item>
+        <item name="preferenceCategoryStyle">@style/SettingsLibPreference.Category.Expressive</item>
+        <item name="preferenceScreenStyle">@style/SettingsLibPreference.PreferenceScreen.Expressive</item>
+        <item name="preferenceStyle">@style/SettingsLibPreference.Expressive</item>
+        <item name="switchPreferenceCompatStyle">@style/SettingsLibPreference.SwitchPreferenceCompat.Expressive</item>
+        <item name="preferenceCategoryTitleTextAppearance">@style/TextAppearance.CategoryTitle.SettingsLib</item>
+        <item name="preferenceCategoryTitleTextColor">@color/settingslib_materialColorPrimary</item>
+    </style>
+</resources>
\ No newline at end of file
diff --git a/packages/SettingsLib/res/drawable/ic_media_microphone.xml b/packages/SettingsLib/res/drawable/ic_media_microphone.xml
new file mode 100644
index 0000000..209dea5
--- /dev/null
+++ b/packages/SettingsLib/res/drawable/ic_media_microphone.xml
@@ -0,0 +1,25 @@
+<!--
+    Copyright (C) 2024 The Android Open Source Project
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+         http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+        android:width="24dp"
+        android:height="24dp"
+        android:tint="?attr/colorControlNormal"
+        android:viewportHeight="960"
+        android:viewportWidth="960">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M480,560Q430,560 395,525Q360,490 360,440L360,200Q360,150 395,115Q430,80 480,80Q530,80 565,115Q600,150 600,200L600,440Q600,490 565,525Q530,560 480,560ZM480,320Q480,320 480,320Q480,320 480,320L480,320Q480,320 480,320Q480,320 480,320Q480,320 480,320Q480,320 480,320L480,320Q480,320 480,320Q480,320 480,320ZM440,840L440,717Q336,703 268,624Q200,545 200,440L280,440Q280,523 338.5,581.5Q397,640 480,640Q563,640 621.5,581.5Q680,523 680,440L760,440Q760,545 692,624Q624,703 520,717L520,840L440,840ZM480,480Q497,480 508.5,468.5Q520,457 520,440L520,200Q520,183 508.5,171.5Q497,160 480,160Q463,160 451.5,171.5Q440,183 440,200L440,440Q440,457 451.5,468.5Q463,480 480,480Z" />
+</vector>
\ No newline at end of file
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/InputMediaDevice.java b/packages/SettingsLib/src/com/android/settingslib/media/InputMediaDevice.java
index 9dd2dbb..dae69e6 100644
--- a/packages/SettingsLib/src/com/android/settingslib/media/InputMediaDevice.java
+++ b/packages/SettingsLib/src/com/android/settingslib/media/InputMediaDevice.java
@@ -132,8 +132,7 @@
 
     @VisibleForTesting
     int getDrawableResId() {
-        // TODO(b/357122624): check with UX to obtain the icon for desktop devices.
-        return R.drawable.ic_media_tablet;
+        return R.drawable.ic_media_microphone;
     }
 
     @Override
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputMediaDeviceTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputMediaDeviceTest.java
index bc1ea6c..088d554 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputMediaDeviceTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/media/InputMediaDeviceTest.java
@@ -18,9 +18,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
 import android.content.Context;
 import android.media.AudioDeviceInfo;
 import android.platform.test.flag.junit.SetFlagsRule;
@@ -64,7 +61,7 @@
                         CURRENT_VOLUME,
                         IS_VOLUME_FIXED);
         assertThat(builtinMediaDevice).isNotNull();
-        assertThat(builtinMediaDevice.getDrawableResId()).isEqualTo(R.drawable.ic_media_tablet);
+        assertThat(builtinMediaDevice.getDrawableResId()).isEqualTo(R.drawable.ic_media_microphone);
     }
 
     @Test
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeOverlay.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeOverlay.kt
index f8d0588..8e6cb3f 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeOverlay.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeOverlay.kt
@@ -22,6 +22,7 @@
 import androidx.compose.animation.fadeOut
 import androidx.compose.animation.togetherWith
 import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
@@ -34,6 +35,7 @@
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.battery.BatteryMeterViewController
@@ -41,6 +43,7 @@
 import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.lifecycle.rememberViewModel
+import com.android.systemui.qs.composefragment.ui.GridAnchor
 import com.android.systemui.qs.panels.ui.compose.EditMode
 import com.android.systemui.qs.panels.ui.compose.TileGrid
 import com.android.systemui.qs.ui.viewmodel.QuickSettingsContainerViewModel
@@ -79,16 +82,11 @@
     }
 
     @Composable
-    override fun ContentScope.Content(
-        modifier: Modifier,
-    ) {
+    override fun ContentScope.Content(modifier: Modifier) {
         val viewModel =
             rememberViewModel("QuickSettingsShadeOverlay") { contentViewModelFactory.create() }
 
-        OverlayShade(
-            modifier = modifier,
-            onScrimClicked = viewModel::onScrimClicked,
-        ) {
+        OverlayShade(modifier = modifier, onScrimClicked = viewModel::onScrimClicked) {
             Column {
                 ExpandedShadeHeader(
                     viewModelFactory = viewModel.shadeHeaderViewModelFactory,
@@ -98,40 +96,36 @@
                     modifier = Modifier.padding(QuickSettingsShade.Dimensions.Padding),
                 )
 
-                ShadeBody(
-                    viewModel = viewModel.quickSettingsContainerViewModel,
-                )
+                ShadeBody(viewModel = viewModel.quickSettingsContainerViewModel)
             }
         }
     }
 }
 
 @Composable
-fun ShadeBody(
-    viewModel: QuickSettingsContainerViewModel,
-) {
+fun SceneScope.ShadeBody(viewModel: QuickSettingsContainerViewModel) {
     val isEditing by viewModel.editModeViewModel.isEditing.collectAsStateWithLifecycle()
 
     AnimatedContent(
         targetState = isEditing,
-        transitionSpec = { fadeIn(tween(500)) togetherWith fadeOut(tween(500)) }
+        transitionSpec = { fadeIn(tween(500)) togetherWith fadeOut(tween(500)) },
     ) { editing ->
         if (editing) {
             EditMode(
                 viewModel = viewModel.editModeViewModel,
-                modifier = Modifier.fillMaxWidth().padding(QuickSettingsShade.Dimensions.Padding)
+                modifier = Modifier.fillMaxWidth().padding(QuickSettingsShade.Dimensions.Padding),
             )
         } else {
             QuickSettingsLayout(
                 viewModel = viewModel,
-                modifier = Modifier.sysuiResTag("quick_settings_panel")
+                modifier = Modifier.sysuiResTag("quick_settings_panel"),
             )
         }
     }
 }
 
 @Composable
-private fun QuickSettingsLayout(
+private fun SceneScope.QuickSettingsLayout(
     viewModel: QuickSettingsContainerViewModel,
     modifier: Modifier = Modifier,
 ) {
@@ -143,15 +137,18 @@
         BrightnessSliderContainer(
             viewModel = viewModel.brightnessSliderViewModel,
             modifier =
-                Modifier.fillMaxWidth()
-                    .height(QuickSettingsShade.Dimensions.BrightnessSliderHeight),
+                Modifier.fillMaxWidth().height(QuickSettingsShade.Dimensions.BrightnessSliderHeight),
         )
-        TileGrid(
-            viewModel = viewModel.tileGridViewModel,
-            modifier =
-                Modifier.fillMaxWidth().heightIn(max = QuickSettingsShade.Dimensions.GridMaxHeight),
-            viewModel.editModeViewModel::startEditing,
-        )
+        Box {
+            GridAnchor()
+            TileGrid(
+                viewModel = viewModel.tileGridViewModel,
+                modifier =
+                    Modifier.fillMaxWidth()
+                        .heightIn(max = QuickSettingsShade.Dimensions.GridMaxHeight),
+                viewModel.editModeViewModel::startEditing,
+            )
+        }
     }
 }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
index 7203b61..6f20e70 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
@@ -78,8 +78,6 @@
         Dispatchers.resetMain()
     }
 
-    // For now the state changes at 0.5f expansion. This will change once we implement animation
-    // (and this test will fail)
     @Test
     fun qsExpansionValueChanges_correctExpansionState() =
         with(kosmos) {
@@ -87,18 +85,27 @@
                 val expansionState by collectLastValue(underTest.expansionState)
 
                 underTest.qsExpansionValue = 0f
-                assertThat(expansionState)
-                    .isEqualTo(QSFragmentComposeViewModel.QSExpansionState.QQS)
+                assertThat(expansionState!!.progress).isEqualTo(0f)
 
                 underTest.qsExpansionValue = 0.3f
-                assertThat(expansionState)
-                    .isEqualTo(QSFragmentComposeViewModel.QSExpansionState.QQS)
-
-                underTest.qsExpansionValue = 0.7f
-                assertThat(expansionState).isEqualTo(QSFragmentComposeViewModel.QSExpansionState.QS)
+                assertThat(expansionState!!.progress).isEqualTo(0.3f)
 
                 underTest.qsExpansionValue = 1f
-                assertThat(expansionState).isEqualTo(QSFragmentComposeViewModel.QSExpansionState.QS)
+                assertThat(expansionState!!.progress).isEqualTo(1f)
+            }
+        }
+
+    @Test
+    fun qsExpansionValueChanges_clamped() =
+        with(kosmos) {
+            testScope.testWithinLifecycle {
+                val expansionState by collectLastValue(underTest.expansionState)
+
+                underTest.qsExpansionValue = -1f
+                assertThat(expansionState!!.progress).isEqualTo(0f)
+
+                underTest.qsExpansionValue = 2f
+                assertThat(expansionState!!.progress).isEqualTo(1f)
             }
         }
 
@@ -110,7 +117,7 @@
 
                 testableContext.orCreateTestableResources.addOverride(
                     R.bool.config_use_large_screen_shade_header,
-                    true
+                    true,
                 )
                 fakeConfigurationRepository.onConfigurationChange()
 
@@ -126,7 +133,7 @@
 
                 testableContext.orCreateTestableResources.addOverride(
                     R.bool.config_use_large_screen_shade_header,
-                    false
+                    false,
                 )
                 fakeConfigurationRepository.onConfigurationChange()
 
diff --git a/packages/SystemUI/res/layout/status_bar.xml b/packages/SystemUI/res/layout/status_bar.xml
index 32bcca1..1f4dea9 100644
--- a/packages/SystemUI/res/layout/status_bar.xml
+++ b/packages/SystemUI/res/layout/status_bar.xml
@@ -63,10 +63,12 @@
             <!-- Container that is wrapped around the views on the start half of the status bar.
                  Its width will change with the number of visible children and sub-children.
                  It is useful when we want to know the visible bounds of the content. -->
+            <!-- IMPORTANT: The height of this view *must* be match_parent so that the activity
+                 chips don't get cropped when they appear. See b/302160300 and b/366988057. -->
             <FrameLayout
                 android:id="@+id/status_bar_start_side_content"
                 android:layout_width="wrap_content"
-                android:layout_height="wrap_content"
+                android:layout_height="match_parent"
                 android:layout_gravity="center_vertical|start"
                 android:clipChildren="false">
 
@@ -75,6 +77,8 @@
                 <!-- The alpha of the start side is controlled by PhoneStatusBarTransitions, and the
                      individual views are controlled by StatusBarManager disable flags DISABLE_CLOCK
                      and DISABLE_NOTIFICATION_ICONS, respectively -->
+                <!-- IMPORTANT: The height of this view *must* be match_parent so that the activity
+                 chips don't get cropped when they appear. See b/302160300 and b/366988057. -->
                 <LinearLayout
                     android:id="@+id/status_bar_start_side_except_heads_up"
                     android:layout_height="match_parent"
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/domain/interactor/UdfpsOverlayInteractor.kt b/packages/SystemUI/src/com/android/systemui/biometrics/domain/interactor/UdfpsOverlayInteractor.kt
index bb450c0..18a7739 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/domain/interactor/UdfpsOverlayInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/domain/interactor/UdfpsOverlayInteractor.kt
@@ -83,7 +83,7 @@
 
     /** Sets whether Udfps overlay should handle touches */
     fun setHandleTouches(shouldHandle: Boolean = true) {
-        if (authController.isUltrasonicUdfpsSupported
+        if (authController.isUdfpsSupported
                 && shouldHandle != _shouldHandleTouches.value) {
             fingerprintManager?.setIgnoreDisplayTouches(
                 requestId.value,
diff --git a/packages/SystemUI/src/com/android/systemui/grid/ui/compose/SpannedGrids.kt b/packages/SystemUI/src/com/android/systemui/grid/ui/compose/SpannedGrids.kt
new file mode 100644
index 0000000..62ab18b
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/grid/ui/compose/SpannedGrids.kt
@@ -0,0 +1,321 @@
+/*
+ * 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.grid.ui.compose
+
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.BoxScope
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.remember
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.layout.Layout
+import androidx.compose.ui.layout.Placeable
+import androidx.compose.ui.semantics.CollectionInfo
+import androidx.compose.ui.semantics.CollectionItemInfo
+import androidx.compose.ui.semantics.collectionInfo
+import androidx.compose.ui.semantics.collectionItemInfo
+import androidx.compose.ui.semantics.semantics
+import androidx.compose.ui.unit.Constraints
+import androidx.compose.ui.unit.Dp
+import androidx.compose.ui.unit.LayoutDirection
+import androidx.compose.ui.unit.dp
+import kotlin.math.max
+
+/**
+ * Horizontal (non lazy) grid that supports [spans] for its elements.
+ *
+ * The elements will be laid down vertically first, and then by columns. So assuming LTR layout, it
+ * will be (for a span list `[2, 1, 2, 1, 1, 1, 1, 1]` and 4 rows):
+ * ```
+ * 0  2  5
+ * 0  2  6
+ * 1  3  7
+ *    4
+ * ```
+ *
+ * where repeated numbers show larger span. If an element doesn't fit in a column due to its span,
+ * it will start a new column.
+ *
+ * Elements in [spans] must be in the interval `[1, rows]` ([rows] > 0), and the composables are
+ * associated with the corresponding span based on their index.
+ *
+ * Due to the fact that elements are seen as a linear list that's laid out in a grid, the semantics
+ * represent the collection as a list of elements.
+ */
+@Composable
+fun HorizontalSpannedGrid(
+    rows: Int,
+    columnSpacing: Dp,
+    rowSpacing: Dp,
+    spans: List<Int>,
+    modifier: Modifier = Modifier,
+    composables: @Composable BoxScope.(spanIndex: Int) -> Unit,
+) {
+    SpannedGrid(
+        primarySpaces = rows,
+        crossAxisSpacing = rowSpacing,
+        mainAxisSpacing = columnSpacing,
+        spans = spans,
+        isVertical = false,
+        modifier = modifier,
+        composables = composables,
+    )
+}
+
+/**
+ * Horizontal (non lazy) grid that supports [spans] for its elements.
+ *
+ * The elements will be laid down horizontally first, and then by rows. So assuming LTR layout, it
+ * will be (for a span list `[2, 1, 2, 1, 1, 1, 1, 1]` and 4 columns):
+ * ```
+ * 0  0  1
+ * 2  2  3  4
+ * 5  6  7
+ * ```
+ *
+ * where repeated numbers show larger span. If an element doesn't fit in a row due to its span, it
+ * will start a new row.
+ *
+ * Elements in [spans] must be in the interval `[1, columns]` ([columns] > 0), and the composables
+ * are associated with the corresponding span based on their index.
+ *
+ * Due to the fact that elements are seen as a linear list that's laid out in a grid, the semantics
+ * represent the collection as a list of elements.
+ */
+@Composable
+fun VerticalSpannedGrid(
+    columns: Int,
+    columnSpacing: Dp,
+    rowSpacing: Dp,
+    spans: List<Int>,
+    modifier: Modifier = Modifier,
+    composables: @Composable BoxScope.(spanIndex: Int) -> Unit,
+) {
+    SpannedGrid(
+        primarySpaces = columns,
+        crossAxisSpacing = columnSpacing,
+        mainAxisSpacing = rowSpacing,
+        spans = spans,
+        isVertical = true,
+        modifier = modifier,
+        composables = composables,
+    )
+}
+
+@Composable
+private fun SpannedGrid(
+    primarySpaces: Int,
+    crossAxisSpacing: Dp,
+    mainAxisSpacing: Dp,
+    spans: List<Int>,
+    isVertical: Boolean,
+    modifier: Modifier = Modifier,
+    composables: @Composable BoxScope.(spanIndex: Int) -> Unit,
+) {
+    val crossAxisArrangement = Arrangement.spacedBy(crossAxisSpacing)
+    spans.forEachIndexed { index, span ->
+        check(span in 1..primarySpaces) {
+            "Span out of bounds. Span at index $index has value of $span which is outside of the " +
+                "expected rance of [1, $primarySpaces]"
+        }
+    }
+
+    if (isVertical) {
+        check(crossAxisSpacing >= 0.dp) { "Negative columnSpacing $crossAxisSpacing" }
+        check(mainAxisSpacing >= 0.dp) { "Negative rowSpacing $mainAxisSpacing" }
+    } else {
+        check(mainAxisSpacing >= 0.dp) { "Negative columnSpacing $mainAxisSpacing" }
+        check(crossAxisSpacing >= 0.dp) { "Negative rowSpacing $crossAxisSpacing" }
+    }
+
+    val totalMainAxisGroups: Int =
+        remember(primarySpaces, spans) {
+            var currentAccumulated = 0
+            var groups = 1
+            spans.forEach { span ->
+                if (currentAccumulated + span <= primarySpaces) {
+                    currentAccumulated += span
+                } else {
+                    groups += 1
+                    currentAccumulated = span
+                }
+            }
+            groups
+        }
+
+    val slotPositionsAndSizesCache = remember {
+        object {
+            var sizes = IntArray(0)
+            var positions = IntArray(0)
+        }
+    }
+
+    Layout(
+        {
+            (0 until spans.size).map { spanIndex ->
+                Box(
+                    Modifier.semantics {
+                        collectionItemInfo =
+                            if (isVertical) {
+                                CollectionItemInfo(spanIndex, 1, 0, 1)
+                            } else {
+                                CollectionItemInfo(0, 1, spanIndex, 1)
+                            }
+                    }
+                ) {
+                    composables(spanIndex)
+                }
+            }
+        },
+        modifier.semantics { collectionInfo = CollectionInfo(spans.size, 1) },
+    ) { measurables, constraints ->
+        check(measurables.size == spans.size)
+        val crossAxisSize = if (isVertical) constraints.maxWidth else constraints.maxHeight
+        check(crossAxisSize != Constraints.Infinity) { "Width must be constrained" }
+        if (slotPositionsAndSizesCache.sizes.size != primarySpaces) {
+            slotPositionsAndSizesCache.sizes = IntArray(primarySpaces)
+            slotPositionsAndSizesCache.positions = IntArray(primarySpaces)
+        }
+        calculateCellsCrossAxisSize(
+            crossAxisSize,
+            primarySpaces,
+            crossAxisSpacing.roundToPx(),
+            slotPositionsAndSizesCache.sizes,
+        )
+        val cellSizesInCrossAxis = slotPositionsAndSizesCache.sizes
+
+        // with is needed because of the double receiver (Density, Arrangement).
+        with(crossAxisArrangement) {
+            arrange(
+                crossAxisSize,
+                slotPositionsAndSizesCache.sizes,
+                LayoutDirection.Ltr,
+                slotPositionsAndSizesCache.positions,
+            )
+        }
+        val startPositions = slotPositionsAndSizesCache.positions
+
+        val mainAxisSpacingPx = mainAxisSpacing.roundToPx()
+        val mainAxisTotalGaps = (totalMainAxisGroups - 1) * mainAxisSpacingPx
+        val mainAxisSize = if (isVertical) constraints.maxHeight else constraints.maxWidth
+        val mainAxisElementConstraint =
+            if (mainAxisSize == Constraints.Infinity) {
+                Constraints.Infinity
+            } else {
+                max(0, (mainAxisSize - mainAxisTotalGaps) / totalMainAxisGroups)
+            }
+
+        val mainAxisSizes = IntArray(totalMainAxisGroups) { 0 }
+
+        var currentSlot = 0
+        var mainAxisGroup = 0
+        val placeables =
+            measurables.mapIndexed { index, measurable ->
+                val span = spans[index]
+                if (currentSlot + span > primarySpaces) {
+                    currentSlot = 0
+                    mainAxisGroup += 1
+                }
+                val crossAxisConstraint =
+                    calculateWidth(cellSizesInCrossAxis, startPositions, currentSlot, span)
+                PlaceResult(
+                        measurable.measure(
+                            makeConstraint(
+                                isVertical,
+                                mainAxisElementConstraint,
+                                crossAxisConstraint,
+                            )
+                        ),
+                        currentSlot,
+                        mainAxisGroup,
+                    )
+                    .also {
+                        currentSlot += span
+                        mainAxisSizes[mainAxisGroup] =
+                            max(
+                                mainAxisSizes[mainAxisGroup],
+                                if (isVertical) it.placeable.height else it.placeable.width,
+                            )
+                    }
+            }
+
+        val mainAxisTotalSize = mainAxisTotalGaps + mainAxisSizes.sum()
+        val mainAxisStartingPoints =
+            mainAxisSizes.runningFold(0) { acc, value -> acc + value + mainAxisSpacingPx }
+        val height = if (isVertical) mainAxisTotalSize else crossAxisSize
+        val width = if (isVertical) crossAxisSize else mainAxisTotalSize
+
+        layout(width, height) {
+            placeables.forEach { (placeable, slot, mainAxisGroup) ->
+                val x =
+                    if (isVertical) {
+                        startPositions[slot]
+                    } else {
+                        mainAxisStartingPoints[mainAxisGroup]
+                    }
+                val y =
+                    if (isVertical) {
+                        mainAxisStartingPoints[mainAxisGroup]
+                    } else {
+                        startPositions[slot]
+                    }
+                placeable.placeRelative(x, y)
+            }
+        }
+    }
+}
+
+fun makeConstraint(isVertical: Boolean, mainAxisSize: Int, crossAxisSize: Int): Constraints {
+    return if (isVertical) {
+        Constraints(maxHeight = mainAxisSize, minWidth = crossAxisSize, maxWidth = crossAxisSize)
+    } else {
+        Constraints(maxWidth = mainAxisSize, minHeight = crossAxisSize, maxHeight = crossAxisSize)
+    }
+}
+
+private fun calculateWidth(sizes: IntArray, positions: IntArray, startSlot: Int, span: Int): Int {
+    val crossAxisSize =
+        if (span == 1) {
+                sizes[startSlot]
+            } else {
+                val endSlot = startSlot + span - 1
+                positions[endSlot] + sizes[endSlot] - positions[startSlot]
+            }
+            .coerceAtLeast(0)
+    return crossAxisSize
+}
+
+private fun calculateCellsCrossAxisSize(
+    gridSize: Int,
+    slotCount: Int,
+    spacingPx: Int,
+    outArray: IntArray,
+) {
+    check(outArray.size == slotCount)
+    val gridSizeWithoutSpacing = gridSize - spacingPx * (slotCount - 1)
+    val slotSize = gridSizeWithoutSpacing / slotCount
+    val remainingPixels = gridSizeWithoutSpacing % slotCount
+    outArray.indices.forEach { index ->
+        outArray[index] = slotSize + if (index < remainingPixels) 1 else 0
+    }
+}
+
+private data class PlaceResult(
+    val placeable: Placeable,
+    val slotIndex: Int,
+    val mainAxisGroup: Int,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt
index 44460ed..eff5fc0 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt
@@ -75,7 +75,9 @@
      * [NoteTaskController], ensure custom actions can be triggered (i.e., keyboard shortcut).
      */
     private fun initializeHandleSystemKey() {
-        commandQueue.addCallback(callbacks)
+        if (!useKeyGestureEventHandler()) {
+            commandQueue.addCallback(callbacks)
+        }
     }
 
     /**
@@ -130,6 +132,11 @@
             InputManager.KeyGestureEventHandler {
 
             override fun handleSystemKey(key: KeyEvent) {
+                if (useKeyGestureEventHandler()) {
+                    throw IllegalStateException(
+                        "handleSystemKey must not be used when KeyGestureEventHandler is used"
+                    )
+                }
                 key.toNoteTaskEntryPointOrNull()?.let(controller::showNoteTask)
             }
 
@@ -151,13 +158,13 @@
 
             override fun handleKeyGestureEvent(
                 event: KeyGestureEvent,
-                focusedToken: IBinder?
+                focusedToken: IBinder?,
             ): Boolean {
                 return [email protected](event)
             }
 
             override fun isKeyGestureSupported(gestureType: Int): Boolean {
-                return [email protected](gestureType);
+                return [email protected](gestureType)
             }
         }
 
@@ -209,8 +216,20 @@
             "handleKeyGestureEvent: Received OPEN_NOTES gesture event from keycodes: " +
                 event.keycodes.contentToString()
         }
-        backgroundExecutor.execute { controller.showNoteTask(KEYBOARD_SHORTCUT) }
-        return true
+        if (
+            event.keycodes.contains(KEYCODE_N) &&
+                event.hasModifiers(KeyEvent.META_CTRL_ON or KeyEvent.META_META_ON)
+        ) {
+            debugLog { "Note task triggered by keyboard shortcut" }
+            backgroundExecutor.execute { controller.showNoteTask(KEYBOARD_SHORTCUT) }
+            return true
+        }
+        if (event.keycodes.size == 1 && event.keycodes[0] == KEYCODE_STYLUS_BUTTON_TAIL) {
+            debugLog { "Note task triggered by stylus tail button" }
+            backgroundExecutor.execute { controller.showNoteTask(TAIL_BUTTON) }
+            return true
+        }
+        return false
     }
 
     private fun isKeyGestureSupported(gestureType: Int): Boolean {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
index af167d4..c174038 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
@@ -26,7 +26,6 @@
 import androidx.activity.OnBackPressedDispatcher
 import androidx.activity.OnBackPressedDispatcherOwner
 import androidx.activity.setViewTreeOnBackPressedDispatcherOwner
-import androidx.compose.animation.AnimatedContent
 import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
@@ -38,10 +37,14 @@
 import androidx.compose.foundation.layout.windowInsetsPadding
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableFloatStateOf
 import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.graphicsLayer
 import androidx.compose.ui.layout.onPlaced
 import androidx.compose.ui.layout.onSizeChanged
 import androidx.compose.ui.layout.positionInRoot
@@ -51,11 +54,18 @@
 import androidx.compose.ui.semantics.CustomAccessibilityAction
 import androidx.compose.ui.semantics.customActions
 import androidx.compose.ui.semantics.semantics
+import androidx.compose.ui.unit.dp
 import androidx.compose.ui.unit.round
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
+import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
+import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.SceneScope
+import com.android.compose.animation.scene.SceneTransitionLayout
+import com.android.compose.animation.scene.content.state.TransitionState
+import com.android.compose.animation.scene.transitions
 import com.android.compose.modifiers.height
 import com.android.compose.modifiers.padding
 import com.android.compose.modifiers.thenIf
@@ -70,11 +80,17 @@
 import com.android.systemui.media.dagger.MediaModule.QUICK_QS_PANEL
 import com.android.systemui.plugins.qs.QS
 import com.android.systemui.plugins.qs.QSContainerController
+import com.android.systemui.qs.composefragment.SceneKeys.QuickQuickSettings
+import com.android.systemui.qs.composefragment.SceneKeys.QuickSettings
+import com.android.systemui.qs.composefragment.SceneKeys.toIdleSceneKey
 import com.android.systemui.qs.composefragment.ui.notificationScrimClip
+import com.android.systemui.qs.composefragment.ui.quickQuickSettingsToQuickSettings
 import com.android.systemui.qs.composefragment.viewmodel.QSFragmentComposeViewModel
 import com.android.systemui.qs.flags.QSComposeFragment
 import com.android.systemui.qs.footer.ui.compose.FooterActions
 import com.android.systemui.qs.panels.ui.compose.QuickQuickSettings
+import com.android.systemui.qs.shared.ui.ElementKeys
+import com.android.systemui.qs.ui.composable.QuickSettingsShade
 import com.android.systemui.qs.ui.composable.QuickSettingsTheme
 import com.android.systemui.qs.ui.composable.ShadeBody
 import com.android.systemui.res.R
@@ -86,11 +102,13 @@
 import java.util.function.Consumer
 import javax.inject.Inject
 import javax.inject.Named
+import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
 @SuppressLint("ValidFragment")
@@ -166,33 +184,48 @@
             setContent {
                 PlatformTheme {
                     val visible by viewModel.qsVisible.collectAsStateWithLifecycle()
-                    val qsState by viewModel.expansionState.collectAsStateWithLifecycle()
 
                     AnimatedVisibility(
                         visible = visible,
                         modifier =
-                            Modifier.windowInsetsPadding(WindowInsets.navigationBars).thenIf(
-                                notificationScrimClippingParams.isEnabled
-                            ) {
-                                Modifier.notificationScrimClip(
-                                    notificationScrimClippingParams.leftInset,
-                                    notificationScrimClippingParams.top,
-                                    notificationScrimClippingParams.rightInset,
-                                    notificationScrimClippingParams.bottom,
-                                    notificationScrimClippingParams.radius,
-                                )
-                            },
+                            Modifier.windowInsetsPadding(WindowInsets.navigationBars)
+                                .thenIf(notificationScrimClippingParams.isEnabled) {
+                                    Modifier.notificationScrimClip(
+                                        notificationScrimClippingParams.leftInset,
+                                        notificationScrimClippingParams.top,
+                                        notificationScrimClippingParams.rightInset,
+                                        notificationScrimClippingParams.bottom,
+                                        notificationScrimClippingParams.radius,
+                                    )
+                                }
+                                .graphicsLayer { elevation = 4.dp.toPx() },
                     ) {
-                        AnimatedContent(targetState = qsState) {
-                            when (it) {
-                                QSFragmentComposeViewModel.QSExpansionState.QQS -> {
-                                    QuickQuickSettingsElement()
-                                }
-                                QSFragmentComposeViewModel.QSExpansionState.QS -> {
-                                    QuickSettingsElement()
-                                }
-                                else -> {}
-                            }
+                        val sceneState = remember {
+                            MutableSceneTransitionLayoutState(
+                                viewModel.expansionState.value.toIdleSceneKey(),
+                                transitions =
+                                    transitions {
+                                        from(QuickQuickSettings, QuickSettings) {
+                                            quickQuickSettingsToQuickSettings()
+                                        }
+                                    },
+                            )
+                        }
+
+                        LaunchedEffect(Unit) {
+                            synchronizeQsState(
+                                sceneState,
+                                viewModel.expansionState.map { it.progress },
+                            )
+                        }
+
+                        SceneTransitionLayout(
+                            state = sceneState,
+                            modifier = Modifier.fillMaxSize(),
+                        ) {
+                            scene(QuickSettings) { QuickSettingsElement() }
+
+                            scene(QuickQuickSettings) { QuickQuickSettingsElement() }
                         }
                     }
                 }
@@ -420,7 +453,7 @@
     }
 
     @Composable
-    private fun QuickQuickSettingsElement() {
+    private fun SceneScope.QuickQuickSettingsElement() {
         val qqsPadding by viewModel.qqsHeaderHeight.collectAsStateWithLifecycle()
         val bottomPadding = dimensionResource(id = R.dimen.qqs_layout_padding_bottom)
         DisposableEffect(Unit) {
@@ -450,8 +483,15 @@
                         viewModel = viewModel.containerViewModel.quickQuickSettingsViewModel,
                         modifier =
                             Modifier.collapseExpandSemanticAction(
-                                stringResource(id = R.string.accessibility_quick_settings_expand)
-                            ),
+                                    stringResource(
+                                        id = R.string.accessibility_quick_settings_expand
+                                    )
+                                )
+                                .padding(
+                                    horizontal = {
+                                        QuickSettingsShade.Dimensions.Padding.roundToPx()
+                                    }
+                                ),
                     )
                 }
             }
@@ -460,7 +500,7 @@
     }
 
     @Composable
-    private fun QuickSettingsElement() {
+    private fun SceneScope.QuickSettingsElement() {
         val qqsPadding by viewModel.qqsHeaderHeight.collectAsStateWithLifecycle()
         val qsExtraPadding = dimensionResource(R.dimen.qs_panel_padding_top)
         Column(
@@ -471,7 +511,10 @@
         ) {
             val qsEnabled by viewModel.qsEnabled.collectAsStateWithLifecycle()
             if (qsEnabled) {
-                Box(modifier = Modifier.fillMaxSize().weight(1f)) {
+                Box(
+                    modifier =
+                        Modifier.element(ElementKeys.QuickSettingsContent).fillMaxSize().weight(1f)
+                ) {
                     Column {
                         Spacer(
                             modifier = Modifier.height { qqsPadding + qsExtraPadding.roundToPx() }
@@ -483,7 +526,9 @@
                     FooterActions(
                         viewModel = viewModel.footerActionsViewModel,
                         qsVisibilityLifecycleOwner = this@QSFragmentCompose,
-                        modifier = Modifier.sysuiResTag("qs_footer_actions"),
+                        modifier =
+                            Modifier.sysuiResTag("qs_footer_actions")
+                                .element(ElementKeys.FooterActions),
                     )
                 }
             }
@@ -590,3 +635,85 @@
             return currentId++
         }
     }
+
+object SceneKeys {
+    val QuickQuickSettings = SceneKey("QuickQuickSettingsScene")
+    val QuickSettings = SceneKey("QuickSettingsScene")
+
+    fun QSFragmentComposeViewModel.QSExpansionState.toIdleSceneKey(): SceneKey {
+        return when {
+            progress < 0.5f -> QuickQuickSettings
+            else -> QuickSettings
+        }
+    }
+}
+
+suspend fun synchronizeQsState(state: MutableSceneTransitionLayoutState, expansion: Flow<Float>) {
+    coroutineScope {
+        val animationScope = this
+
+        var currentTransition: ExpansionTransition? = null
+
+        fun snapTo(scene: SceneKey) {
+            state.snapToScene(scene)
+            currentTransition = null
+        }
+
+        expansion.collectLatest { progress ->
+            when (progress) {
+                0f -> snapTo(QuickQuickSettings)
+                1f -> snapTo(QuickSettings)
+                else -> {
+                    val transition = currentTransition
+                    if (transition != null) {
+                        transition.progress = progress
+                        return@collectLatest
+                    }
+
+                    val newTransition =
+                        ExpansionTransition(progress).also { currentTransition = it }
+                    state.startTransitionImmediately(
+                        animationScope = animationScope,
+                        transition = newTransition,
+                    )
+                }
+            }
+        }
+    }
+}
+
+private class ExpansionTransition(currentProgress: Float) :
+    TransitionState.Transition.ChangeScene(
+        fromScene = QuickQuickSettings,
+        toScene = QuickSettings,
+    ) {
+    override val currentScene: SceneKey
+        get() {
+            // This should return the logical scene. If the QS STLState is only driven by
+            // synchronizeQSState() then it probably does not matter which one we return, this is
+            // only used to compute the current user actions of a STL.
+            return QuickQuickSettings
+        }
+
+    override var progress: Float by mutableFloatStateOf(currentProgress)
+
+    override val progressVelocity: Float
+        get() = 0f
+
+    override val isInitiatedByUserInput: Boolean
+        get() = true
+
+    override val isUserInputOngoing: Boolean
+        get() = true
+
+    private val finishCompletable = CompletableDeferred<Unit>()
+
+    override suspend fun run() {
+        // This transition runs until it is interrupted by another one.
+        finishCompletable.await()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        finishCompletable.complete(Unit)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/ui/FromQuickQuickSettingsToQuickSettings.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/ui/FromQuickQuickSettingsToQuickSettings.kt
new file mode 100644
index 0000000..1514986
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/ui/FromQuickQuickSettingsToQuickSettings.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.composefragment.ui
+
+import com.android.compose.animation.scene.TransitionBuilder
+import com.android.systemui.qs.shared.ui.ElementKeys
+
+fun TransitionBuilder.quickQuickSettingsToQuickSettings() {
+
+    fractionRange(start = 0.5f) { fade(ElementKeys.QuickSettingsContent) }
+
+    fractionRange(start = 0.9f) { fade(ElementKeys.FooterActions) }
+
+    anchoredTranslate(ElementKeys.QuickSettingsContent, ElementKeys.GridAnchor)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/ui/GridAnchor.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/ui/GridAnchor.kt
new file mode 100644
index 0000000..f0f46d3
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/ui/GridAnchor.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.composefragment.ui
+
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.SceneScope
+import com.android.systemui.qs.shared.ui.ElementKeys
+
+/**
+ * This composable is used at the start of the tiles in QQS and QS to anchor the expansion and be
+ * able to have relative anchor translation of elements that appear in QS.
+ */
+@Composable
+fun SceneScope.GridAnchor(modifier: Modifier = Modifier) {
+    // The size of this anchor does not matter, as the tiles don't change size on expansion.
+    Spacer(modifier.element(ElementKeys.GridAnchor))
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
index 7ab11d2..7300ee1 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
@@ -147,7 +147,7 @@
             .stateIn(
                 lifecycleScope,
                 SharingStarted.WhileSubscribed(),
-                disableFlagsRepository.disableFlags.value.isQuickSettingsEnabled()
+                disableFlagsRepository.disableFlags.value.isQuickSettingsEnabled(),
             )
 
     private val _showCollapsedOnKeyguard = MutableStateFlow(false)
@@ -213,19 +213,11 @@
         }
 
     val expansionState: StateFlow<QSExpansionState> =
-        combine(
-                _stackScrollerOverscrolling,
-                _qsExpanded,
-                _qsExpansion,
-            ) { args: Array<Any> ->
+        combine(_stackScrollerOverscrolling, _qsExpanded, _qsExpansion) { args: Array<Any> ->
                 val expansion = args[2] as Float
-                if (expansion > 0.5f) {
-                    QSExpansionState.QS
-                } else {
-                    QSExpansionState.QQS
-                }
+                QSExpansionState(expansion.coerceIn(0f, 1f))
             }
-            .stateIn(lifecycleScope, SharingStarted.WhileSubscribed(), QSExpansionState.QQS)
+            .stateIn(lifecycleScope, SharingStarted.WhileSubscribed(), QSExpansionState(0f))
 
     /**
      * Accessibility action for collapsing/expanding QS. The provided runnable is responsible for
@@ -262,13 +254,6 @@
         fun create(lifecycleScope: LifecycleCoroutineScope): QSFragmentComposeViewModel
     }
 
-    sealed interface QSExpansionState {
-        data object QQS : QSExpansionState
-
-        data object QS : QSExpansionState
-
-        @JvmInline value class Expanding(val progress: Float) : QSExpansionState
-
-        @JvmInline value class Collapsing(val progress: Float) : QSExpansionState
-    }
+    // In the future, this will have other relevant elements like squishiness.
+    data class QSExpansionState(@FloatRange(0.0, 1.0) val progress: Float)
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
index fd276c2..0c02b40 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
@@ -18,6 +18,7 @@
 
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.qs.panels.shared.model.SizedTile
 import com.android.systemui.qs.panels.shared.model.TileRow
 import com.android.systemui.qs.panels.ui.viewmodel.EditTileViewModel
@@ -27,7 +28,7 @@
 /** A layout of tiles, indicating how they should be composed when showing in QS or in edit mode. */
 interface GridLayout {
     @Composable
-    fun TileGrid(
+    fun SceneScope.TileGrid(
         tiles: List<TileViewModel>,
         modifier: Modifier,
         editModeStart: () -> Unit,
@@ -66,7 +67,7 @@
          */
         fun splitInRows(
             tiles: List<SizedTile<TileViewModel>>,
-            columns: Int
+            columns: Int,
         ): List<List<SizedTile<TileViewModel>>> {
             val row = TileRow<TileViewModel>(columns)
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/PaginatedGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/PaginatedGridLayout.kt
index 08a56bf..083f529 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/PaginatedGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/PaginatedGridLayout.kt
@@ -39,6 +39,7 @@
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.qs.panels.dagger.PaginatedBaseLayoutType
 import com.android.systemui.qs.panels.ui.compose.PaginatedGridLayout.Dimensions.FooterHeight
@@ -55,7 +56,7 @@
     @PaginatedBaseLayoutType private val delegateGridLayout: PaginatableGridLayout,
 ) : GridLayout by delegateGridLayout {
     @Composable
-    override fun TileGrid(
+    override fun SceneScope.TileGrid(
         tiles: List<TileViewModel>,
         modifier: Modifier,
         editModeStart: () -> Unit,
@@ -85,16 +86,16 @@
             ) {
                 val page = pages[it]
 
-                delegateGridLayout.TileGrid(tiles = page, modifier = Modifier, editModeStart = {})
+                with(delegateGridLayout) {
+                    TileGrid(tiles = page, modifier = Modifier, editModeStart = {})
+                }
             }
-            Box(
-                modifier = Modifier.height(FooterHeight).fillMaxWidth(),
-            ) {
+            Box(modifier = Modifier.height(FooterHeight).fillMaxWidth()) {
                 PagerDots(
                     pagerState = pagerState,
                     activeColor = MaterialTheme.colorScheme.primary,
                     nonActiveColor = MaterialTheme.colorScheme.surfaceVariant,
-                    modifier = Modifier.align(Alignment.Center)
+                    modifier = Modifier.align(Alignment.Center),
                 )
                 CompositionLocalProvider(value = LocalContentColor provides Color.White) {
                     IconButton(
@@ -103,7 +104,7 @@
                     ) {
                         Icon(
                             imageVector = Icons.Default.Edit,
-                            contentDescription = stringResource(id = R.string.qs_edit)
+                            contentDescription = stringResource(id = R.string.qs_edit),
                         )
                     }
                 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
index f4acbec..8998a7f 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
@@ -16,21 +16,28 @@
 
 package com.android.systemui.qs.panels.ui.compose
 
-import androidx.compose.foundation.lazy.grid.GridCells
-import androidx.compose.foundation.lazy.grid.GridItemSpan
+import androidx.compose.foundation.layout.Box
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.res.dimensionResource
 import androidx.compose.ui.util.fastMap
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.compose.modifiers.sysuiResTag
+import com.android.systemui.grid.ui.compose.VerticalSpannedGrid
+import com.android.systemui.qs.composefragment.ui.GridAnchor
 import com.android.systemui.qs.panels.ui.compose.infinitegrid.Tile
-import com.android.systemui.qs.panels.ui.compose.infinitegrid.TileLazyGrid
 import com.android.systemui.qs.panels.ui.viewmodel.QuickQuickSettingsViewModel
+import com.android.systemui.qs.shared.ui.ElementKeys.toElementKey
+import com.android.systemui.res.R
 
 @Composable
-fun QuickQuickSettings(viewModel: QuickQuickSettingsViewModel, modifier: Modifier = Modifier) {
+fun SceneScope.QuickQuickSettings(
+    viewModel: QuickQuickSettingsViewModel,
+    modifier: Modifier = Modifier,
+) {
     val sizedTiles by
         viewModel.tileViewModels.collectAsStateWithLifecycle(initialValue = emptyList())
     val tiles = sizedTiles.fastMap { it.tile }
@@ -41,20 +48,20 @@
         onDispose { tiles.forEach { it.stopListening(token) } }
     }
     val columns by viewModel.columns.collectAsStateWithLifecycle()
-
-    TileLazyGrid(
-        modifier = modifier.sysuiResTag("qqs_tile_layout"),
-        columns = GridCells.Fixed(columns),
-    ) {
-        items(
-            sizedTiles.size,
-            key = { index -> sizedTiles[index].tile.spec.spec },
-            span = { index -> GridItemSpan(sizedTiles[index].width) },
-        ) { index ->
+    Box(modifier = modifier) {
+        GridAnchor()
+        VerticalSpannedGrid(
+            columns = columns,
+            columnSpacing = dimensionResource(R.dimen.qs_tile_margin_horizontal),
+            rowSpacing = dimensionResource(R.dimen.qs_tile_margin_vertical),
+            spans = sizedTiles.fastMap { it.width },
+            modifier = Modifier.sysuiResTag("qqs_tile_layout"),
+        ) { spanIndex ->
+            val it = sizedTiles[spanIndex]
             Tile(
-                tile = sizedTiles[index].tile,
-                iconOnly = sizedTiles[index].isIcon,
-                modifier = Modifier,
+                tile = it.tile,
+                iconOnly = it.isIcon,
+                modifier = Modifier.element(it.tile.spec.toElementKey(spanIndex)),
             )
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
index 8c57d41..1a5297b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/TileGrid.kt
@@ -20,16 +20,17 @@
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.qs.panels.ui.viewmodel.TileGridViewModel
 
 @Composable
-fun TileGrid(
+fun SceneScope.TileGrid(
     viewModel: TileGridViewModel,
     modifier: Modifier = Modifier,
-    editModeStart: () -> Unit
+    editModeStart: () -> Unit,
 ) {
     val gridLayout by viewModel.gridLayout.collectAsStateWithLifecycle()
     val tiles by viewModel.tileViewModels.collectAsStateWithLifecycle(emptyList())
 
-    gridLayout.TileGrid(tiles, modifier, editModeStart)
+    with(gridLayout) { TileGrid(tiles, modifier, editModeStart) }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
index 4946c01..8a96065 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/infinitegrid/InfiniteGridLayout.kt
@@ -16,15 +16,17 @@
 
 package com.android.systemui.qs.panels.ui.compose.infinitegrid
 
-import androidx.compose.foundation.lazy.grid.GridCells
-import androidx.compose.foundation.lazy.grid.GridItemSpan
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.res.dimensionResource
+import androidx.compose.ui.util.fastMap
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.grid.ui.compose.VerticalSpannedGrid
 import com.android.systemui.qs.panels.shared.model.SizedTileImpl
 import com.android.systemui.qs.panels.ui.compose.PaginatableGridLayout
 import com.android.systemui.qs.panels.ui.compose.rememberEditListState
@@ -33,6 +35,8 @@
 import com.android.systemui.qs.panels.ui.viewmodel.IconTilesViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
 import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.qs.shared.ui.ElementKeys.toElementKey
+import com.android.systemui.res.R
 import javax.inject.Inject
 
 @SysUISingleton
@@ -44,7 +48,7 @@
 ) : PaginatableGridLayout {
 
     @Composable
-    override fun TileGrid(
+    override fun SceneScope.TileGrid(
         tiles: List<TileViewModel>,
         modifier: Modifier,
         editModeStart: () -> Unit,
@@ -57,15 +61,18 @@
         val columns by gridSizeViewModel.columns.collectAsStateWithLifecycle()
         val sizedTiles = tiles.map { SizedTileImpl(it, it.spec.width()) }
 
-        TileLazyGrid(modifier = modifier, columns = GridCells.Fixed(columns)) {
-            items(sizedTiles.size, span = { index -> GridItemSpan(sizedTiles[index].width) }) {
-                index ->
-                Tile(
-                    tile = sizedTiles[index].tile,
-                    iconOnly = iconTilesViewModel.isIconTile(sizedTiles[index].tile.spec),
-                    modifier = Modifier,
-                )
-            }
+        VerticalSpannedGrid(
+            columns = columns,
+            columnSpacing = dimensionResource(R.dimen.qs_tile_margin_horizontal),
+            rowSpacing = dimensionResource(R.dimen.qs_tile_margin_vertical),
+            spans = sizedTiles.fastMap { it.width },
+        ) { spanIndex ->
+            val it = sizedTiles[spanIndex]
+            Tile(
+                tile = it.tile,
+                iconOnly = iconTilesViewModel.isIconTile(it.tile.spec),
+                modifier = Modifier.element(it.tile.spec.toElementKey(spanIndex)),
+            )
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/shared/ui/ElementKeys.kt b/packages/SystemUI/src/com/android/systemui/qs/shared/ui/ElementKeys.kt
new file mode 100644
index 0000000..625459d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/shared/ui/ElementKeys.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.shared.ui
+
+import com.android.compose.animation.scene.ElementKey
+import com.android.systemui.qs.pipeline.shared.TileSpec
+
+/** Element keys to be used by the compose implementation of QS for animations. */
+object ElementKeys {
+    val QuickSettingsContent = ElementKey("QuickSettingsContent")
+    val GridAnchor = ElementKey("QuickSettingsGridAnchor")
+    val FooterActions = ElementKey("FooterActions")
+
+    class TileElementKey(spec: TileSpec, val position: Int) : ElementKey(spec.spec, spec.spec)
+
+    fun TileSpec.toElementKey(positionInGrid: Int) = TileElementKey(this, positionInGrid)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
index f69b0cb..7724abd4 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotController.java
@@ -505,8 +505,8 @@
             return;
         }
         // delay starting scroll capture to make sure scrim is up before the app moves
-        mViewProxy.prepareScrollingTransition(response, mScreenBitmap, newScreenshot,
-                mScreenshotTakenInPortrait, () -> executeBatchScrollCapture(response, owner));
+        mViewProxy.prepareScrollingTransition(response, newScreenshot, mScreenshotTakenInPortrait,
+                () -> executeBatchScrollCapture(response, owner));
     }
 
     private void executeBatchScrollCapture(ScrollCaptureResponse response, UserHandle owner) {
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
deleted file mode 100644
index fe58bc9..0000000
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ /dev/null
@@ -1,663 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.screenshot;
-
-import static android.content.res.Configuration.ORIENTATION_PORTRAIT;
-
-import static com.android.systemui.screenshot.LogConfig.DEBUG_ANIM;
-import static com.android.systemui.screenshot.LogConfig.DEBUG_CALLBACK;
-import static com.android.systemui.screenshot.LogConfig.DEBUG_INPUT;
-import static com.android.systemui.screenshot.LogConfig.DEBUG_UI;
-import static com.android.systemui.screenshot.LogConfig.DEBUG_WINDOW;
-import static com.android.systemui.screenshot.LogConfig.logTag;
-import static com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER;
-import static com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_INTERACTION_TIMEOUT;
-
-import android.animation.Animator;
-import android.animation.AnimatorListenerAdapter;
-import android.annotation.NonNull;
-import android.annotation.Nullable;
-import android.content.BroadcastReceiver;
-import android.content.Context;
-import android.content.Intent;
-import android.content.IntentFilter;
-import android.content.pm.ActivityInfo;
-import android.content.res.Configuration;
-import android.graphics.Bitmap;
-import android.graphics.Insets;
-import android.graphics.Rect;
-import android.net.Uri;
-import android.os.Process;
-import android.os.UserHandle;
-import android.os.UserManager;
-import android.provider.Settings;
-import android.util.DisplayMetrics;
-import android.util.Log;
-import android.view.Display;
-import android.view.ScrollCaptureResponse;
-import android.view.ViewRootImpl;
-import android.view.WindowManager;
-import android.widget.Toast;
-import android.window.WindowContext;
-
-import com.android.internal.logging.UiEventLogger;
-import com.android.settingslib.applications.InterestingConfigChanges;
-import com.android.systemui.broadcast.BroadcastDispatcher;
-import com.android.systemui.broadcast.BroadcastSender;
-import com.android.systemui.clipboardoverlay.ClipboardOverlayController;
-import com.android.systemui.dagger.qualifiers.Main;
-import com.android.systemui.flags.FeatureFlags;
-import com.android.systemui.res.R;
-import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback;
-import com.android.systemui.screenshot.scroll.ScrollCaptureExecutor;
-import com.android.systemui.util.Assert;
-
-import com.google.common.util.concurrent.ListenableFuture;
-
-import dagger.assisted.Assisted;
-import dagger.assisted.AssistedFactory;
-import dagger.assisted.AssistedInject;
-
-import kotlin.Unit;
-
-import java.util.UUID;
-import java.util.concurrent.Executor;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.function.Consumer;
-
-import javax.inject.Provider;
-
-/**
- * Controls the state and flow for screenshots.
- */
-public class ScreenshotController implements InteractiveScreenshotHandler {
-    private static final String TAG = logTag(ScreenshotController.class);
-
-    // From WizardManagerHelper.java
-    private static final String SETTINGS_SECURE_USER_SETUP_COMPLETE = "user_setup_complete";
-
-    static final int SCREENSHOT_CORNER_DEFAULT_TIMEOUT_MILLIS = 6000;
-
-    private final WindowContext mContext;
-    private final FeatureFlags mFlags;
-    private final ScreenshotShelfViewProxy mViewProxy;
-    private final ScreenshotNotificationsController mNotificationsController;
-    private final ScreenshotSmartActions mScreenshotSmartActions;
-    private final UiEventLogger mUiEventLogger;
-    private final ImageExporter mImageExporter;
-    private final ImageCapture mImageCapture;
-    private final Executor mMainExecutor;
-    private final ExecutorService mBgExecutor;
-    private final BroadcastSender mBroadcastSender;
-    private final BroadcastDispatcher mBroadcastDispatcher;
-    private final ScreenshotActionsController mActionsController;
-
-    @Nullable
-    private final ScreenshotSoundController mScreenshotSoundController;
-    private final ScreenshotWindow mWindow;
-    private final Display mDisplay;
-    private final ScrollCaptureExecutor mScrollCaptureExecutor;
-    private final ScreenshotNotificationSmartActionsProvider
-            mScreenshotNotificationSmartActionsProvider;
-    private final TimeoutHandler mScreenshotHandler;
-    private final UserManager mUserManager;
-    private final AssistContentRequester mAssistContentRequester;
-    private final ActionExecutor mActionExecutor;
-
-
-    private final MessageContainerController mMessageContainerController;
-    private final AnnouncementResolver mAnnouncementResolver;
-    private Bitmap mScreenBitmap;
-    private boolean mScreenshotTakenInPortrait;
-    private Animator mScreenshotAnimation;
-    private RequestCallback mCurrentRequestCallback;
-    private String mPackageName = "";
-    private final BroadcastReceiver mCopyBroadcastReceiver;
-
-    /** Tracks config changes that require re-creating UI */
-    private final InterestingConfigChanges mConfigChanges = new InterestingConfigChanges(
-            ActivityInfo.CONFIG_ORIENTATION
-                    | ActivityInfo.CONFIG_LAYOUT_DIRECTION
-                    | ActivityInfo.CONFIG_LOCALE
-                    | ActivityInfo.CONFIG_UI_MODE
-                    | ActivityInfo.CONFIG_SCREEN_LAYOUT
-                    | ActivityInfo.CONFIG_ASSETS_PATHS);
-
-
-    @AssistedInject
-    ScreenshotController(
-            Context context,
-            ScreenshotWindow.Factory screenshotWindowFactory,
-            FeatureFlags flags,
-            ScreenshotShelfViewProxy.Factory viewProxyFactory,
-            ScreenshotSmartActions screenshotSmartActions,
-            ScreenshotNotificationsController.Factory screenshotNotificationsControllerFactory,
-            UiEventLogger uiEventLogger,
-            ImageExporter imageExporter,
-            ImageCapture imageCapture,
-            @Main Executor mainExecutor,
-            ScrollCaptureExecutor scrollCaptureExecutor,
-            TimeoutHandler timeoutHandler,
-            BroadcastSender broadcastSender,
-            BroadcastDispatcher broadcastDispatcher,
-            ScreenshotNotificationSmartActionsProvider screenshotNotificationSmartActionsProvider,
-            ScreenshotActionsController.Factory screenshotActionsControllerFactory,
-            ActionExecutor.Factory actionExecutorFactory,
-            UserManager userManager,
-            AssistContentRequester assistContentRequester,
-            MessageContainerController messageContainerController,
-            Provider<ScreenshotSoundController> screenshotSoundController,
-            AnnouncementResolver announcementResolver,
-            @Assisted Display display
-    ) {
-        mScreenshotSmartActions = screenshotSmartActions;
-        mNotificationsController = screenshotNotificationsControllerFactory.create(
-                display.getDisplayId());
-        mUiEventLogger = uiEventLogger;
-        mImageExporter = imageExporter;
-        mImageCapture = imageCapture;
-        mMainExecutor = mainExecutor;
-        mScrollCaptureExecutor = scrollCaptureExecutor;
-        mScreenshotNotificationSmartActionsProvider = screenshotNotificationSmartActionsProvider;
-        mBgExecutor = Executors.newSingleThreadExecutor();
-        mBroadcastSender = broadcastSender;
-        mBroadcastDispatcher = broadcastDispatcher;
-
-        mScreenshotHandler = timeoutHandler;
-        mScreenshotHandler.setDefaultTimeoutMillis(SCREENSHOT_CORNER_DEFAULT_TIMEOUT_MILLIS);
-
-        mDisplay = display;
-        mWindow = screenshotWindowFactory.create(mDisplay);
-        mContext = mWindow.getContext();
-        mFlags = flags;
-        mUserManager = userManager;
-        mMessageContainerController = messageContainerController;
-        mAssistContentRequester = assistContentRequester;
-        mAnnouncementResolver = announcementResolver;
-
-        mViewProxy = viewProxyFactory.getProxy(mContext, mDisplay.getDisplayId());
-
-        mScreenshotHandler.setOnTimeoutRunnable(() -> {
-            if (DEBUG_UI) {
-                Log.d(TAG, "Corner timeout hit");
-            }
-            mViewProxy.requestDismissal(SCREENSHOT_INTERACTION_TIMEOUT);
-        });
-
-        mConfigChanges.applyNewConfig(context.getResources());
-        reloadAssets();
-
-        mActionExecutor = actionExecutorFactory.create(mWindow.getWindow(), mViewProxy,
-                () -> {
-                    finishDismiss();
-                    return Unit.INSTANCE;
-                });
-        mActionsController = screenshotActionsControllerFactory.getController(mActionExecutor);
-
-
-        // Sound is only reproduced from the controller of the default display.
-        if (mDisplay.getDisplayId() == Display.DEFAULT_DISPLAY) {
-            mScreenshotSoundController = screenshotSoundController.get();
-        } else {
-            mScreenshotSoundController = null;
-        }
-
-        mCopyBroadcastReceiver = new BroadcastReceiver() {
-            @Override
-            public void onReceive(Context context, Intent intent) {
-                if (ClipboardOverlayController.COPY_OVERLAY_ACTION.equals(intent.getAction())) {
-                    mViewProxy.requestDismissal(SCREENSHOT_DISMISSED_OTHER);
-                }
-            }
-        };
-        mBroadcastDispatcher.registerReceiver(mCopyBroadcastReceiver, new IntentFilter(
-                        ClipboardOverlayController.COPY_OVERLAY_ACTION), null, null,
-                Context.RECEIVER_NOT_EXPORTED, ClipboardOverlayController.SELF_PERMISSION);
-    }
-
-    @Override
-    public void handleScreenshot(ScreenshotData screenshot, Consumer<Uri> finisher,
-            RequestCallback requestCallback) {
-        Assert.isMainThread();
-
-        mCurrentRequestCallback = requestCallback;
-        if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_FULLSCREEN
-                && screenshot.getBitmap() == null) {
-            Rect bounds = getFullScreenRect();
-            screenshot.setBitmap(mImageCapture.captureDisplay(mDisplay.getDisplayId(), bounds));
-            screenshot.setScreenBounds(bounds);
-        }
-
-        if (screenshot.getBitmap() == null) {
-            Log.e(TAG, "handleScreenshot: Screenshot bitmap was null");
-            mNotificationsController.notifyScreenshotError(
-                    R.string.screenshot_failed_to_capture_text);
-            if (mCurrentRequestCallback != null) {
-                mCurrentRequestCallback.reportError();
-            }
-            return;
-        }
-
-        mScreenBitmap = screenshot.getBitmap();
-        String oldPackageName = mPackageName;
-        mPackageName = screenshot.getPackageNameString();
-
-        if (!isUserSetupComplete(Process.myUserHandle())) {
-            Log.w(TAG, "User setup not complete, displaying toast only");
-            // User setup isn't complete, so we don't want to show any UI beyond a toast, as editing
-            // and sharing shouldn't be exposed to the user.
-            saveScreenshotAndToast(screenshot, finisher);
-            return;
-        }
-
-        mBroadcastSender.sendBroadcast(new Intent(ClipboardOverlayController.SCREENSHOT_ACTION),
-                ClipboardOverlayController.SELF_PERMISSION);
-
-        mScreenshotTakenInPortrait =
-                mContext.getResources().getConfiguration().orientation == ORIENTATION_PORTRAIT;
-
-        // Optimizations
-        mScreenBitmap.setHasAlpha(false);
-        mScreenBitmap.prepareToDraw();
-
-        prepareViewForNewScreenshot(screenshot, oldPackageName);
-
-        final UUID requestId;
-        requestId = mActionsController.setCurrentScreenshot(screenshot);
-        saveScreenshotInBackground(screenshot, requestId, finisher, result -> {
-            if (result.uri != null) {
-                ScreenshotSavedResult savedScreenshot = new ScreenshotSavedResult(
-                        result.uri, screenshot.getUserOrDefault(), result.timestamp);
-                mActionsController.setCompletedScreenshot(requestId, savedScreenshot);
-            }
-        });
-
-        if (screenshot.getTaskId() >= 0) {
-            mAssistContentRequester.requestAssistContent(
-                    screenshot.getTaskId(),
-                    assistContent ->
-                            mActionsController.onAssistContent(requestId, assistContent));
-        } else {
-            mActionsController.onAssistContent(requestId, null);
-        }
-
-        // The window is focusable by default
-        mWindow.setFocusable(true);
-        mViewProxy.requestFocus();
-
-        enqueueScrollCaptureRequest(requestId, screenshot.getUserHandle());
-
-        mWindow.attachWindow();
-
-        boolean showFlash;
-        if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-            if (screenshot.getScreenBounds() != null
-                    && aspectRatiosMatch(screenshot.getBitmap(), screenshot.getInsets(),
-                    screenshot.getScreenBounds())) {
-                showFlash = false;
-            } else {
-                showFlash = true;
-                screenshot.setInsets(Insets.NONE);
-                screenshot.setScreenBounds(new Rect(0, 0, screenshot.getBitmap().getWidth(),
-                        screenshot.getBitmap().getHeight()));
-            }
-        } else {
-            showFlash = true;
-        }
-
-        mViewProxy.prepareEntranceAnimation(
-                () -> startAnimation(screenshot.getScreenBounds(), showFlash,
-                        () -> mMessageContainerController.onScreenshotTaken(screenshot)));
-
-        mViewProxy.setScreenshot(screenshot);
-
-    }
-
-    void prepareViewForNewScreenshot(@NonNull ScreenshotData screenshot, String oldPackageName) {
-        mWindow.whenWindowAttached(() -> {
-            mAnnouncementResolver.getScreenshotAnnouncement(
-                    screenshot.getUserHandle().getIdentifier(),
-                    announcement -> {
-                        mViewProxy.announceForAccessibility(announcement);
-                    });
-        });
-
-        mViewProxy.reset();
-
-        if (mViewProxy.isAttachedToWindow()) {
-            // if we didn't already dismiss for another reason
-            if (!mViewProxy.isDismissing()) {
-                mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_REENTERED, 0,
-                        oldPackageName);
-            }
-            if (DEBUG_WINDOW) {
-                Log.d(TAG, "saveScreenshot: screenshotView is already attached, resetting. "
-                        + "(dismissing=" + mViewProxy.isDismissing() + ")");
-            }
-        }
-
-        mViewProxy.setPackageName(mPackageName);
-    }
-
-    /**
-     * Requests the view to dismiss the current screenshot (may be ignored, if screenshot is already
-     * being dismissed)
-     */
-    @Override
-    public void requestDismissal(ScreenshotEvent event) {
-        mViewProxy.requestDismissal(event);
-    }
-
-    @Override
-    public boolean isPendingSharedTransition() {
-        return mActionExecutor.isPendingSharedTransition();
-    }
-
-    // Any cleanup needed when the service is being destroyed.
-    @Override
-    public void onDestroy() {
-        removeWindow();
-        releaseMediaPlayer();
-        releaseContext();
-        mBgExecutor.shutdown();
-    }
-
-    /**
-     * Release the constructed window context.
-     */
-    private void releaseContext() {
-        mBroadcastDispatcher.unregisterReceiver(mCopyBroadcastReceiver);
-        mContext.release();
-    }
-
-    private void releaseMediaPlayer() {
-        if (mScreenshotSoundController == null) return;
-        mScreenshotSoundController.releaseScreenshotSoundAsync();
-    }
-
-    /**
-     * Update resources on configuration change. Reinflate for theme/color changes.
-     */
-    private void reloadAssets() {
-        if (DEBUG_UI) {
-            Log.d(TAG, "reloadAssets()");
-        }
-
-        mMessageContainerController.setView(mViewProxy.getView());
-        mViewProxy.setCallbacks(new ScreenshotShelfViewProxy.ScreenshotViewCallback() {
-            @Override
-            public void onUserInteraction() {
-                if (DEBUG_INPUT) {
-                    Log.d(TAG, "onUserInteraction");
-                }
-                mScreenshotHandler.resetTimeout();
-            }
-
-            @Override
-            public void onDismiss() {
-                finishDismiss();
-            }
-
-            @Override
-            public void onTouchOutside() {
-                // TODO(159460485): Remove this when focus is handled properly in the system
-                mWindow.setFocusable(false);
-            }
-        });
-
-        if (DEBUG_WINDOW) {
-            Log.d(TAG, "setContentView: " + mViewProxy.getView());
-        }
-        mWindow.setContentView(mViewProxy.getView());
-    }
-
-    private void enqueueScrollCaptureRequest(UUID requestId, UserHandle owner) {
-        // Wait until this window is attached to request because it is
-        // the reference used to locate the target window (below).
-        mWindow.whenWindowAttached(() -> {
-            requestScrollCapture(requestId, owner);
-            mWindow.setActivityConfigCallback(
-                    new ViewRootImpl.ActivityConfigCallback() {
-                        @Override
-                        public void onConfigurationChanged(Configuration overrideConfig,
-                                int newDisplayId) {
-                            if (mConfigChanges.applyNewConfig(mContext.getResources())) {
-                                // Hide the scroll chip until we know it's available in this
-                                // orientation
-                                mActionsController.onScrollChipInvalidated();
-                                // Delay scroll capture eval a bit to allow the underlying activity
-                                // to set up in the new orientation.
-                                mScreenshotHandler.postDelayed(
-                                        () -> requestScrollCapture(requestId, owner), 150);
-                                mViewProxy.updateInsets(mWindow.getWindowInsets());
-                                // Screenshot animation calculations won't be valid anymore,
-                                // so just end
-                                if (mScreenshotAnimation != null
-                                        && mScreenshotAnimation.isRunning()) {
-                                    mScreenshotAnimation.end();
-                                }
-                            }
-                        }
-                    });
-        });
-    }
-
-    private void requestScrollCapture(UUID requestId, UserHandle owner) {
-        mScrollCaptureExecutor.requestScrollCapture(
-                mDisplay.getDisplayId(),
-                mWindow.getWindowToken(),
-                (response) -> {
-                    mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_IMPRESSION,
-                            0, response.getPackageName());
-                    mActionsController.onScrollChipReady(requestId,
-                            () -> onScrollButtonClicked(owner, response));
-                    return Unit.INSTANCE;
-                }
-        );
-    }
-
-    private void onScrollButtonClicked(UserHandle owner, ScrollCaptureResponse response) {
-        if (DEBUG_INPUT) {
-            Log.d(TAG, "scroll chip tapped");
-        }
-        mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_REQUESTED, 0,
-                response.getPackageName());
-        Bitmap newScreenshot = mImageCapture.captureDisplay(mDisplay.getDisplayId(),
-                getFullScreenRect());
-        if (newScreenshot == null) {
-            Log.e(TAG, "Failed to capture current screenshot for scroll transition!");
-            return;
-        }
-        // delay starting scroll capture to make sure scrim is up before the app moves
-        mViewProxy.prepareScrollingTransition(response, mScreenBitmap, newScreenshot,
-                mScreenshotTakenInPortrait, () -> executeBatchScrollCapture(response, owner));
-    }
-
-    private void executeBatchScrollCapture(ScrollCaptureResponse response, UserHandle owner) {
-        mScrollCaptureExecutor.executeBatchScrollCapture(response,
-                () -> {
-                    final Intent intent = ActionIntentCreator.INSTANCE.createLongScreenshotIntent(
-                            owner, mContext);
-                    mContext.startActivity(intent);
-                },
-                mViewProxy::restoreNonScrollingUi,
-                mViewProxy::startLongScreenshotTransition);
-    }
-
-    @Override
-    public void removeWindow() {
-        mWindow.removeWindow();
-        mViewProxy.stopInputListening();
-    }
-
-    private void playCameraSoundIfNeeded() {
-        if (mScreenshotSoundController == null) return;
-        // the controller is not-null only on the default display controller
-        mScreenshotSoundController.playScreenshotSoundAsync();
-    }
-
-    /**
-     * Save the bitmap but don't show the normal screenshot UI.. just a toast (or notification on
-     * failure).
-     */
-    private void saveScreenshotAndToast(ScreenshotData screenshot, Consumer<Uri> finisher) {
-        // Play the shutter sound to notify that we've taken a screenshot
-        playCameraSoundIfNeeded();
-
-        saveScreenshotInBackground(screenshot, UUID.randomUUID(), finisher, result -> {
-            if (result.uri != null) {
-                mScreenshotHandler.post(() -> Toast.makeText(mContext,
-                        R.string.screenshot_saved_title, Toast.LENGTH_SHORT).show());
-            }
-        });
-    }
-
-    /**
-     * Starts the animation after taking the screenshot
-     */
-    private void startAnimation(Rect screenRect, boolean showFlash, Runnable onAnimationComplete) {
-        if (mScreenshotAnimation != null && mScreenshotAnimation.isRunning()) {
-            mScreenshotAnimation.cancel();
-        }
-
-        mScreenshotAnimation =
-                mViewProxy.createScreenshotDropInAnimation(screenRect, showFlash);
-        if (onAnimationComplete != null) {
-            mScreenshotAnimation.addListener(new AnimatorListenerAdapter() {
-                @Override
-                public void onAnimationEnd(Animator animation) {
-                    super.onAnimationEnd(animation);
-                    onAnimationComplete.run();
-                }
-            });
-        }
-
-        // Play the shutter sound to notify that we've taken a screenshot
-        playCameraSoundIfNeeded();
-
-        if (DEBUG_ANIM) {
-            Log.d(TAG, "starting post-screenshot animation");
-        }
-        mScreenshotAnimation.start();
-    }
-
-    /** Reset screenshot view and then call onCompleteRunnable */
-    private void finishDismiss() {
-        Log.d(TAG, "finishDismiss");
-        mActionsController.endScreenshotSession();
-        mScrollCaptureExecutor.close();
-        if (mCurrentRequestCallback != null) {
-            mCurrentRequestCallback.onFinish();
-            mCurrentRequestCallback = null;
-        }
-        mViewProxy.reset();
-        removeWindow();
-        mScreenshotHandler.cancelTimeout();
-    }
-
-    private void saveScreenshotInBackground(ScreenshotData screenshot, UUID requestId,
-            Consumer<Uri> finisher, Consumer<ImageExporter.Result> onResult) {
-        ListenableFuture<ImageExporter.Result> future = mImageExporter.export(mBgExecutor,
-                requestId, screenshot.getBitmap(), screenshot.getUserOrDefault(),
-                mDisplay.getDisplayId());
-        future.addListener(() -> {
-            try {
-                ImageExporter.Result result = future.get();
-                Log.d(TAG, "Saved screenshot: " + result);
-                logScreenshotResultStatus(result.uri, screenshot.getUserHandle());
-                onResult.accept(result);
-                if (DEBUG_CALLBACK) {
-                    Log.d(TAG, "finished background processing, Calling (Consumer<Uri>) "
-                            + "finisher.accept(\"" + result.uri + "\"");
-                }
-                finisher.accept(result.uri);
-            } catch (Exception e) {
-                Log.d(TAG, "Failed to store screenshot", e);
-                if (DEBUG_CALLBACK) {
-                    Log.d(TAG, "Calling (Consumer<Uri>) finisher.accept(null)");
-                }
-                finisher.accept(null);
-            }
-        }, mMainExecutor);
-    }
-
-    /**
-     * Logs success/failure of the screenshot saving task, and shows an error if it failed.
-     */
-    private void logScreenshotResultStatus(Uri uri, UserHandle owner) {
-        if (uri == null) {
-            mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_NOT_SAVED, 0, mPackageName);
-            mNotificationsController.notifyScreenshotError(
-                    R.string.screenshot_failed_to_save_text);
-        } else {
-            mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_SAVED, 0, mPackageName);
-            if (mUserManager.isManagedProfile(owner.getIdentifier())) {
-                mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_SAVED_TO_WORK_PROFILE, 0,
-                        mPackageName);
-            }
-        }
-    }
-
-    private boolean isUserSetupComplete(UserHandle owner) {
-        return Settings.Secure.getInt(mContext.createContextAsUser(owner, 0)
-                .getContentResolver(), SETTINGS_SECURE_USER_SETUP_COMPLETE, 0) == 1;
-    }
-
-    private Rect getFullScreenRect() {
-        DisplayMetrics displayMetrics = new DisplayMetrics();
-        mDisplay.getRealMetrics(displayMetrics);
-        return new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels);
-    }
-
-    /** Does the aspect ratio of the bitmap with insets removed match the bounds. */
-    private static boolean aspectRatiosMatch(Bitmap bitmap, Insets bitmapInsets,
-            Rect screenBounds) {
-        int insettedWidth = bitmap.getWidth() - bitmapInsets.left - bitmapInsets.right;
-        int insettedHeight = bitmap.getHeight() - bitmapInsets.top - bitmapInsets.bottom;
-
-        if (insettedHeight == 0 || insettedWidth == 0 || bitmap.getWidth() == 0
-                || bitmap.getHeight() == 0) {
-            if (DEBUG_UI) {
-                Log.e(TAG, "Provided bitmap and insets create degenerate region: "
-                        + bitmap.getWidth() + "x" + bitmap.getHeight() + " " + bitmapInsets);
-            }
-            return false;
-        }
-
-        float insettedBitmapAspect = ((float) insettedWidth) / insettedHeight;
-        float boundsAspect = ((float) screenBounds.width()) / screenBounds.height();
-
-        boolean matchWithinTolerance = Math.abs(insettedBitmapAspect - boundsAspect) < 0.1f;
-        if (DEBUG_UI) {
-            Log.d(TAG, "aspectRatiosMatch: don't match bitmap: " + insettedBitmapAspect
-                    + ", bounds: " + boundsAspect);
-        }
-        return matchWithinTolerance;
-    }
-
-    /** Injectable factory to create screenshot controller instances for a specific display. */
-    @AssistedFactory
-    public interface Factory extends InteractiveScreenshotHandler.Factory {
-        /**
-         * Creates an instance of the controller for that specific display.
-         *
-         * @param display                 display to capture
-         */
-        ScreenshotController create(Display display);
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
new file mode 100644
index 0000000..29208f8
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.kt
@@ -0,0 +1,632 @@
+/*
+ * 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
+
+import android.animation.Animator
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.content.pm.ActivityInfo
+import android.content.res.Configuration
+import android.graphics.Bitmap
+import android.graphics.Insets
+import android.graphics.Rect
+import android.net.Uri
+import android.os.Process
+import android.os.UserHandle
+import android.os.UserManager
+import android.provider.Settings
+import android.util.DisplayMetrics
+import android.util.Log
+import android.view.Display
+import android.view.ScrollCaptureResponse
+import android.view.ViewRootImpl.ActivityConfigCallback
+import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
+import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
+import android.widget.Toast
+import android.window.WindowContext
+import androidx.core.animation.doOnEnd
+import com.android.internal.logging.UiEventLogger
+import com.android.settingslib.applications.InterestingConfigChanges
+import com.android.systemui.broadcast.BroadcastDispatcher
+import com.android.systemui.broadcast.BroadcastSender
+import com.android.systemui.clipboardoverlay.ClipboardOverlayController
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.res.R
+import com.android.systemui.screenshot.ActionIntentCreator.createLongScreenshotIntent
+import com.android.systemui.screenshot.ScreenshotShelfViewProxy.ScreenshotViewCallback
+import com.android.systemui.screenshot.scroll.ScrollCaptureController.LongScreenshot
+import com.android.systemui.screenshot.scroll.ScrollCaptureExecutor
+import com.android.systemui.util.Assert
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import java.util.UUID
+import java.util.concurrent.Executor
+import java.util.concurrent.ExecutorService
+import java.util.concurrent.Executors
+import java.util.function.Consumer
+import javax.inject.Provider
+import kotlin.math.abs
+
+/** Controls the state and flow for screenshots. */
+class ScreenshotController
+@AssistedInject
+internal constructor(
+    appContext: Context,
+    screenshotWindowFactory: ScreenshotWindow.Factory,
+    viewProxyFactory: ScreenshotShelfViewProxy.Factory,
+    screenshotNotificationsControllerFactory: ScreenshotNotificationsController.Factory,
+    screenshotActionsControllerFactory: ScreenshotActionsController.Factory,
+    actionExecutorFactory: ActionExecutor.Factory,
+    screenshotSoundControllerProvider: Provider<ScreenshotSoundController?>,
+    private val uiEventLogger: UiEventLogger,
+    private val imageExporter: ImageExporter,
+    private val imageCapture: ImageCapture,
+    private val scrollCaptureExecutor: ScrollCaptureExecutor,
+    private val screenshotHandler: TimeoutHandler,
+    private val broadcastSender: BroadcastSender,
+    private val broadcastDispatcher: BroadcastDispatcher,
+    private val userManager: UserManager,
+    private val assistContentRequester: AssistContentRequester,
+    private val messageContainerController: MessageContainerController,
+    private val announcementResolver: AnnouncementResolver,
+    @Main private val mainExecutor: Executor,
+    @Assisted private val display: Display,
+) : InteractiveScreenshotHandler {
+    private val context: WindowContext
+    private val viewProxy: ScreenshotShelfViewProxy
+    private val notificationController =
+        screenshotNotificationsControllerFactory.create(display.displayId)
+    private val bgExecutor: ExecutorService = Executors.newSingleThreadExecutor()
+    private val actionsController: ScreenshotActionsController
+    private val window: ScreenshotWindow
+    private val actionExecutor: ActionExecutor
+    private val copyBroadcastReceiver: BroadcastReceiver
+
+    private var screenshotSoundController: ScreenshotSoundController? = null
+    private var screenBitmap: Bitmap? = null
+    private var screenshotTakenInPortrait = false
+    private var screenshotAnimation: Animator? = null
+    private var currentRequestCallback: TakeScreenshotService.RequestCallback? = null
+    private var packageName = ""
+
+    /** Tracks config changes that require re-creating UI */
+    private val configChanges =
+        InterestingConfigChanges(
+            ActivityInfo.CONFIG_ORIENTATION or
+                ActivityInfo.CONFIG_LAYOUT_DIRECTION or
+                ActivityInfo.CONFIG_LOCALE or
+                ActivityInfo.CONFIG_UI_MODE or
+                ActivityInfo.CONFIG_SCREEN_LAYOUT or
+                ActivityInfo.CONFIG_ASSETS_PATHS
+        )
+
+    init {
+        screenshotHandler.defaultTimeoutMillis = SCREENSHOT_CORNER_DEFAULT_TIMEOUT_MILLIS
+
+        window = screenshotWindowFactory.create(display)
+        context = window.getContext()
+
+        viewProxy = viewProxyFactory.getProxy(context, display.displayId)
+
+        screenshotHandler.setOnTimeoutRunnable {
+            if (LogConfig.DEBUG_UI) {
+                Log.d(TAG, "Corner timeout hit")
+            }
+            viewProxy.requestDismissal(ScreenshotEvent.SCREENSHOT_INTERACTION_TIMEOUT)
+        }
+
+        configChanges.applyNewConfig(appContext.resources)
+        reloadAssets()
+
+        actionExecutor = actionExecutorFactory.create(window.window, viewProxy) { finishDismiss() }
+        actionsController = screenshotActionsControllerFactory.getController(actionExecutor)
+
+        // Sound is only reproduced from the controller of the default display.
+        screenshotSoundController =
+            if (display.displayId == Display.DEFAULT_DISPLAY) {
+                screenshotSoundControllerProvider.get()
+            } else {
+                null
+            }
+
+        copyBroadcastReceiver =
+            object : BroadcastReceiver() {
+                override fun onReceive(context: Context, intent: Intent) {
+                    if (ClipboardOverlayController.COPY_OVERLAY_ACTION == intent.action) {
+                        viewProxy.requestDismissal(ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER)
+                    }
+                }
+            }
+        broadcastDispatcher.registerReceiver(
+            copyBroadcastReceiver,
+            IntentFilter(ClipboardOverlayController.COPY_OVERLAY_ACTION),
+            null,
+            null,
+            Context.RECEIVER_NOT_EXPORTED,
+            ClipboardOverlayController.SELF_PERMISSION,
+        )
+    }
+
+    override fun handleScreenshot(
+        screenshot: ScreenshotData,
+        finisher: Consumer<Uri?>,
+        requestCallback: TakeScreenshotService.RequestCallback,
+    ) {
+        Assert.isMainThread()
+
+        currentRequestCallback = requestCallback
+        if (screenshot.type == TAKE_SCREENSHOT_FULLSCREEN && screenshot.bitmap == null) {
+            val bounds = fullScreenRect
+            screenshot.bitmap = imageCapture.captureDisplay(display.displayId, bounds)
+            screenshot.screenBounds = bounds
+        }
+
+        val currentBitmap = screenshot.bitmap
+        if (currentBitmap == null) {
+            Log.e(TAG, "handleScreenshot: Screenshot bitmap was null")
+            notificationController.notifyScreenshotError(R.string.screenshot_failed_to_capture_text)
+            currentRequestCallback?.reportError()
+            return
+        }
+
+        screenBitmap = currentBitmap
+        val oldPackageName = packageName
+        packageName = screenshot.packageNameString
+
+        if (!isUserSetupComplete(Process.myUserHandle())) {
+            Log.w(TAG, "User setup not complete, displaying toast only")
+            // User setup isn't complete, so we don't want to show any UI beyond a toast, as editing
+            // and sharing shouldn't be exposed to the user.
+            saveScreenshotAndToast(screenshot, finisher)
+            return
+        }
+
+        broadcastSender.sendBroadcast(
+            Intent(ClipboardOverlayController.SCREENSHOT_ACTION),
+            ClipboardOverlayController.SELF_PERMISSION,
+        )
+
+        screenshotTakenInPortrait =
+            context.resources.configuration.orientation == Configuration.ORIENTATION_PORTRAIT
+
+        // Optimizations
+        currentBitmap.setHasAlpha(false)
+        currentBitmap.prepareToDraw()
+
+        prepareViewForNewScreenshot(screenshot, oldPackageName)
+        val requestId = actionsController.setCurrentScreenshot(screenshot)
+        saveScreenshotInBackground(screenshot, requestId, finisher) { result ->
+            if (result.uri != null) {
+                val savedScreenshot =
+                    ScreenshotSavedResult(
+                        result.uri,
+                        screenshot.getUserOrDefault(),
+                        result.timestamp,
+                    )
+                actionsController.setCompletedScreenshot(requestId, savedScreenshot)
+            }
+        }
+
+        if (screenshot.taskId >= 0) {
+            assistContentRequester.requestAssistContent(screenshot.taskId) { assistContent ->
+                actionsController.onAssistContent(requestId, assistContent)
+            }
+        } else {
+            actionsController.onAssistContent(requestId, null)
+        }
+
+        // The window is focusable by default
+        window.setFocusable(true)
+        viewProxy.requestFocus()
+
+        enqueueScrollCaptureRequest(requestId, screenshot.userHandle!!)
+
+        window.attachWindow()
+
+        val showFlash: Boolean
+        if (screenshot.type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
+            if (aspectRatiosMatch(currentBitmap, screenshot.insets, screenshot.screenBounds)) {
+                showFlash = false
+            } else {
+                showFlash = true
+                screenshot.insets = Insets.NONE
+                screenshot.screenBounds = Rect(0, 0, currentBitmap.width, currentBitmap.height)
+            }
+        } else {
+            showFlash = true
+        }
+
+        // screenshot.screenBounds is expected to be non-null in all cases at this point
+        val bounds =
+            screenshot.screenBounds ?: Rect(0, 0, currentBitmap.width, currentBitmap.height)
+
+        viewProxy.prepareEntranceAnimation {
+            startAnimation(bounds, showFlash) {
+                messageContainerController.onScreenshotTaken(screenshot)
+            }
+        }
+
+        viewProxy.screenshot = screenshot
+    }
+
+    private fun prepareViewForNewScreenshot(screenshot: ScreenshotData, oldPackageName: String?) {
+        window.whenWindowAttached {
+            announcementResolver.getScreenshotAnnouncement(screenshot.userHandle!!.identifier) {
+                viewProxy.announceForAccessibility(it)
+            }
+        }
+
+        viewProxy.reset()
+
+        if (viewProxy.isAttachedToWindow) {
+            // if we didn't already dismiss for another reason
+            if (!viewProxy.isDismissing) {
+                uiEventLogger.log(ScreenshotEvent.SCREENSHOT_REENTERED, 0, oldPackageName)
+            }
+            if (LogConfig.DEBUG_WINDOW) {
+                Log.d(
+                    TAG,
+                    "saveScreenshot: screenshotView is already attached, resetting. " +
+                        "(dismissing=${viewProxy.isDismissing})",
+                )
+            }
+        }
+
+        viewProxy.packageName = packageName
+    }
+
+    /**
+     * Requests the view to dismiss the current screenshot (may be ignored, if screenshot is already
+     * being dismissed)
+     */
+    override fun requestDismissal(event: ScreenshotEvent) {
+        viewProxy.requestDismissal(event)
+    }
+
+    override fun isPendingSharedTransition(): Boolean {
+        return actionExecutor.isPendingSharedTransition
+    }
+
+    // Any cleanup needed when the service is being destroyed.
+    override fun onDestroy() {
+        removeWindow()
+        releaseMediaPlayer()
+        releaseContext()
+        bgExecutor.shutdown()
+    }
+
+    /** Release the constructed window context. */
+    private fun releaseContext() {
+        broadcastDispatcher.unregisterReceiver(copyBroadcastReceiver)
+        context.release()
+    }
+
+    private fun releaseMediaPlayer() {
+        screenshotSoundController?.releaseScreenshotSoundAsync()
+    }
+
+    /** Update resources on configuration change. Reinflate for theme/color changes. */
+    private fun reloadAssets() {
+        if (LogConfig.DEBUG_UI) {
+            Log.d(TAG, "reloadAssets()")
+        }
+
+        messageContainerController.setView(viewProxy.view)
+        viewProxy.callbacks =
+            object : ScreenshotViewCallback {
+                override fun onUserInteraction() {
+                    if (LogConfig.DEBUG_INPUT) {
+                        Log.d(TAG, "onUserInteraction")
+                    }
+                    screenshotHandler.resetTimeout()
+                }
+
+                override fun onDismiss() {
+                    finishDismiss()
+                }
+
+                override fun onTouchOutside() {
+                    // TODO(159460485): Remove this when focus is handled properly in the system
+                    window.setFocusable(false)
+                }
+            }
+
+        if (LogConfig.DEBUG_WINDOW) {
+            Log.d(TAG, "setContentView: " + viewProxy.view)
+        }
+        window.setContentView(viewProxy.view)
+    }
+
+    private fun enqueueScrollCaptureRequest(requestId: UUID, owner: UserHandle) {
+        // Wait until this window is attached to request because it is
+        // the reference used to locate the target window (below).
+        window.whenWindowAttached {
+            requestScrollCapture(requestId, owner)
+            window.setActivityConfigCallback(
+                object : ActivityConfigCallback {
+                    override fun onConfigurationChanged(
+                        overrideConfig: Configuration,
+                        newDisplayId: Int,
+                    ) {
+                        if (configChanges.applyNewConfig(context.resources)) {
+                            // Hide the scroll chip until we know it's available in this
+                            // orientation
+                            actionsController.onScrollChipInvalidated()
+                            // Delay scroll capture eval a bit to allow the underlying activity
+                            // to set up in the new orientation.
+                            screenshotHandler.postDelayed(
+                                { requestScrollCapture(requestId, owner) },
+                                150,
+                            )
+                            viewProxy.updateInsets(window.getWindowInsets())
+                            // Screenshot animation calculations won't be valid anymore, so just end
+                            screenshotAnimation?.let { currentAnimation ->
+                                if (currentAnimation.isRunning) {
+                                    currentAnimation.end()
+                                }
+                            }
+                        }
+                    }
+                }
+            )
+        }
+    }
+
+    private fun requestScrollCapture(requestId: UUID, owner: UserHandle) {
+        scrollCaptureExecutor.requestScrollCapture(display.displayId, window.getWindowToken()) {
+            response: ScrollCaptureResponse ->
+            uiEventLogger.log(
+                ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_IMPRESSION,
+                0,
+                response.packageName,
+            )
+            actionsController.onScrollChipReady(requestId) {
+                onScrollButtonClicked(owner, response)
+            }
+        }
+    }
+
+    private fun onScrollButtonClicked(owner: UserHandle, response: ScrollCaptureResponse) {
+        if (LogConfig.DEBUG_INPUT) {
+            Log.d(TAG, "scroll chip tapped")
+        }
+        uiEventLogger.log(
+            ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_REQUESTED,
+            0,
+            response.packageName,
+        )
+        val newScreenshot = imageCapture.captureDisplay(display.displayId, null)
+        if (newScreenshot == null) {
+            Log.e(TAG, "Failed to capture current screenshot for scroll transition!")
+            return
+        }
+        // delay starting scroll capture to make sure scrim is up before the app moves
+        viewProxy.prepareScrollingTransition(response, newScreenshot, screenshotTakenInPortrait) {
+            executeBatchScrollCapture(response, owner)
+        }
+    }
+
+    private fun executeBatchScrollCapture(response: ScrollCaptureResponse, owner: UserHandle) {
+        scrollCaptureExecutor.executeBatchScrollCapture(
+            response,
+            {
+                val intent = createLongScreenshotIntent(owner, context)
+                context.startActivity(intent)
+            },
+            { viewProxy.restoreNonScrollingUi() },
+            { transitionDestination: Rect, onTransitionEnd: Runnable, longScreenshot: LongScreenshot
+                ->
+                viewProxy.startLongScreenshotTransition(
+                    transitionDestination,
+                    onTransitionEnd,
+                    longScreenshot,
+                )
+            },
+        )
+    }
+
+    override fun removeWindow() {
+        window.removeWindow()
+        viewProxy.stopInputListening()
+    }
+
+    private fun playCameraSoundIfNeeded() {
+        // the controller is not-null only on the default display controller
+        screenshotSoundController?.playScreenshotSoundAsync()
+    }
+
+    /**
+     * Save the bitmap but don't show the normal screenshot UI.. just a toast (or notification on
+     * failure).
+     */
+    private fun saveScreenshotAndToast(screenshot: ScreenshotData, finisher: Consumer<Uri?>) {
+        // Play the shutter sound to notify that we've taken a screenshot
+        playCameraSoundIfNeeded()
+
+        saveScreenshotInBackground(screenshot, UUID.randomUUID(), finisher) {
+            result: ImageExporter.Result ->
+            if (result.uri != null) {
+                screenshotHandler.post {
+                    Toast.makeText(context, R.string.screenshot_saved_title, Toast.LENGTH_SHORT)
+                        .show()
+                }
+            }
+        }
+    }
+
+    /** Starts the animation after taking the screenshot */
+    private fun startAnimation(
+        screenRect: Rect,
+        showFlash: Boolean,
+        onAnimationComplete: Runnable?,
+    ) {
+        screenshotAnimation?.let { currentAnimation ->
+            if (currentAnimation.isRunning) {
+                currentAnimation.cancel()
+            }
+        }
+
+        screenshotAnimation =
+            viewProxy.createScreenshotDropInAnimation(screenRect, showFlash).apply {
+                doOnEnd { onAnimationComplete?.run() }
+                // Play the shutter sound to notify that we've taken a screenshot
+                playCameraSoundIfNeeded()
+                if (LogConfig.DEBUG_ANIM) {
+                    Log.d(TAG, "starting post-screenshot animation")
+                }
+                start()
+            }
+    }
+
+    /** Reset screenshot view and then call onCompleteRunnable */
+    private fun finishDismiss() {
+        Log.d(TAG, "finishDismiss")
+        actionsController.endScreenshotSession()
+        scrollCaptureExecutor.close()
+        currentRequestCallback?.onFinish()
+        currentRequestCallback = null
+        viewProxy.reset()
+        removeWindow()
+        screenshotHandler.cancelTimeout()
+    }
+
+    private fun saveScreenshotInBackground(
+        screenshot: ScreenshotData,
+        requestId: UUID,
+        finisher: Consumer<Uri?>,
+        onResult: Consumer<ImageExporter.Result>,
+    ) {
+        val future =
+            imageExporter.export(
+                bgExecutor,
+                requestId,
+                screenshot.bitmap,
+                screenshot.getUserOrDefault(),
+                display.displayId,
+            )
+        future.addListener(
+            {
+                try {
+                    val result = future.get()
+                    Log.d(TAG, "Saved screenshot: $result")
+                    logScreenshotResultStatus(result.uri, screenshot.userHandle!!)
+                    onResult.accept(result)
+                    if (LogConfig.DEBUG_CALLBACK) {
+                        Log.d(TAG, "finished bg processing, calling back with uri: ${result.uri}")
+                    }
+                    finisher.accept(result.uri)
+                } catch (e: Exception) {
+                    Log.d(TAG, "Failed to store screenshot", e)
+                    if (LogConfig.DEBUG_CALLBACK) {
+                        Log.d(TAG, "calling back with uri: null")
+                    }
+                    finisher.accept(null)
+                }
+            },
+            mainExecutor,
+        )
+    }
+
+    /** Logs success/failure of the screenshot saving task, and shows an error if it failed. */
+    private fun logScreenshotResultStatus(uri: Uri?, owner: UserHandle) {
+        if (uri == null) {
+            uiEventLogger.log(ScreenshotEvent.SCREENSHOT_NOT_SAVED, 0, packageName)
+            notificationController.notifyScreenshotError(R.string.screenshot_failed_to_save_text)
+        } else {
+            uiEventLogger.log(ScreenshotEvent.SCREENSHOT_SAVED, 0, packageName)
+            if (userManager.isManagedProfile(owner.identifier)) {
+                uiEventLogger.log(ScreenshotEvent.SCREENSHOT_SAVED_TO_WORK_PROFILE, 0, packageName)
+            }
+        }
+    }
+
+    private fun isUserSetupComplete(owner: UserHandle): Boolean {
+        return Settings.Secure.getInt(
+            context.createContextAsUser(owner, 0).contentResolver,
+            SETTINGS_SECURE_USER_SETUP_COMPLETE,
+            0,
+        ) == 1
+    }
+
+    private val fullScreenRect: Rect
+        get() {
+            val displayMetrics = DisplayMetrics()
+            display.getRealMetrics(displayMetrics)
+            return Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels)
+        }
+
+    /** Injectable factory to create screenshot controller instances for a specific display. */
+    @AssistedFactory
+    interface Factory : InteractiveScreenshotHandler.Factory {
+        /**
+         * Creates an instance of the controller for that specific display.
+         *
+         * @param display display to capture
+         */
+        override fun create(display: Display): ScreenshotController
+    }
+
+    companion object {
+        private val TAG: String = LogConfig.logTag(ScreenshotController::class.java)
+
+        // From WizardManagerHelper.java
+        private const val SETTINGS_SECURE_USER_SETUP_COMPLETE = "user_setup_complete"
+
+        const val SCREENSHOT_CORNER_DEFAULT_TIMEOUT_MILLIS: Int = 6000
+
+        /** Does the aspect ratio of the bitmap with insets removed match the bounds. */
+        private fun aspectRatiosMatch(
+            bitmap: Bitmap,
+            bitmapInsets: Insets,
+            screenBounds: Rect?,
+        ): Boolean {
+            if (screenBounds == null) {
+                return false
+            }
+            val insettedWidth = bitmap.width - bitmapInsets.left - bitmapInsets.right
+            val insettedHeight = bitmap.height - bitmapInsets.top - bitmapInsets.bottom
+
+            if (
+                insettedHeight == 0 || insettedWidth == 0 || bitmap.width == 0 || bitmap.height == 0
+            ) {
+                if (LogConfig.DEBUG_UI) {
+                    Log.e(
+                        TAG,
+                        "Provided bitmap and insets create degenerate region: " +
+                            "${bitmap.width} x ${bitmap.height} $bitmapInsets",
+                    )
+                }
+                return false
+            }
+
+            val insettedBitmapAspect = insettedWidth.toFloat() / insettedHeight
+            val boundsAspect = screenBounds.width().toFloat() / screenBounds.height()
+
+            val matchWithinTolerance = abs((insettedBitmapAspect - boundsAspect).toDouble()) < 0.1f
+            if (LogConfig.DEBUG_UI) {
+                Log.d(
+                    TAG,
+                    "aspectRatiosMatch: don't match bitmap: " +
+                        "$insettedBitmapAspect, bounds: $boundsAspect",
+                )
+            }
+            return matchWithinTolerance
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
index 50215af..c1ea3ad 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
@@ -67,7 +67,7 @@
     shelfViewBinder: ScreenshotShelfViewBinder,
     private val thumbnailObserver: ThumbnailObserver,
     @Assisted private val context: Context,
-    @Assisted private val displayId: Int
+    @Assisted private val displayId: Int,
 ) {
 
     interface ScreenshotViewCallback {
@@ -117,7 +117,7 @@
             animationController,
             LayoutInflater.from(context),
             onDismissalRequested = { event, velocity -> requestDismissal(event, velocity) },
-            onUserInteraction = { callbacks?.onUserInteraction() }
+            onUserInteraction = { callbacks?.onUserInteraction() },
         )
         view.updateInsets(windowManager.currentWindowMetrics.windowInsets)
         addPredictiveBackListener { requestDismissal(SCREENSHOT_DISMISSED_OTHER) }
@@ -130,7 +130,7 @@
         screenshotPreview = view.screenshotPreview
         thumbnailObserver.setViews(
             view.blurredScreenshotPreview,
-            view.requireViewById(R.id.screenshot_preview_border)
+            view.requireViewById(R.id.screenshot_preview_border),
         )
         view.addOnAttachStateChangeListener(
             object : View.OnAttachStateChangeListener {
@@ -204,7 +204,6 @@
 
     fun prepareScrollingTransition(
         response: ScrollCaptureResponse,
-        screenBitmap: Bitmap, // unused
         newScreenshot: Bitmap,
         screenshotTakenInPortrait: Boolean,
         onTransitionPrepared: Runnable,
@@ -224,7 +223,7 @@
                 0,
                 0,
                 context.resources.displayMetrics.widthPixels,
-                context.resources.displayMetrics.heightPixels
+                context.resources.displayMetrics.heightPixels,
             )
         )
         return r
@@ -239,7 +238,7 @@
             animationController.runLongScreenshotTransition(
                 transitionDestination,
                 longScreenshot,
-                onTransitionEnd
+                onTransitionEnd,
             )
         transitionAnimation.doOnEnd { callbacks?.onDismiss() }
         transitionAnimation.start()
@@ -295,7 +294,7 @@
                         .findOnBackInvokedDispatcher()
                         ?.registerOnBackInvokedCallback(
                             OnBackInvokedDispatcher.PRIORITY_DEFAULT,
-                            onBackInvokedCallback
+                            onBackInvokedCallback,
                         )
                 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowStateController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowStateController.kt
index 8f424b2..fa9c6b2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowStateController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/window/StatusBarWindowStateController.kt
@@ -17,9 +17,9 @@
 package com.android.systemui.statusbar.window
 
 import android.app.StatusBarManager
-import android.app.StatusBarManager.WindowVisibleState
 import android.app.StatusBarManager.WINDOW_STATE_SHOWING
 import android.app.StatusBarManager.WINDOW_STATUS_BAR
+import android.app.StatusBarManager.WindowVisibleState
 import android.app.StatusBarManager.windowStateToString
 import android.util.Log
 import com.android.systemui.dagger.SysUISingleton
@@ -31,23 +31,27 @@
 /**
  * A centralized class maintaining the state of the status bar window.
  *
+ * @deprecated use
+ *   [com.android.systemui.statusbar.window.data.repository.StatusBarWindowStateRepository] instead.
+ *
  * Classes that want to get updates about the status bar window state should subscribe to this class
  * via [addListener] and should NOT add their own callback on [CommandQueue].
  */
 @SysUISingleton
-class StatusBarWindowStateController @Inject constructor(
-    @DisplayId private val thisDisplayId: Int,
-    commandQueue: CommandQueue
-) {
-    private val commandQueueCallback = object : CommandQueue.Callbacks {
-        override fun setWindowState(
-            displayId: Int,
-            @StatusBarManager.WindowType window: Int,
-            @WindowVisibleState state: Int
-        ) {
-            [email protected](displayId, window, state)
+@Deprecated("Use StatusBarWindowRepository instead")
+class StatusBarWindowStateController
+@Inject
+constructor(@DisplayId private val thisDisplayId: Int, commandQueue: CommandQueue) {
+    private val commandQueueCallback =
+        object : CommandQueue.Callbacks {
+            override fun setWindowState(
+                displayId: Int,
+                @StatusBarManager.WindowType window: Int,
+                @WindowVisibleState state: Int,
+            ) {
+                [email protected](displayId, window, state)
+            }
         }
-    }
     private val listeners: MutableSet<StatusBarWindowStateListener> = HashSet()
 
     @WindowVisibleState private var windowState: Int = WINDOW_STATE_SHOWING
@@ -71,7 +75,7 @@
     private fun setWindowState(
         displayId: Int,
         @StatusBarManager.WindowType window: Int,
-        @WindowVisibleState state: Int
+        @WindowVisibleState state: Int,
     ) {
         if (displayId != thisDisplayId) {
             return
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepository.kt
new file mode 100644
index 0000000..678576d
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepository.kt
@@ -0,0 +1,86 @@
+/*
+ * 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.statusbar.window.data.repository
+
+import android.app.StatusBarManager
+import android.app.StatusBarManager.WINDOW_STATE_HIDDEN
+import android.app.StatusBarManager.WINDOW_STATE_HIDING
+import android.app.StatusBarManager.WINDOW_STATE_SHOWING
+import android.app.StatusBarManager.WINDOW_STATUS_BAR
+import android.app.StatusBarManager.WindowVisibleState
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.DisplayId
+import com.android.systemui.statusbar.CommandQueue
+import com.android.systemui.statusbar.window.data.model.StatusBarWindowState
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.stateIn
+
+/**
+ * A centralized class maintaining the state of the status bar window.
+ *
+ * Classes that want to get updates about the status bar window state should subscribe to
+ * [windowState] and should NOT add their own callback on [CommandQueue].
+ */
+@SysUISingleton
+class StatusBarWindowStateRepository
+@Inject
+constructor(
+    private val commandQueue: CommandQueue,
+    @DisplayId private val thisDisplayId: Int,
+    @Application private val scope: CoroutineScope,
+) {
+    val windowState: StateFlow<StatusBarWindowState> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : CommandQueue.Callbacks {
+                        override fun setWindowState(
+                            displayId: Int,
+                            @StatusBarManager.WindowType window: Int,
+                            @WindowVisibleState state: Int,
+                        ) {
+                            // TODO(b/364360986): Log the window state changes.
+                            if (displayId != thisDisplayId) {
+                                return
+                            }
+                            if (window != WINDOW_STATUS_BAR) {
+                                return
+                            }
+                            trySend(state.toWindowState())
+                        }
+                    }
+                commandQueue.addCallback(callback)
+                awaitClose { commandQueue.removeCallback(callback) }
+            }
+            // Use Eagerly because we always need to know about the status bar window state
+            .stateIn(scope, SharingStarted.Eagerly, StatusBarWindowState.Hidden)
+
+    @WindowVisibleState
+    private fun Int.toWindowState(): StatusBarWindowState {
+        return when (this) {
+            WINDOW_STATE_SHOWING -> StatusBarWindowState.Showing
+            WINDOW_STATE_HIDING -> StatusBarWindowState.Hiding
+            WINDOW_STATE_HIDDEN -> StatusBarWindowState.Hidden
+            else -> StatusBarWindowState.Hidden
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/window/shared/model/StatusBarWindowState.kt b/packages/SystemUI/src/com/android/systemui/statusbar/window/shared/model/StatusBarWindowState.kt
new file mode 100644
index 0000000..a99046e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/window/shared/model/StatusBarWindowState.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.window.data.model
+
+/**
+ * Represents the state of the status bar *window* as a whole (as opposed to individual views within
+ * the status bar).
+ */
+enum class StatusBarWindowState {
+    Showing,
+    Hiding,
+    Hidden,
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt
index b86d571..ab846f1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt
@@ -21,6 +21,7 @@
 import android.hardware.input.KeyGestureEvent
 import android.os.UserHandle
 import android.os.UserManager
+import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
 import android.view.KeyEvent
@@ -32,6 +33,8 @@
 import androidx.test.filters.SmallTest
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.notetask.NoteTaskEntryPoint.KEYBOARD_SHORTCUT
+import com.android.systemui.notetask.NoteTaskEntryPoint.TAIL_BUTTON
 import com.android.systemui.settings.FakeUserTracker
 import com.android.systemui.statusbar.CommandQueue
 import com.android.systemui.util.concurrency.FakeExecutor
@@ -62,8 +65,7 @@
 @RunWith(AndroidJUnit4::class)
 internal class NoteTaskInitializerTest : SysuiTestCase() {
 
-    @get:Rule
-    val setFlagsRule = SetFlagsRule()
+    @get:Rule val setFlagsRule = SetFlagsRule()
 
     @Mock lateinit var commandQueue: CommandQueue
     @Mock lateinit var inputManager: InputManager
@@ -83,10 +85,7 @@
         whenever(keyguardMonitor.isUserUnlocked(userTracker.userId)).thenReturn(true)
     }
 
-    private fun createUnderTest(
-        isEnabled: Boolean,
-        bubbles: Bubbles?,
-    ): NoteTaskInitializer =
+    private fun createUnderTest(isEnabled: Boolean, bubbles: Bubbles?): NoteTaskInitializer =
         NoteTaskInitializer(
             controller = controller,
             commandQueue = commandQueue,
@@ -104,7 +103,7 @@
         code: Int,
         downTime: Long = 0L,
         eventTime: Long = 0L,
-        metaState: Int = 0
+        metaState: Int = 0,
     ): KeyEvent = KeyEvent(downTime, eventTime, action, code, 0 /*repeat*/, metaState)
 
     @Test
@@ -113,7 +112,6 @@
 
         createUnderTest(isEnabled = true, bubbles = bubbles).initialize()
 
-        verify(commandQueue).addCallback(any())
         verify(roleManager).addOnRoleHoldersChangedListenerAsUser(any(), any(), any())
         verify(controller).updateNoteTaskForCurrentUserAndManagedProfiles()
         verify(keyguardMonitor).registerCallback(any())
@@ -125,7 +123,6 @@
 
         createUnderTest(isEnabled = true, bubbles = bubbles).initialize()
 
-        verify(commandQueue).addCallback(any())
         verify(roleManager).addOnRoleHoldersChangedListenerAsUser(any(), any(), any())
         verify(controller, never()).setNoteTaskShortcutEnabled(any(), any())
         verify(keyguardMonitor).registerCallback(any())
@@ -165,12 +162,13 @@
     }
 
     @Test
+    @DisableFlags(com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
     fun initialize_handleSystemKey() {
         val expectedKeyEvent =
             createKeyEvent(
                 ACTION_DOWN,
                 KEYCODE_N,
-                metaState = KeyEvent.META_META_ON or KeyEvent.META_CTRL_ON
+                metaState = KeyEvent.META_META_ON or KeyEvent.META_CTRL_ON,
             )
         val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
         underTest.initialize()
@@ -183,22 +181,66 @@
 
     @Test
     @EnableFlags(com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
-    fun initialize_handleKeyGestureEvent() {
-        val gestureEvent = KeyGestureEvent.Builder()
-            .setKeycodes(intArrayOf(KeyEvent.KEYCODE_N))
-            .setModifierState(KeyEvent.META_META_ON or KeyEvent.META_CTRL_ON)
-            .setKeyGestureType(KeyGestureEvent.KEY_GESTURE_TYPE_OPEN_NOTES)
-            .setAction(KeyGestureEvent.ACTION_GESTURE_COMPLETE)
-            .build()
+    fun handlesShortcut_metaCtrlN() {
+        val gestureEvent =
+            KeyGestureEvent.Builder()
+                .setKeycodes(intArrayOf(KeyEvent.KEYCODE_N))
+                .setModifierState(KeyEvent.META_META_ON or KeyEvent.META_CTRL_ON)
+                .setKeyGestureType(KeyGestureEvent.KEY_GESTURE_TYPE_OPEN_NOTES)
+                .setAction(KeyGestureEvent.ACTION_GESTURE_COMPLETE)
+                .build()
         val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
         underTest.initialize()
-        val callback =
-            withArgCaptor { verify(inputManager).registerKeyGestureEventHandler(capture()) }
+        val callback = withArgCaptor {
+            verify(inputManager).registerKeyGestureEventHandler(capture())
+        }
 
         assertThat(callback.handleKeyGestureEvent(gestureEvent, null)).isTrue()
 
         executor.runAllReady()
-        verify(controller).showNoteTask(any())
+        verify(controller).showNoteTask(eq(KEYBOARD_SHORTCUT))
+    }
+
+    @Test
+    @EnableFlags(com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
+    fun handlesShortcut_stylusTailButton() {
+        val gestureEvent =
+            KeyGestureEvent.Builder()
+                .setKeycodes(intArrayOf(KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL))
+                .setKeyGestureType(KeyGestureEvent.KEY_GESTURE_TYPE_OPEN_NOTES)
+                .setAction(KeyGestureEvent.ACTION_GESTURE_COMPLETE)
+                .build()
+        val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
+        underTest.initialize()
+        val callback = withArgCaptor {
+            verify(inputManager).registerKeyGestureEventHandler(capture())
+        }
+
+        assertThat(callback.handleKeyGestureEvent(gestureEvent, null)).isTrue()
+
+        executor.runAllReady()
+        verify(controller).showNoteTask(eq(TAIL_BUTTON))
+    }
+
+    @Test
+    @EnableFlags(com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
+    fun ignoresUnrelatedShortcuts() {
+        val gestureEvent =
+            KeyGestureEvent.Builder()
+                .setKeycodes(intArrayOf(KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL))
+                .setKeyGestureType(KeyGestureEvent.KEY_GESTURE_TYPE_HOME)
+                .setAction(KeyGestureEvent.ACTION_GESTURE_COMPLETE)
+                .build()
+        val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
+        underTest.initialize()
+        val callback = withArgCaptor {
+            verify(inputManager).registerKeyGestureEventHandler(capture())
+        }
+
+        assertThat(callback.handleKeyGestureEvent(gestureEvent, null)).isFalse()
+
+        executor.runAllReady()
+        verify(controller, never()).showNoteTask(any())
     }
 
     @Test
@@ -249,6 +291,7 @@
     }
 
     @Test
+    @DisableFlags(com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
     fun tailButtonGestureDetection_singlePress_shouldShowNoteTaskOnUp() {
         val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
         underTest.initialize()
@@ -267,6 +310,7 @@
     }
 
     @Test
+    @DisableFlags(com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
     fun tailButtonGestureDetection_doublePress_shouldNotShowNoteTaskTwice() {
         val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
         underTest.initialize()
@@ -289,6 +333,7 @@
     }
 
     @Test
+    @DisableFlags(com.android.hardware.input.Flags.FLAG_USE_KEY_GESTURE_EVENT_HANDLER)
     fun tailButtonGestureDetection_longPress_shouldNotShowNoteTask() {
         val underTest = createUnderTest(isEnabled = true, bubbles = bubbles)
         underTest.initialize()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryTest.kt
new file mode 100644
index 0000000..38e04bb
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/window/data/repository/StatusBarWindowStateRepositoryTest.kt
@@ -0,0 +1,112 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.window.data.repository
+
+import android.app.StatusBarManager.WINDOW_NAVIGATION_BAR
+import android.app.StatusBarManager.WINDOW_STATE_HIDDEN
+import android.app.StatusBarManager.WINDOW_STATE_HIDING
+import android.app.StatusBarManager.WINDOW_STATE_SHOWING
+import android.app.StatusBarManager.WINDOW_STATUS_BAR
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.statusbar.CommandQueue
+import com.android.systemui.statusbar.commandQueue
+import com.android.systemui.statusbar.window.data.model.StatusBarWindowState
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.mockito.Mockito.verify
+import org.mockito.kotlin.argumentCaptor
+
+@SmallTest
+@OptIn(ExperimentalCoroutinesApi::class)
+class StatusBarWindowStateRepositoryTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val commandQueue = kosmos.commandQueue
+    private val underTest =
+        StatusBarWindowStateRepository(commandQueue, DISPLAY_ID, testScope.backgroundScope)
+
+    private val callback: CommandQueue.Callbacks
+        get() {
+            testScope.runCurrent()
+            val callbackCaptor = argumentCaptor<CommandQueue.Callbacks>()
+            verify(commandQueue).addCallback(callbackCaptor.capture())
+            return callbackCaptor.firstValue
+        }
+
+    @Test
+    fun windowState_notSameDisplayId_notUpdated() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.windowState)
+            assertThat(latest).isEqualTo(StatusBarWindowState.Hidden)
+
+            callback.setWindowState(DISPLAY_ID + 1, WINDOW_STATUS_BAR, WINDOW_STATE_SHOWING)
+
+            assertThat(latest).isEqualTo(StatusBarWindowState.Hidden)
+        }
+
+    @Test
+    fun windowState_notStatusBarWindow_notUpdated() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.windowState)
+            assertThat(latest).isEqualTo(StatusBarWindowState.Hidden)
+
+            callback.setWindowState(DISPLAY_ID, WINDOW_NAVIGATION_BAR, WINDOW_STATE_SHOWING)
+
+            assertThat(latest).isEqualTo(StatusBarWindowState.Hidden)
+        }
+
+    @Test
+    fun windowState_showing_updated() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.windowState)
+
+            callback.setWindowState(DISPLAY_ID, WINDOW_STATUS_BAR, WINDOW_STATE_SHOWING)
+
+            assertThat(latest).isEqualTo(StatusBarWindowState.Showing)
+        }
+
+    @Test
+    fun windowState_hiding_updated() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.windowState)
+
+            callback.setWindowState(DISPLAY_ID, WINDOW_STATUS_BAR, WINDOW_STATE_HIDING)
+
+            assertThat(latest).isEqualTo(StatusBarWindowState.Hiding)
+        }
+
+    @Test
+    fun windowState_hidden_updated() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.windowState)
+            callback.setWindowState(DISPLAY_ID, WINDOW_STATUS_BAR, WINDOW_STATE_SHOWING)
+            assertThat(latest).isEqualTo(StatusBarWindowState.Showing)
+
+            callback.setWindowState(DISPLAY_ID, WINDOW_STATUS_BAR, WINDOW_STATE_HIDDEN)
+
+            assertThat(latest).isEqualTo(StatusBarWindowState.Hidden)
+        }
+}
+
+private const val DISPLAY_ID = 10
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
index 2417262..90bb93d 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuntimeEnvironmentController.java
@@ -16,12 +16,11 @@
 
 package android.platform.test.ravenwood;
 
+import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_INST_RESOURCE_APK;
 import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_RESOURCE_APK;
 import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_VERBOSE_LOGGING;
 import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_VERSION_JAVA_SYSPROP;
 
-import static org.junit.Assert.fail;
-
 import android.app.ActivityManager;
 import android.app.Instrumentation;
 import android.app.ResourcesManager;
@@ -211,23 +210,21 @@
             var file = new File(RAVENWOOD_RESOURCE_APK);
             return config.mState.loadResources(file.exists() ? file : null);
         };
-        // Set up test context's resources.
+
+        // Set up test context's (== instrumentation context's) resources.
         // If the target package name == test package name, then we use the main resources.
-        // Otherwise, we don't simulate loading resources from the test APK yet.
-        // (we need to add `test_resource_apk` to `android_ravenwood_test`)
-        final Supplier<Resources> testResourcesLoader;
+        final Supplier<Resources> instResourcesLoader;
         if (isSelfInstrumenting) {
-            testResourcesLoader = targetResourcesLoader;
+            instResourcesLoader = targetResourcesLoader;
         } else {
-            testResourcesLoader = () -> {
-                fail("Cannot load resources from the test context (yet)."
-                        + " Use target context's resources instead.");
-                return null; // unreachable.
+            instResourcesLoader = () -> {
+                var file = new File(RAVENWOOD_INST_RESOURCE_APK);
+                return config.mState.loadResources(file.exists() ? file : null);
             };
         }
 
-        var testContext = new RavenwoodContext(
-                config.mTestPackageName, main, testResourcesLoader);
+        var instContext = new RavenwoodContext(
+                config.mTestPackageName, main, instResourcesLoader);
         var targetContext = new RavenwoodContext(
                 config.mTargetPackageName, main, targetResourcesLoader);
 
@@ -236,18 +233,18 @@
                 config.mTargetPackageName, main, targetResourcesLoader);
         appContext.setApplicationContext(appContext);
         if (isSelfInstrumenting) {
-            testContext.setApplicationContext(appContext);
+            instContext.setApplicationContext(appContext);
             targetContext.setApplicationContext(appContext);
         } else {
             // When instrumenting into another APK, the test context doesn't have an app context.
             targetContext.setApplicationContext(appContext);
         }
-        config.mTestContext = testContext;
+        config.mInstContext = instContext;
         config.mTargetContext = targetContext;
 
         // Prepare other fields.
         config.mInstrumentation = new Instrumentation();
-        config.mInstrumentation.basicInit(config.mTestContext, config.mTargetContext);
+        config.mInstrumentation.basicInit(config.mInstContext, config.mTargetContext);
         InstrumentationRegistry.registerInstance(config.mInstrumentation, Bundle.EMPTY);
 
         RavenwoodSystemServer.init(config);
@@ -284,13 +281,13 @@
 
         InstrumentationRegistry.registerInstance(null, Bundle.EMPTY);
         config.mInstrumentation = null;
-        if (config.mTestContext != null) {
-            ((RavenwoodContext) config.mTestContext).cleanUp();
+        if (config.mInstContext != null) {
+            ((RavenwoodContext) config.mInstContext).cleanUp();
         }
         if (config.mTargetContext != null) {
             ((RavenwoodContext) config.mTargetContext).cleanUp();
         }
-        config.mTestContext = null;
+        config.mInstContext = null;
         config.mTargetContext = null;
 
         if (config.mProvideMainThread) {
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java
index d4090e2..3946dd84 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodSystemServer.java
@@ -67,7 +67,7 @@
 
         sStartedServices = new ArraySet<>();
         sTimings = new TimingsTraceAndSlog();
-        sServiceManager = new SystemServiceManager(config.mTestContext);
+        sServiceManager = new SystemServiceManager(config.mInstContext);
         sServiceManager.setStartInfo(false,
                 SystemClock.elapsedRealtime(),
                 SystemClock.uptimeMillis());
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java
index ea33aa6..446f819 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodConfig.java
@@ -74,7 +74,7 @@
 
     final List<Class<?>> mServicesRequired = new ArrayList<>();
 
-    volatile Context mTestContext;
+    volatile Context mInstContext;
     volatile Context mTargetContext;
     volatile Instrumentation mInstrumentation;
 
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
index 984106b..4196d8e 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
@@ -216,7 +216,7 @@
      */
     @Deprecated
     public Context getContext() {
-        return Objects.requireNonNull(mConfiguration.mTestContext,
+        return Objects.requireNonNull(mConfiguration.mInstContext,
                 "Context is only available during @Test execution");
     }
 
diff --git a/ravenwood/runtime-common-src/com/android/ravenwood/common/RavenwoodCommonUtils.java b/ravenwood/runtime-common-src/com/android/ravenwood/common/RavenwoodCommonUtils.java
index 96746c6..989bb6b 100644
--- a/ravenwood/runtime-common-src/com/android/ravenwood/common/RavenwoodCommonUtils.java
+++ b/ravenwood/runtime-common-src/com/android/ravenwood/common/RavenwoodCommonUtils.java
@@ -63,6 +63,8 @@
     public static final String RAVENWOOD_SYSPROP = "ro.is_on_ravenwood";
 
     public static final String RAVENWOOD_RESOURCE_APK = "ravenwood-res-apks/ravenwood-res.apk";
+    public static final String RAVENWOOD_INST_RESOURCE_APK =
+            "ravenwood-res-apks/ravenwood-inst-res.apk";
 
     public static final String RAVENWOOD_EMPTY_RESOURCES_APK =
             RAVENWOOD_RUNTIME_PATH + "ravenwood-data/ravenwood-empty-res.apk";
diff --git a/ravenwood/tests/bivalentinst/Android.bp b/ravenwood/tests/bivalentinst/Android.bp
index 38d1b299..41e45e5 100644
--- a/ravenwood/tests/bivalentinst/Android.bp
+++ b/ravenwood/tests/bivalentinst/Android.bp
@@ -27,8 +27,7 @@
         "junit",
         "truth",
     ],
-    // TODO(b/366246777) uncomment it and the test.
-    // resource_apk: "RavenwoodBivalentInstTest_self_inst_device",
+    resource_apk: "RavenwoodBivalentInstTest_self_inst_device",
     auto_gen_config: true,
 }
 
@@ -53,8 +52,8 @@
         "junit",
         "truth",
     ],
-    // TODO(b/366246777) uncomment it and the test.
-    // resource_apk: "RavenwoodBivalentInstTestTarget",
+    resource_apk: "RavenwoodBivalentInstTestTarget",
+    inst_resource_apk: "RavenwoodBivalentInstTest_nonself_inst_device",
     auto_gen_config: true,
 }
 
diff --git a/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_nonself.java b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_nonself.java
index 9f3ca6f..92d43d7 100644
--- a/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_nonself.java
+++ b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_nonself.java
@@ -19,7 +19,6 @@
 
 import android.app.Instrumentation;
 import android.content.Context;
-import android.platform.test.annotations.DisabledOnRavenwood;
 import android.platform.test.ravenwood.RavenwoodConfig;
 import android.platform.test.ravenwood.RavenwoodConfig.Config;
 
@@ -97,7 +96,6 @@
     }
 
     @Test
-    @DisabledOnRavenwood(reason = "b/366246777")
     public void testTargetAppResource() {
         assertThat(sTargetContext.getString(
                 com.android.ravenwood.bivalentinst_target_app.R.string.test_string_in_target))
@@ -105,8 +103,6 @@
     }
 
     @Test
-    @DisabledOnRavenwood(
-            reason = "Loading resources from non-self-instrumenting test APK isn't supported yet")
     public void testTestAppResource() {
         assertThat(sTestContext.getString(
                 com.android.ravenwood.bivalentinsttest_nonself_inst.R.string.test_string_in_test))
diff --git a/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_self.java b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_self.java
index fdff222..2f35923 100644
--- a/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_self.java
+++ b/ravenwood/tests/bivalentinst/test/com/android/ravenwoodtest/bivalentinst/RavenwoodInstrumentationTest_self.java
@@ -19,7 +19,6 @@
 
 import android.app.Instrumentation;
 import android.content.Context;
-import android.platform.test.annotations.DisabledOnRavenwood;
 import android.platform.test.ravenwood.RavenwoodConfig;
 import android.platform.test.ravenwood.RavenwoodConfig.Config;
 
@@ -109,7 +108,6 @@
     }
 
     @Test
-    @DisabledOnRavenwood(reason = "b/366246777")
     public void testTargetAppResource() {
         assertThat(sTargetContext.getString(
                 com.android.ravenwood.bivalentinsttest_self_inst.R.string.test_string_in_test))
@@ -117,7 +115,6 @@
     }
 
     @Test
-    @DisabledOnRavenwood(reason = "b/366246777")
     public void testTestAppResource() {
         assertThat(sTestContext.getString(
                 com.android.ravenwood.bivalentinsttest_self_inst.R.string.test_string_in_test))
diff --git a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionExecutors.java b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionExecutors.java
index c3b7087..1f98334 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionExecutors.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionExecutors.java
@@ -16,7 +16,15 @@
 
 package com.android.server.appfunctions;
 
+import android.annotation.NonNull;
+import android.os.UserHandle;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.GuardedBy;
+
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -33,5 +41,50 @@
                     /* unit= */ TimeUnit.SECONDS,
                     /* workQueue= */ new LinkedBlockingQueue<>());
 
+    /** A map of per-user executors for queued work. */
+    @GuardedBy("sLock")
+    private static final SparseArray<ExecutorService> mPerUserExecutorsLocked = new SparseArray<>();
+
+    private static final Object sLock = new Object();
+
+    /**
+     * Returns a per-user executor for queued metadata sync request.
+     *
+     * <p>The work submitted to these executor (Sync request) needs to be synchronous per user hence
+     * the use of a single thread.
+     *
+     * <p>Note: Use a different executor if not calling {@code submitSyncRequest} on a {@code
+     * MetadataSyncAdapter}.
+     */
+    // TODO(b/357551503): Restrict the scope of this executor to the MetadataSyncAdapter itself.
+    public static ExecutorService getPerUserSyncExecutor(@NonNull UserHandle user) {
+        synchronized (sLock) {
+            ExecutorService executor = mPerUserExecutorsLocked.get(user.getIdentifier(), null);
+            if (executor == null) {
+                executor = Executors.newSingleThreadExecutor();
+                mPerUserExecutorsLocked.put(user.getIdentifier(), executor);
+            }
+            return executor;
+        }
+    }
+
+    /**
+     * Shuts down and removes the per-user executor for queued work.
+     *
+     * <p>This should be called when the user is removed.
+     */
+    public static void shutDownAndRemoveUserExecutor(@NonNull UserHandle user)
+            throws InterruptedException {
+        ExecutorService executor;
+        synchronized (sLock) {
+            executor = mPerUserExecutorsLocked.get(user.getIdentifier());
+            mPerUserExecutorsLocked.remove(user.getIdentifier());
+        }
+        if (executor != null) {
+            executor.shutdown();
+            var unused = executor.awaitTermination(30, TimeUnit.SECONDS);
+        }
+    }
+
     private AppFunctionExecutors() {}
 }
diff --git a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
index 1e723b5..b4713d9 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
@@ -95,7 +95,12 @@
     public void onUserStopping(@NonNull TargetUser user) {
         Objects.requireNonNull(user);
 
-        MetadataSyncPerUser.removeUserSyncAdapter(user.getUserHandle());
+        try {
+            AppFunctionExecutors.shutDownAndRemoveUserExecutor(user.getUserHandle());
+            MetadataSyncPerUser.removeUserSyncAdapter(user.getUserHandle());
+        } catch (InterruptedException e) {
+            Slog.e(TAG, "Unable to remove data for: " + user.getUserHandle(), e);
+        }
     }
 
     @Override
diff --git a/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java b/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java
index 759f02e..e29b6e4 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncAdapter.java
@@ -42,7 +42,6 @@
 import android.util.ArraySet;
 import android.util.Slog;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.infra.AndroidFuture;
 import com.android.server.appfunctions.FutureAppSearchSession.FutureSearchResults;
@@ -53,12 +52,8 @@
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.Executor;
 
 /**
  * This class implements helper methods for synchronously interacting with AppSearch while
@@ -68,14 +63,9 @@
  */
 public class MetadataSyncAdapter {
     private static final String TAG = MetadataSyncAdapter.class.getSimpleName();
-
-    private final ExecutorService mExecutor;
-
+    private final Executor mSyncExecutor;
     private final AppSearchManager mAppSearchManager;
     private final PackageManager mPackageManager;
-    private final Object mLock = new Object();
-    @GuardedBy("mLock")
-    private Future<AndroidFuture<Boolean>> mCurrentSyncTask;
 
     // Hidden constants in {@link SetSchemaRequest} that restricts runtime metadata visibility
     // by permissions.
@@ -83,10 +73,12 @@
     public static final int EXECUTE_APP_FUNCTIONS_TRUSTED = 10;
 
     public MetadataSyncAdapter(
-            @NonNull PackageManager packageManager, @NonNull AppSearchManager appSearchManager) {
+            @NonNull Executor syncExecutor,
+            @NonNull PackageManager packageManager,
+            @NonNull AppSearchManager appSearchManager) {
+        mSyncExecutor = Objects.requireNonNull(syncExecutor);
         mPackageManager = Objects.requireNonNull(packageManager);
         mAppSearchManager = Objects.requireNonNull(appSearchManager);
-        mExecutor = Executors.newSingleThreadExecutor();
     }
 
     /**
@@ -105,7 +97,7 @@
                                 AppFunctionRuntimeMetadata.APP_FUNCTION_RUNTIME_METADATA_DB)
                         .build();
         AndroidFuture<Boolean> settableSyncStatus = new AndroidFuture<>();
-        Callable<AndroidFuture<Boolean>> callableTask =
+        mSyncExecutor.execute(
                 () -> {
                     try (FutureAppSearchSession staticMetadataSearchSession =
                                     new FutureAppSearchSessionImpl(
@@ -125,28 +117,10 @@
                     } catch (Exception ex) {
                         settableSyncStatus.completeExceptionally(ex);
                     }
-                    return settableSyncStatus;
-                };
-
-        synchronized (mLock) {
-            if (mCurrentSyncTask != null && !mCurrentSyncTask.isDone()) {
-                boolean cancel = mCurrentSyncTask.cancel(false);
-            }
-            mCurrentSyncTask = mExecutor.submit(callableTask);
-        }
-
+                });
         return settableSyncStatus;
     }
 
-    /** This method shuts down the {@link MetadataSyncAdapter} scheduler. */
-    public void shutDown() {
-        try {
-            var unused = mExecutor.awaitTermination(30, TimeUnit.SECONDS);
-        } catch (InterruptedException e) {
-            Slog.e(TAG, "Error shutting down MetadataSyncAdapter scheduler", e);
-        }
-    }
-
     @WorkerThread
     @VisibleForTesting
     void trySyncAppFunctionMetadataBlocking(
diff --git a/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncPerUser.java b/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncPerUser.java
index e933ec1..f421527 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncPerUser.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/MetadataSyncPerUser.java
@@ -55,7 +55,10 @@
                 PackageManager perUserPackageManager = userContext.getPackageManager();
                 if (perUserAppSearchManager != null) {
                     metadataSyncAdapter =
-                            new MetadataSyncAdapter(perUserPackageManager, perUserAppSearchManager);
+                            new MetadataSyncAdapter(
+                                    AppFunctionExecutors.getPerUserSyncExecutor(user),
+                                    perUserPackageManager,
+                                    perUserAppSearchManager);
                     sPerUserMetadataSyncAdapter.put(user.getIdentifier(), metadataSyncAdapter);
                     return metadataSyncAdapter;
                 }
@@ -71,12 +74,7 @@
      */
     public static void removeUserSyncAdapter(UserHandle user) {
         synchronized (sLock) {
-            MetadataSyncAdapter metadataSyncAdapter =
-                    sPerUserMetadataSyncAdapter.get(user.getIdentifier(), null);
-            if (metadataSyncAdapter != null) {
-                metadataSyncAdapter.shutDown();
-                sPerUserMetadataSyncAdapter.remove(user.getIdentifier());
-            }
+            sPerUserMetadataSyncAdapter.remove(user.getIdentifier());
         }
     }
 }
diff --git a/services/core/java/com/android/server/audio/AudioService.java b/services/core/java/com/android/server/audio/AudioService.java
index 99404428..9e11081 100644
--- a/services/core/java/com/android/server/audio/AudioService.java
+++ b/services/core/java/com/android/server/audio/AudioService.java
@@ -13383,19 +13383,39 @@
     }
 
     @android.annotation.EnforcePermission(MODIFY_AUDIO_ROUTING)
-    /** @see AudioPolicy#getFocusStack() */
+    /* @see AudioPolicy#getFocusStack() */
     public List<AudioFocusInfo> getFocusStack() {
         super.getFocusStack_enforcePermission();
 
         return mMediaFocusControl.getFocusStack();
     }
 
-    /** @see AudioPolicy#sendFocusLoss */
+    /**
+     * @param focusLoser non-null entry that may be in the stack
+     * @see AudioPolicy#sendFocusLossAndUpdate(AudioFocusInfo)
+     */
+    @android.annotation.EnforcePermission(MODIFY_AUDIO_ROUTING)
+    public void sendFocusLossAndUpdate(@NonNull AudioFocusInfo focusLoser,
+            @NonNull IAudioPolicyCallback apcb) {
+        super.sendFocusLossAndUpdate_enforcePermission();
+        Objects.requireNonNull(apcb);
+        if (!mAudioPolicies.containsKey(apcb.asBinder())) {
+            throw new IllegalStateException("Only registered AudioPolicy can change focus");
+        }
+        if (!mAudioPolicies.get(apcb.asBinder()).mHasFocusListener) {
+            throw new IllegalStateException("AudioPolicy must have focus listener to change focus");
+        }
+
+        mMediaFocusControl.sendFocusLossAndUpdate(Objects.requireNonNull(focusLoser));
+    }
+
+    /* @see AudioPolicy#sendFocusLoss(AudioFocusInfo)  */
+    @android.annotation.EnforcePermission(MODIFY_AUDIO_ROUTING)
     public boolean sendFocusLoss(@NonNull AudioFocusInfo focusLoser,
             @NonNull IAudioPolicyCallback apcb) {
+        super.sendFocusLoss_enforcePermission();
         Objects.requireNonNull(focusLoser);
         Objects.requireNonNull(apcb);
-        enforceModifyAudioRoutingPermission();
         if (!mAudioPolicies.containsKey(apcb.asBinder())) {
             throw new IllegalStateException("Only registered AudioPolicy can change focus");
         }
diff --git a/services/core/java/com/android/server/audio/MediaFocusControl.java b/services/core/java/com/android/server/audio/MediaFocusControl.java
index 7e26356..b4af46e 100644
--- a/services/core/java/com/android/server/audio/MediaFocusControl.java
+++ b/services/core/java/com/android/server/audio/MediaFocusControl.java
@@ -280,6 +280,37 @@
     }
 
     /**
+     * Like {@link #sendFocusLoss(AudioFocusInfo)} but if the loser was at the top of stack,
+     * make the next entry gain focus with {@link AudioManager#AUDIOFOCUS_GAIN}.
+     * @param focusInfo the focus owner to discard
+     * @see AudioPolicy#sendFocusLossAndUpdate(AudioFocusInfo)
+     */
+    protected void sendFocusLossAndUpdate(@NonNull AudioFocusInfo focusInfo) {
+        synchronized (mAudioFocusLock) {
+            if (mFocusStack.isEmpty()) {
+                return;
+            }
+            final FocusRequester currentFocusOwner = mFocusStack.peek();
+            if (currentFocusOwner.toAudioFocusInfo().equals(focusInfo)) {
+                // focus loss is for the top of the stack
+                currentFocusOwner.handleFocusLoss(AudioManager.AUDIOFOCUS_LOSS, null,
+                            false /*forceDuck*/);
+                currentFocusOwner.release();
+
+                mFocusStack.pop();
+                // is there a new focus owner?
+                if (!mFocusStack.isEmpty()) {
+                    mFocusStack.peek().handleFocusGain(AudioManager.AUDIOFOCUS_GAIN);
+                }
+            } else {
+                // focus loss if for another entry that's not at the top of the stack,
+                // just remove it from the stack and make it lose focus
+                sendFocusLoss(focusInfo);
+            }
+        }
+    }
+
+    /**
      * Return a copy of the focus stack for external consumption (composed of AudioFocusInfo
      * instead of FocusRequester instances)
      * @return a SystemApi-friendly version of the focus stack, in the same order (last entry
diff --git a/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamper.java b/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamper.java
index cf44ac0..a1fd164 100644
--- a/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamper.java
+++ b/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamper.java
@@ -74,6 +74,5 @@
 
     protected enum Type {
         POWER,
-        WEAR_BEDTIME_MODE,
     }
 }
diff --git a/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java b/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java
index 9404034..a10094f 100644
--- a/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java
+++ b/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java
@@ -218,9 +218,7 @@
             return BrightnessInfo.BRIGHTNESS_MAX_REASON_NONE;
         } else if (mClamperType == Type.POWER) {
             return BrightnessInfo.BRIGHTNESS_MAX_REASON_POWER_IC;
-        } else if (mClamperType == Type.WEAR_BEDTIME_MODE) {
-            return BrightnessInfo.BRIGHTNESS_MAX_REASON_WEAR_BEDTIME_MODE;
-        } else {
+        }  else {
             Slog.wtf(TAG, "BrightnessMaxReason not mapped for type=" + mClamperType);
             return BrightnessInfo.BRIGHTNESS_MAX_REASON_NONE;
         }
@@ -350,10 +348,6 @@
                             data, currentBrightness));
                 }
             }
-            if (flags.isBrightnessWearBedtimeModeClamperEnabled()) {
-                clampers.add(new BrightnessWearBedtimeModeClamper(handler, context,
-                        clamperChangeListener, data));
-            }
             return clampers;
         }
 
@@ -362,6 +356,10 @@
                 DisplayDeviceData data) {
             List<BrightnessStateModifier> modifiers = new ArrayList<>();
             modifiers.add(new BrightnessThermalModifier(handler, listener, data));
+            if (flags.isBrightnessWearBedtimeModeClamperEnabled()) {
+                modifiers.add(new BrightnessWearBedtimeModeModifier(handler, context,
+                        listener, data));
+            }
 
             modifiers.add(new DisplayDimModifier(context));
             modifiers.add(new BrightnessLowPowerModeModifier());
@@ -395,7 +393,7 @@
      */
     public static class DisplayDeviceData implements BrightnessThermalModifier.ThermalData,
             BrightnessPowerClamper.PowerData,
-            BrightnessWearBedtimeModeClamper.WearBedtimeModeData {
+            BrightnessWearBedtimeModeModifier.WearBedtimeModeData {
         @NonNull
         private final String mUniqueDisplayId;
         @NonNull
diff --git a/services/core/java/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeClamper.java b/services/core/java/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeClamper.java
deleted file mode 100644
index 1902e35..0000000
--- a/services/core/java/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeClamper.java
+++ /dev/null
@@ -1,99 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.server.display.brightness.clamper;
-
-import android.annotation.NonNull;
-import android.content.ContentResolver;
-import android.content.Context;
-import android.database.ContentObserver;
-import android.os.Handler;
-import android.os.UserHandle;
-import android.provider.Settings;
-
-import com.android.internal.annotations.VisibleForTesting;
-
-public class BrightnessWearBedtimeModeClamper extends
-        BrightnessClamper<BrightnessWearBedtimeModeClamper.WearBedtimeModeData> {
-
-    public static final int BEDTIME_MODE_OFF = 0;
-    public static final int BEDTIME_MODE_ON = 1;
-
-    private final Context mContext;
-
-    private final ContentObserver mSettingsObserver;
-
-    BrightnessWearBedtimeModeClamper(Handler handler, Context context,
-            BrightnessClamperController.ClamperChangeListener listener, WearBedtimeModeData data) {
-        this(new Injector(), handler, context, listener, data);
-    }
-
-    @VisibleForTesting
-    BrightnessWearBedtimeModeClamper(Injector injector, Handler handler, Context context,
-            BrightnessClamperController.ClamperChangeListener listener, WearBedtimeModeData data) {
-        super(handler, listener);
-        mContext = context;
-        mBrightnessCap = data.getBrightnessWearBedtimeModeCap();
-        mSettingsObserver = new ContentObserver(mHandler) {
-            @Override
-            public void onChange(boolean selfChange) {
-                final int bedtimeModeSetting = Settings.Global.getInt(
-                        mContext.getContentResolver(),
-                        Settings.Global.Wearable.BEDTIME_MODE,
-                        BEDTIME_MODE_OFF);
-                mIsActive = bedtimeModeSetting == BEDTIME_MODE_ON;
-                mChangeListener.onChanged();
-            }
-        };
-        injector.registerBedtimeModeObserver(context.getContentResolver(), mSettingsObserver);
-    }
-
-    @NonNull
-    @Override
-    Type getType() {
-        return Type.WEAR_BEDTIME_MODE;
-    }
-
-    @Override
-    void onDeviceConfigChanged() {}
-
-    @Override
-    void onDisplayChanged(WearBedtimeModeData displayData) {
-        mHandler.post(() -> {
-            mBrightnessCap = displayData.getBrightnessWearBedtimeModeCap();
-            mChangeListener.onChanged();
-        });
-    }
-
-    @Override
-    void stop() {
-        mContext.getContentResolver().unregisterContentObserver(mSettingsObserver);
-    }
-
-    interface WearBedtimeModeData {
-        float getBrightnessWearBedtimeModeCap();
-    }
-
-    @VisibleForTesting
-    static class Injector {
-        void registerBedtimeModeObserver(@NonNull ContentResolver cr,
-                @NonNull ContentObserver observer) {
-            cr.registerContentObserver(
-                    Settings.Global.getUriFor(Settings.Global.Wearable.BEDTIME_MODE),
-                    /* notifyForDescendants= */ false, observer, UserHandle.USER_ALL);
-        }
-    }
-}
diff --git a/services/core/java/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeModifier.java b/services/core/java/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeModifier.java
new file mode 100644
index 0000000..c9c8c33
--- /dev/null
+++ b/services/core/java/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeModifier.java
@@ -0,0 +1,158 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.brightness.clamper;
+
+import android.annotation.NonNull;
+import android.content.ContentResolver;
+import android.content.Context;
+import android.database.ContentObserver;
+import android.hardware.display.BrightnessInfo;
+import android.hardware.display.DisplayManagerInternal;
+import android.os.Handler;
+import android.os.UserHandle;
+import android.provider.Settings;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.display.DisplayBrightnessState;
+import com.android.server.display.brightness.BrightnessReason;
+
+import java.io.PrintWriter;
+
+public class BrightnessWearBedtimeModeModifier implements BrightnessStateModifier,
+        BrightnessClamperController.DisplayDeviceDataListener,
+        BrightnessClamperController.StatefulModifier {
+
+    public static final int BEDTIME_MODE_OFF = 0;
+    public static final int BEDTIME_MODE_ON = 1;
+
+    private final Context mContext;
+
+    private final ContentObserver mSettingsObserver;
+    protected final Handler mHandler;
+    protected final BrightnessClamperController.ClamperChangeListener mChangeListener;
+
+    private float mBrightnessCap;
+    private boolean mIsActive = false;
+    private boolean mApplied = false;
+
+    BrightnessWearBedtimeModeModifier(Handler handler, Context context,
+            BrightnessClamperController.ClamperChangeListener listener, WearBedtimeModeData data) {
+        this(new Injector(), handler, context, listener, data);
+    }
+
+    @VisibleForTesting
+    BrightnessWearBedtimeModeModifier(Injector injector, Handler handler, Context context,
+            BrightnessClamperController.ClamperChangeListener listener, WearBedtimeModeData data) {
+        mHandler = handler;
+        mChangeListener = listener;
+        mContext = context;
+        mBrightnessCap = data.getBrightnessWearBedtimeModeCap();
+        mSettingsObserver = new ContentObserver(mHandler) {
+            @Override
+            public void onChange(boolean selfChange) {
+                final int bedtimeModeSetting = Settings.Global.getInt(
+                        mContext.getContentResolver(),
+                        Settings.Global.Wearable.BEDTIME_MODE,
+                        BEDTIME_MODE_OFF);
+                mIsActive = bedtimeModeSetting == BEDTIME_MODE_ON;
+                mChangeListener.onChanged();
+            }
+        };
+        injector.registerBedtimeModeObserver(context.getContentResolver(), mSettingsObserver);
+    }
+
+    //region BrightnessStateModifier
+    @Override
+    public void apply(DisplayManagerInternal.DisplayPowerRequest request,
+            DisplayBrightnessState.Builder stateBuilder) {
+        if (mIsActive && stateBuilder.getMaxBrightness() > mBrightnessCap) {
+            stateBuilder.setMaxBrightness(mBrightnessCap);
+            stateBuilder.setBrightness(Math.min(stateBuilder.getBrightness(), mBrightnessCap));
+            stateBuilder.setBrightnessMaxReason(
+                    BrightnessInfo.BRIGHTNESS_MAX_REASON_WEAR_BEDTIME_MODE);
+            stateBuilder.getBrightnessReason().addModifier(BrightnessReason.MODIFIER_THROTTLED);
+            // set fast change only when modifier is activated.
+            // this will allow auto brightness to apply slow change even when modifier is active
+            if (!mApplied) {
+                stateBuilder.setIsSlowChange(false);
+            }
+            mApplied = true;
+        } else {
+            mApplied = false;
+        }
+    }
+
+    @Override
+    public void stop() {
+        mContext.getContentResolver().unregisterContentObserver(mSettingsObserver);
+    }
+
+    @Override
+    public void dump(PrintWriter writer) {
+        writer.println("BrightnessWearBedtimeModeModifier:");
+        writer.println("  mBrightnessCap: " + mBrightnessCap);
+        writer.println("  mIsActive: " + mIsActive);
+        writer.println("  mApplied: " + mApplied);
+    }
+
+    @Override
+    public boolean shouldListenToLightSensor() {
+        return false;
+    }
+
+    @Override
+    public void setAmbientLux(float lux) {
+        // noop
+    }
+    //endregion
+
+    //region DisplayDeviceDataListener
+    @Override
+    public void onDisplayChanged(BrightnessClamperController.DisplayDeviceData data) {
+        mHandler.post(() -> {
+            mBrightnessCap = data.getBrightnessWearBedtimeModeCap();
+            mChangeListener.onChanged();
+        });
+    }
+    //endregion
+
+    //region StatefulModifier
+    @Override
+    public void applyStateChange(
+            BrightnessClamperController.ModifiersAggregatedState aggregatedState) {
+        if (mIsActive && aggregatedState.mMaxBrightness > mBrightnessCap) {
+            aggregatedState.mMaxBrightness = mBrightnessCap;
+            aggregatedState.mMaxBrightnessReason =
+                    BrightnessInfo.BRIGHTNESS_MAX_REASON_WEAR_BEDTIME_MODE;
+        }
+    }
+    //endregion
+
+    interface WearBedtimeModeData {
+        float getBrightnessWearBedtimeModeCap();
+    }
+
+    @VisibleForTesting
+    static class Injector {
+        void registerBedtimeModeObserver(@NonNull ContentResolver cr,
+                @NonNull ContentObserver observer) {
+            cr.registerContentObserver(
+                    Settings.Global.getUriFor(Settings.Global.Wearable.BEDTIME_MODE),
+                    /* notifyForDescendants= */ false, observer, UserHandle.USER_ALL);
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/input/InputManagerInternal.java b/services/core/java/com/android/server/input/InputManagerInternal.java
index 73f18d1..9281267 100644
--- a/services/core/java/com/android/server/input/InputManagerInternal.java
+++ b/services/core/java/com/android/server/input/InputManagerInternal.java
@@ -232,6 +232,9 @@
     /**
      * Notify key gesture was completed by the user.
      *
+     * NOTE: This is a temporary API added to assist in a long-term refactor, and is not meant for
+     * general use by system services.
+     *
      * @param deviceId the device ID of the keyboard using which the event was completed
      * @param keycodes the keys pressed for the event
      * @param modifierState the modifier state
@@ -240,4 +243,20 @@
      */
     public abstract void notifyKeyGestureCompleted(int deviceId, int[] keycodes, int modifierState,
             @KeyGestureEvent.KeyGestureType int event);
+
+    /**
+     * Notify that a key gesture was detected by another system component, and it should be handled
+     * appropriately by KeyGestureController.
+     *
+     * NOTE: This is a temporary API added to assist in a long-term refactor, and is not meant for
+     * general use by system services.
+     *
+     * @param deviceId the device ID of the keyboard using which the event was completed
+     * @param keycodes the keys pressed for the event
+     * @param modifierState the modifier state
+     * @param event the gesture event that was completed
+     *
+     */
+    public abstract void handleKeyGestureInKeyGestureController(int deviceId, int[] keycodes,
+            int modifierState, @KeyGestureEvent.KeyGestureType int event);
 }
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index 65adaba..fd7479e 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -3408,6 +3408,12 @@
             mKeyGestureController.notifyKeyGestureCompleted(deviceId, keycodes, modifierState,
                     gestureType);
         }
+
+        @Override
+        public void handleKeyGestureInKeyGestureController(int deviceId, int[] keycodes,
+                int modifierState, @KeyGestureEvent.KeyGestureType int gestureType) {
+            mKeyGestureController.handleKeyGesture(deviceId, keycodes, modifierState, gestureType);
+        }
     }
 
     @Override
diff --git a/services/core/java/com/android/server/input/KeyGestureController.java b/services/core/java/com/android/server/input/KeyGestureController.java
index 7fe7891..4538b49 100644
--- a/services/core/java/com/android/server/input/KeyGestureController.java
+++ b/services/core/java/com/android/server/input/KeyGestureController.java
@@ -24,7 +24,6 @@
 import android.annotation.SuppressLint;
 import android.content.Context;
 import android.content.res.Resources;
-
 import android.hardware.input.AidlKeyGestureEvent;
 import android.hardware.input.IKeyGestureEventListener;
 import android.hardware.input.IKeyGestureHandler;
@@ -582,8 +581,11 @@
     boolean handleKeyGesture(int deviceId, int[] keycodes, int modifierState,
             @KeyGestureEvent.KeyGestureType int gestureType, int action, int displayId,
             IBinder focusedToken, int flags) {
-        AidlKeyGestureEvent event = createKeyGestureEvent(deviceId, keycodes,
-                modifierState, gestureType, action, displayId, flags);
+        return handleKeyGesture(createKeyGestureEvent(deviceId, keycodes,
+                modifierState, gestureType, action, displayId, flags), focusedToken);
+    }
+
+    private boolean handleKeyGesture(AidlKeyGestureEvent event, @Nullable IBinder focusedToken) {
         synchronized (mKeyGestureHandlerRecords) {
             for (KeyGestureHandlerRecord handler : mKeyGestureHandlerRecords.values()) {
                 if (handler.handleKeyGesture(event, focusedToken)) {
@@ -616,6 +618,13 @@
         mHandler.obtainMessage(MSG_NOTIFY_KEY_GESTURE_EVENT, event).sendToTarget();
     }
 
+    public void handleKeyGesture(int deviceId, int[] keycodes, int modifierState,
+            @KeyGestureEvent.KeyGestureType int gestureType) {
+        AidlKeyGestureEvent event = createKeyGestureEvent(deviceId, keycodes, modifierState,
+                gestureType, KeyGestureEvent.ACTION_GESTURE_COMPLETE, Display.DEFAULT_DISPLAY, 0);
+        handleKeyGesture(event, null /*focusedToken*/);
+    }
+
     @MainThread
     private void notifyKeyGestureEvent(AidlKeyGestureEvent event) {
         InputDevice device = getInputDevice(event.deviceId);
diff --git a/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java b/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java
index 5aea356..49a6ffd 100644
--- a/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java
+++ b/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java
@@ -191,6 +191,7 @@
     /**
      * Stop player proxy for the ongoing alarm and drop focus for its AudioFocusInfo.
      */
+    @SuppressLint("MissingPermission")
     @VisibleForTesting
     void muteAlarmSounds(Context context) {
         AudioManager audioManager = context.getSystemService(AudioManager.class);
@@ -201,6 +202,11 @@
                 }
             }
         }
+
+        AudioFocusInfo currentAfi = getAudioFocusInfoForNotification();
+        if (currentAfi != null) {
+            mFocusControlAudioPolicy.sendFocusLossAndUpdate(currentAfi);
+        }
     }
 
     /**
diff --git a/services/core/java/com/android/server/pm/UserManagerService.java b/services/core/java/com/android/server/pm/UserManagerService.java
index 89417f3..3e70d92 100644
--- a/services/core/java/com/android/server/pm/UserManagerService.java
+++ b/services/core/java/com/android/server/pm/UserManagerService.java
@@ -2647,11 +2647,15 @@
     }
 
     @Override
-    public int getMainDisplayIdAssignedToUser() {
-        // Not checking for any permission as it returns info about calling user
-        int userId = UserHandle.getUserId(Binder.getCallingUid());
-        int displayId = mUserVisibilityMediator.getMainDisplayAssignedToUser(userId);
-        return displayId;
+    public int getMainDisplayIdAssignedToUser(int userId) {
+        final int callingUserId = UserHandle.getCallingUserId();
+        if (callingUserId != userId
+                && !hasManageUsersOrPermission(android.Manifest.permission.INTERACT_ACROSS_USERS)) {
+            throw new SecurityException("Caller from user " + callingUserId + " needs MANAGE_USERS "
+                    + "or INTERACT_ACROSS_USERS permission to get the main display for (" + userId
+                    + ")");
+        }
+        return mUserVisibilityMediator.getMainDisplayAssignedToUser(userId);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 02c02b0..63491e8 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -47,6 +47,7 @@
 import static android.view.KeyEvent.KEYCODE_HOME;
 import static android.view.KeyEvent.KEYCODE_POWER;
 import static android.view.KeyEvent.KEYCODE_STEM_PRIMARY;
+import static android.view.KeyEvent.KEYCODE_STYLUS_BUTTON_TAIL;
 import static android.view.KeyEvent.KEYCODE_UNKNOWN;
 import static android.view.KeyEvent.KEYCODE_VOLUME_DOWN;
 import static android.view.KeyEvent.KEYCODE_VOLUME_UP;
@@ -2643,7 +2644,7 @@
         }
 
         @Override
-        void onKeyUp(long eventTime, int count, int displayId) {
+        void onKeyUp(long eventTime, int count, int displayId, int deviceId, int metaState) {
             if (mShouldEarlyShortPressOnPower && count == 1) {
                 powerPress(eventTime, 1 /*pressCount*/, displayId);
             }
@@ -2763,7 +2764,7 @@
         }
 
         @Override
-        void onKeyUp(long eventTime, int count, int unusedDisplayId) {
+        void onKeyUp(long eventTime, int count, int displayId, int deviceId, int metaState) {
             if (count == 1) {
                 // Save info about the most recent task on the first press of the stem key. This
                 // may be used later to switch to the most recent app using double press gesture.
@@ -2816,6 +2817,33 @@
         }
     }
 
+    // TODO(b/358569822): Move to KeyGestureController.
+    private final class StylusTailButtonRule extends SingleKeyGestureDetector.SingleKeyRule {
+        StylusTailButtonRule() {
+            super(KEYCODE_STYLUS_BUTTON_TAIL);
+        }
+
+        @Override
+        int getMaxMultiPressCount() {
+            return 2;
+        }
+
+        @Override
+        void onPress(long downTime, int displayId) {
+
+        }
+
+        @Override
+        void onKeyUp(long eventTime, int pressCount, int displayId, int deviceId, int metaState) {
+            if (pressCount != 1) {
+                return;
+            }
+            // Single press on tail button triggers the open notes gesture.
+            handleKeyGestureInKeyGestureController(KeyGestureEvent.KEY_GESTURE_TYPE_OPEN_NOTES,
+                    deviceId, KEYCODE_STYLUS_BUTTON_TAIL, metaState);
+        }
+    }
+
     private void initSingleKeyGestureRules(Looper looper) {
         mSingleKeyGestureDetector = SingleKeyGestureDetector.get(mContext, looper);
         mSingleKeyGestureDetector.addRule(new PowerKeyRule());
@@ -2825,6 +2853,7 @@
         if (hasStemPrimaryBehavior()) {
             mSingleKeyGestureDetector.addRule(new StemPrimaryKeyRule());
         }
+        mSingleKeyGestureDetector.addRule(new StylusTailButtonRule());
     }
 
     /**
@@ -3314,6 +3343,16 @@
                 new int[]{event.getKeyCode()}, event.getMetaState(), gestureType);
     }
 
+    private void handleKeyGestureInKeyGestureController(
+            @KeyGestureEvent.KeyGestureType int gestureType, int deviceId, int keyCode,
+            int metaState) {
+        if (gestureType == KeyGestureEvent.KEY_GESTURE_TYPE_UNSPECIFIED) {
+            return;
+        }
+        mInputManagerInternal.handleKeyGestureInKeyGestureController(deviceId, new int[]{keyCode},
+                metaState, gestureType);
+    }
+
     @Override
     public KeyboardShortcutGroup getApplicationLaunchKeyboardShortcuts(int deviceId) {
         return mModifierShortcutManager.getApplicationLaunchKeyboardShortcuts(deviceId);
diff --git a/services/core/java/com/android/server/policy/SingleKeyGestureDetector.java b/services/core/java/com/android/server/policy/SingleKeyGestureDetector.java
index a060f50..441d3ea 100644
--- a/services/core/java/com/android/server/policy/SingleKeyGestureDetector.java
+++ b/services/core/java/com/android/server/policy/SingleKeyGestureDetector.java
@@ -105,9 +105,9 @@
 
         /**
          *  Maximum count of multi presses.
-         *  Return 1 will trigger onPress immediately when {@link KeyEvent.ACTION_UP}.
+         *  Return 1 will trigger onPress immediately when {@link KeyEvent#ACTION_UP}.
          *  Otherwise trigger onMultiPress immediately when reach max count when
-         *  {@link KeyEvent.ACTION_DOWN}.
+         *  {@link KeyEvent#ACTION_DOWN}.
          */
         int getMaxMultiPressCount() {
             return 1;
@@ -153,8 +153,10 @@
          * @param eventTime  the timestamp of this event
          * @param pressCount the number of presses detected leading up to this key up event
          * @param displayId  the display ID of the event
+         * @param deviceId   the ID of the input device that generated this event
+         * @param metaState  the state of the modifiers when this gesture was detected
          */
-        void onKeyUp(long eventTime, int pressCount, int displayId) {}
+        void onKeyUp(long eventTime, int pressCount, int displayId, int deviceId, int metaState) {}
 
         @Override
         public String toString() {
@@ -183,7 +185,11 @@
     }
 
     private record MessageObject(SingleKeyRule activeRule, int keyCode, int pressCount,
-                                 int displayId) {
+                                 int displayId, int metaState, int deviceId) {
+        MessageObject(SingleKeyRule activeRule, int keyCode, int pressCount, KeyEvent event) {
+            this(activeRule, keyCode, pressCount, event.getDisplayId(), event.getMetaState(),
+                    event.getDeviceId());
+        }
     }
 
     static SingleKeyGestureDetector get(Context context, Looper looper) {
@@ -236,7 +242,7 @@
                 mHandler.removeMessages(MSG_KEY_LONG_PRESS);
                 mHandler.removeMessages(MSG_KEY_VERY_LONG_PRESS);
                 MessageObject object = new MessageObject(mActiveRule, keyCode, /* pressCount= */ 1,
-                        event.getDisplayId());
+                        event);
                 final Message msg = mHandler.obtainMessage(MSG_KEY_LONG_PRESS, object);
                 msg.setAsynchronous(true);
                 mHandler.sendMessage(msg);
@@ -284,7 +290,7 @@
         if (mKeyPressCounter == 1) {
             if (mActiveRule.supportLongPress()) {
                 MessageObject object = new MessageObject(mActiveRule, keyCode, mKeyPressCounter,
-                        event.getDisplayId());
+                        event);
                 final Message msg = mHandler.obtainMessage(MSG_KEY_LONG_PRESS, object);
                 msg.setAsynchronous(true);
                 mHandler.sendMessageDelayed(msg, mActiveRule.getLongPressTimeoutMs());
@@ -292,7 +298,7 @@
 
             if (mActiveRule.supportVeryLongPress()) {
                 MessageObject object = new MessageObject(mActiveRule, keyCode, mKeyPressCounter,
-                        event.getDisplayId());
+                        event);
                 final Message msg = mHandler.obtainMessage(MSG_KEY_VERY_LONG_PRESS, object);
                 msg.setAsynchronous(true);
                 mHandler.sendMessageDelayed(msg, mActiveRule.getVeryLongPressTimeoutMs());
@@ -310,7 +316,7 @@
                             + " reached the max count " + mKeyPressCounter);
                 }
                 MessageObject object = new MessageObject(mActiveRule, keyCode, mKeyPressCounter,
-                        event.getDisplayId());
+                        event);
                 final Message msg = mHandler.obtainMessage(MSG_KEY_DELAYED_PRESS, object);
                 msg.setAsynchronous(true);
                 mHandler.sendMessage(msg);
@@ -351,7 +357,7 @@
         if (event.getKeyCode() == mActiveRule.mKeyCode) {
             // key-up action should always be triggered if not processed by long press.
             MessageObject object = new MessageObject(mActiveRule, mActiveRule.mKeyCode,
-                    mKeyPressCounter, event.getDisplayId());
+                    mKeyPressCounter, event);
             Message msgKeyUp = mHandler.obtainMessage(MSG_KEY_UP, object);
             msgKeyUp.setAsynchronous(true);
             mHandler.sendMessage(msgKeyUp);
@@ -362,7 +368,7 @@
                     Log.i(TAG, "press key " + KeyEvent.keyCodeToString(event.getKeyCode()));
                 }
                 object = new MessageObject(mActiveRule, mActiveRule.mKeyCode,
-                        /* pressCount= */ 1, event.getDisplayId());
+                        /* pressCount= */ 1, event);
                 Message msg = mHandler.obtainMessage(MSG_KEY_DELAYED_PRESS, object);
                 msg.setAsynchronous(true);
                 mHandler.sendMessage(msg);
@@ -373,7 +379,7 @@
             // This could be a multi-press.  Wait a little bit longer to confirm.
             if (mKeyPressCounter < mActiveRule.getMaxMultiPressCount()) {
                 object = new MessageObject(mActiveRule, mActiveRule.mKeyCode,
-                        mKeyPressCounter, event.getDisplayId());
+                        mKeyPressCounter, event);
                 Message msg = mHandler.obtainMessage(MSG_KEY_DELAYED_PRESS, object);
                 msg.setAsynchronous(true);
                 mHandler.sendMessageDelayed(msg, MULTI_PRESS_TIMEOUT);
@@ -452,7 +458,8 @@
                         Log.i(TAG, "Detect key up " + KeyEvent.keyCodeToString(keyCode)
                                 + " on display " + displayId);
                     }
-                    rule.onKeyUp(mLastDownTime, pressCount, displayId);
+                    rule.onKeyUp(mLastDownTime, pressCount, displayId, object.deviceId,
+                            object.metaState);
                     break;
                 case MSG_KEY_LONG_PRESS:
                     if (DEBUG) {
diff --git a/services/core/java/com/android/server/wallpaper/WallpaperDataParser.java b/services/core/java/com/android/server/wallpaper/WallpaperDataParser.java
index e3e83b3..74ca230 100644
--- a/services/core/java/com/android/server/wallpaper/WallpaperDataParser.java
+++ b/services/core/java/com/android/server/wallpaper/WallpaperDataParser.java
@@ -171,63 +171,9 @@
             stream = new FileInputStream(file);
             TypedXmlPullParser parser = Xml.resolvePullParser(stream);
 
-            int type;
-            do {
-                type = parser.next();
-                if (type == XmlPullParser.START_TAG) {
-                    String tag = parser.getName();
-                    if (("wp".equals(tag) && loadSystem) || ("kwp".equals(tag) && loadLock)) {
-                        if ("kwp".equals(tag)) {
-                            lockWallpaper = new WallpaperData(userId, FLAG_LOCK);
-                        }
-                        WallpaperData wallpaperToParse =
-                                "wp".equals(tag) ? wallpaper : lockWallpaper;
+            lockWallpaper = loadSettingsFromSerializer(parser, wallpaper, userId, loadSystem,
+                    loadLock, keepDimensionHints, wpdData);
 
-                        if (!multiCrop()) {
-                            parseWallpaperAttributes(parser, wallpaperToParse, keepDimensionHints);
-                        }
-
-                        String comp = parser.getAttributeValue(null, "component");
-                        if (removeNextWallpaperComponent()) {
-                            wallpaperToParse.setComponent(comp != null
-                                    ? ComponentName.unflattenFromString(comp)
-                                    : null);
-                            if (wallpaperToParse.getComponent() == null
-                                    || "android".equals(wallpaperToParse.getComponent()
-                                    .getPackageName())) {
-                                wallpaperToParse.setComponent(mImageWallpaper);
-                            }
-                        } else {
-                            wallpaperToParse.nextWallpaperComponent = comp != null
-                                    ? ComponentName.unflattenFromString(comp)
-                                    : null;
-                            if (wallpaperToParse.nextWallpaperComponent == null
-                                    || "android".equals(wallpaperToParse.nextWallpaperComponent
-                                    .getPackageName())) {
-                                wallpaperToParse.nextWallpaperComponent = mImageWallpaper;
-                            }
-                        }
-
-                        if (multiCrop()) {
-                            parseWallpaperAttributes(parser, wallpaperToParse, keepDimensionHints);
-                        }
-
-                        if (DEBUG) {
-                            Slog.v(TAG, "mWidth:" + wpdData.mWidth);
-                            Slog.v(TAG, "mHeight:" + wpdData.mHeight);
-                            Slog.v(TAG, "cropRect:" + wallpaper.cropHint);
-                            Slog.v(TAG, "primaryColors:" + wallpaper.primaryColors);
-                            Slog.v(TAG, "mName:" + wallpaper.name);
-                            if (removeNextWallpaperComponent()) {
-                                Slog.v(TAG, "mWallpaperComponent:" + wallpaper.getComponent());
-                            } else {
-                                Slog.v(TAG, "mNextWallpaperComponent:"
-                                        + wallpaper.nextWallpaperComponent);
-                            }
-                        }
-                    }
-                }
-            } while (type != XmlPullParser.END_DOCUMENT);
             success = true;
         } catch (FileNotFoundException e) {
             Slog.w(TAG, "no current wallpaper -- first boot?");
@@ -275,6 +221,75 @@
         return new WallpaperLoadingResult(wallpaper, lockWallpaper, success);
     }
 
+    // This method updates `wallpaper` in place, but returns `lockWallpaper`. This is because
+    // `wallpaper` already exists if it's being read per `loadSystem`, but `lockWallpaper` is
+    // created conditionally if there is lock screen wallpaper data to read.
+    @VisibleForTesting
+    WallpaperData loadSettingsFromSerializer(TypedXmlPullParser parser, WallpaperData wallpaper,
+            int userId, boolean loadSystem, boolean loadLock, boolean keepDimensionHints,
+            DisplayData wpdData) throws IOException, XmlPullParserException {
+        WallpaperData lockWallpaper = null;
+        int type;
+        do {
+            type = parser.next();
+            if (type == XmlPullParser.START_TAG) {
+                String tag = parser.getName();
+                if (("wp".equals(tag) && loadSystem) || ("kwp".equals(tag) && loadLock)) {
+                    if ("kwp".equals(tag)) {
+                        lockWallpaper = new WallpaperData(userId, FLAG_LOCK);
+                    }
+                    WallpaperData wallpaperToParse =
+                            "wp".equals(tag) ? wallpaper : lockWallpaper;
+
+                    if (!multiCrop()) {
+                        parseWallpaperAttributes(parser, wallpaperToParse, keepDimensionHints);
+                    }
+
+                    String comp = parser.getAttributeValue(null, "component");
+                    if (removeNextWallpaperComponent()) {
+                        wallpaperToParse.setComponent(comp != null
+                                ? ComponentName.unflattenFromString(comp)
+                                : null);
+                        if (wallpaperToParse.getComponent() == null
+                                || "android".equals(wallpaperToParse.getComponent()
+                                .getPackageName())) {
+                            wallpaperToParse.setComponent(mImageWallpaper);
+                        }
+                    } else {
+                        wallpaperToParse.nextWallpaperComponent = comp != null
+                                ? ComponentName.unflattenFromString(comp)
+                                : null;
+                        if (wallpaperToParse.nextWallpaperComponent == null
+                                || "android".equals(wallpaperToParse.nextWallpaperComponent
+                                .getPackageName())) {
+                            wallpaperToParse.nextWallpaperComponent = mImageWallpaper;
+                        }
+                    }
+
+                    if (multiCrop()) {
+                        parseWallpaperAttributes(parser, wallpaperToParse, keepDimensionHints);
+                    }
+
+                    if (DEBUG) {
+                        Slog.v(TAG, "mWidth:" + wpdData.mWidth);
+                        Slog.v(TAG, "mHeight:" + wpdData.mHeight);
+                        Slog.v(TAG, "cropRect:" + wallpaper.cropHint);
+                        Slog.v(TAG, "primaryColors:" + wallpaper.primaryColors);
+                        Slog.v(TAG, "mName:" + wallpaper.name);
+                        if (removeNextWallpaperComponent()) {
+                            Slog.v(TAG, "mWallpaperComponent:" + wallpaper.getComponent());
+                        } else {
+                            Slog.v(TAG, "mNextWallpaperComponent:"
+                                    + wallpaper.nextWallpaperComponent);
+                        }
+                    }
+                }
+            }
+        } while (type != XmlPullParser.END_DOCUMENT);
+
+        return lockWallpaper;
+    }
+
     private void ensureSaneWallpaperData(WallpaperData wallpaper) {
         // Only overwrite cropHint if the rectangle is invalid.
         if (wallpaper.cropHint.width() < 0
@@ -449,18 +464,7 @@
         try {
             fstream = new FileOutputStream(journal.chooseForWrite(), false);
             TypedXmlSerializer out = Xml.resolveSerializer(fstream);
-            out.startDocument(null, true);
-
-            if (wallpaper != null) {
-                writeWallpaperAttributes(out, "wp", wallpaper);
-            }
-
-            if (lockWallpaper != null) {
-                writeWallpaperAttributes(out, "kwp", lockWallpaper);
-            }
-
-            out.endDocument();
-
+            saveSettingsToSerializer(out, wallpaper, lockWallpaper);
             fstream.flush();
             FileUtils.sync(fstream);
             fstream.close();
@@ -472,6 +476,22 @@
     }
 
     @VisibleForTesting
+    void saveSettingsToSerializer(TypedXmlSerializer out, WallpaperData wallpaper,
+            WallpaperData lockWallpaper) throws IOException {
+        out.startDocument(null, true);
+
+        if (wallpaper != null) {
+            writeWallpaperAttributes(out, "wp", wallpaper);
+        }
+
+        if (lockWallpaper != null) {
+            writeWallpaperAttributes(out, "kwp", lockWallpaper);
+        }
+
+        out.endDocument();
+    }
+
+    @VisibleForTesting
     void writeWallpaperAttributes(TypedXmlSerializer out, String tag, WallpaperData wallpaper)
             throws IllegalArgumentException, IllegalStateException, IOException {
         if (DEBUG) {
diff --git a/services/core/java/com/android/server/wm/LaunchParamsPersister.java b/services/core/java/com/android/server/wm/LaunchParamsPersister.java
index b84ef37..2394da9 100644
--- a/services/core/java/com/android/server/wm/LaunchParamsPersister.java
+++ b/services/core/java/com/android/server/wm/LaunchParamsPersister.java
@@ -50,7 +50,6 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.CountDownLatch;
 import java.util.function.IntFunction;
 
 /**
@@ -93,12 +92,6 @@
             new SparseArray<>();
 
     /**
-     * A map from user ID to the active {@link LoadingQueueItem} user when we're loading the launch
-     * params for that user.
-     */
-    private final SparseArray<LoadingQueueItem> mLoadingItemMap = new SparseArray<>();
-
-    /**
      * A map from {@link android.content.pm.ActivityInfo.WindowLayout#windowLayoutAffinity} to
      * activity's component name for reverse queries from window layout affinities to activities.
      * Used to decide if we should use another activity's record with the same affinity.
@@ -124,30 +117,112 @@
     }
 
     void onUnlockUser(int userId) {
-        if (mLoadingItemMap.contains(userId)) {
-            Slog.e(TAG, "Duplicate onUnlockUser " + userId);
-            return;
-        }
-        final LoadingQueueItem item = new LoadingQueueItem(userId);
-        mLoadingItemMap.put(userId, item);
-        mPersisterQueue.addItem(item, /* flush */ false);
+        loadLaunchParams(userId);
     }
 
     void onCleanupUser(int userId) {
-        final LoadingQueueItem item = mLoadingItemMap.removeReturnOld(userId);
-        if (item != null) {
-            item.abort();
-
-            mPersisterQueue.removeItems(
-                    queueItem -> queueItem.mUserId == userId, LoadingQueueItem.class);
-        }
         mLaunchParamsMap.remove(userId);
     }
 
-    private void waitForLoading(int userId) {
-        final LoadingQueueItem item = mLoadingItemMap.get(userId);
-        if (item != null) {
-            item.waitUntilFinish();
+    private void loadLaunchParams(int userId) {
+        final List<File> filesToDelete = new ArrayList<>();
+        final File launchParamsFolder = getLaunchParamFolder(userId);
+        if (!launchParamsFolder.isDirectory()) {
+            Slog.i(TAG, "Didn't find launch param folder for user " + userId);
+            return;
+        }
+
+        final Set<String> packages = new ArraySet<>(mPackageList.getPackageNames());
+
+        final File[] paramsFiles = launchParamsFolder.listFiles();
+        final ArrayMap<ComponentName, PersistableLaunchParams> map =
+                new ArrayMap<>(paramsFiles.length);
+        mLaunchParamsMap.put(userId, map);
+
+        for (File paramsFile : paramsFiles) {
+            if (!paramsFile.isFile()) {
+                Slog.w(TAG, paramsFile.getAbsolutePath() + " is not a file.");
+                continue;
+            }
+            if (!paramsFile.getName().endsWith(LAUNCH_PARAMS_FILE_SUFFIX)) {
+                Slog.w(TAG, "Unexpected params file name: " + paramsFile.getName());
+                filesToDelete.add(paramsFile);
+                continue;
+            }
+            String paramsFileName = paramsFile.getName();
+            // Migrate all records from old separator to new separator.
+            final int oldSeparatorIndex =
+                    paramsFileName.indexOf(OLD_ESCAPED_COMPONENT_SEPARATOR);
+            if (oldSeparatorIndex != -1) {
+                if (paramsFileName.indexOf(
+                        OLD_ESCAPED_COMPONENT_SEPARATOR, oldSeparatorIndex + 1) != -1) {
+                    // Rare case. We have more than one old escaped component separator probably
+                    // because this app uses underscore in their package name. We can't distinguish
+                    // which one is the real separator so let's skip it.
+                    filesToDelete.add(paramsFile);
+                    continue;
+                }
+                paramsFileName = paramsFileName.replace(
+                        OLD_ESCAPED_COMPONENT_SEPARATOR, ESCAPED_COMPONENT_SEPARATOR);
+                final File newFile = new File(launchParamsFolder, paramsFileName);
+                if (paramsFile.renameTo(newFile)) {
+                    paramsFile = newFile;
+                } else {
+                    // Rare case. For some reason we can't rename the file. Let's drop this record
+                    // instead.
+                    filesToDelete.add(paramsFile);
+                    continue;
+                }
+            }
+            final String componentNameString = paramsFileName.substring(
+                    0 /* beginIndex */,
+                    paramsFileName.length() - LAUNCH_PARAMS_FILE_SUFFIX.length())
+                    .replace(ESCAPED_COMPONENT_SEPARATOR, ORIGINAL_COMPONENT_SEPARATOR);
+            final ComponentName name = ComponentName.unflattenFromString(
+                    componentNameString);
+            if (name == null) {
+                Slog.w(TAG, "Unexpected file name: " + paramsFileName);
+                filesToDelete.add(paramsFile);
+                continue;
+            }
+
+            if (!packages.contains(name.getPackageName())) {
+                // Rare case. PersisterQueue doesn't have a chance to remove files for removed
+                // packages last time.
+                filesToDelete.add(paramsFile);
+                continue;
+            }
+
+            try (InputStream in = new FileInputStream(paramsFile)) {
+                final PersistableLaunchParams params = new PersistableLaunchParams();
+                final TypedXmlPullParser parser = Xml.resolvePullParser(in);
+                int event;
+                while ((event = parser.next()) != XmlPullParser.END_DOCUMENT
+                        && event != XmlPullParser.END_TAG) {
+                    if (event != XmlPullParser.START_TAG) {
+                        continue;
+                    }
+
+                    final String tagName = parser.getName();
+                    if (!TAG_LAUNCH_PARAMS.equals(tagName)) {
+                        Slog.w(TAG, "Unexpected tag name: " + tagName);
+                        continue;
+                    }
+
+                    params.restore(paramsFile, parser);
+                }
+
+                map.put(name, params);
+                addComponentNameToLaunchParamAffinityMapIfNotNull(
+                        name, params.mWindowLayoutAffinity);
+            } catch (Exception e) {
+                Slog.w(TAG, "Failed to restore launch params for " + name, e);
+                filesToDelete.add(paramsFile);
+            }
+        }
+
+        if (!filesToDelete.isEmpty()) {
+            mPersisterQueue.addItem(new CleanUpComponentQueueItem(filesToDelete), true);
         }
     }
 
@@ -161,7 +236,6 @@
             return;
         }
         final int userId = task.mUserId;
-        waitForLoading(userId);
         PersistableLaunchParams params;
         ArrayMap<ComponentName, PersistableLaunchParams> map = mLaunchParamsMap.get(userId);
         if (map == null) {
@@ -223,7 +297,6 @@
     void getLaunchParams(Task task, ActivityRecord activity, LaunchParams outParams) {
         final ComponentName name = task != null ? task.realActivity : activity.mActivityComponent;
         final int userId = task != null ? task.mUserId : activity.mUserId;
-        waitForLoading(userId);
         final String windowLayoutAffinity;
         if (task != null) {
             windowLayoutAffinity = task.mWindowLayoutAffinity;
@@ -321,156 +394,6 @@
         }
     }
 
-    /**
-     * The work item used to load launch parameters with {@link PersisterQueue} in a background
-     * thread, so that we don't block the thread {@link com.android.server.am.UserController} uses
-     * to broadcast user state changes for I/O operations. See b/365983567 for more details.
-     */
-    private class LoadingQueueItem implements PersisterQueue.QueueItem {
-        private final int mUserId;
-        private final CountDownLatch mLatch = new CountDownLatch(1);
-        private boolean mAborted = false;
-
-        private LoadingQueueItem(int userId) {
-            mUserId = userId;
-        }
-
-        @Override
-        public void process() {
-            try {
-                loadLaunchParams();
-            } finally {
-                synchronized (mSupervisor.mService.getGlobalLock()) {
-                    mLoadingItemMap.remove(mUserId);
-                    mLatch.countDown();
-                }
-            }
-        }
-
-        private void abort() {
-            mAborted = true;
-        }
-
-        private void waitUntilFinish() {
-            if (mAborted) {
-                return;
-            }
-
-            try {
-                mLatch.await();
-            } catch (InterruptedException e) {
-                throw new RuntimeException(e);
-            }
-        }
-
-        private void loadLaunchParams() {
-            final List<File> filesToDelete = new ArrayList<>();
-            final File launchParamsFolder = getLaunchParamFolder(mUserId);
-            if (!launchParamsFolder.isDirectory()) {
-                Slog.i(TAG, "Didn't find launch param folder for user " + mUserId);
-                return;
-            }
-
-            final Set<String> packages = new ArraySet<>(mPackageList.getPackageNames());
-
-            final File[] paramsFiles = launchParamsFolder.listFiles();
-            final ArrayMap<ComponentName, PersistableLaunchParams> map =
-                    new ArrayMap<>(paramsFiles.length);
-
-            for (File paramsFile : paramsFiles) {
-                if (!paramsFile.isFile()) {
-                    Slog.w(TAG, paramsFile.getAbsolutePath() + " is not a file.");
-                    continue;
-                }
-                if (!paramsFile.getName().endsWith(LAUNCH_PARAMS_FILE_SUFFIX)) {
-                    Slog.w(TAG, "Unexpected params file name: " + paramsFile.getName());
-                    filesToDelete.add(paramsFile);
-                    continue;
-                }
-                String paramsFileName = paramsFile.getName();
-                // Migrate all records from old separator to new separator.
-                final int oldSeparatorIndex =
-                        paramsFileName.indexOf(OLD_ESCAPED_COMPONENT_SEPARATOR);
-                if (oldSeparatorIndex != -1) {
-                    if (paramsFileName.indexOf(
-                            OLD_ESCAPED_COMPONENT_SEPARATOR, oldSeparatorIndex + 1) != -1) {
-                        // Rare case. We have more than one old escaped component separator probably
-                        // because this app uses underscore in their package name. We can't
-                        // distinguish which one is the real separator so let's skip it.
-                        filesToDelete.add(paramsFile);
-                        continue;
-                    }
-                    paramsFileName = paramsFileName.replace(
-                            OLD_ESCAPED_COMPONENT_SEPARATOR, ESCAPED_COMPONENT_SEPARATOR);
-                    final File newFile = new File(launchParamsFolder, paramsFileName);
-                    if (paramsFile.renameTo(newFile)) {
-                        paramsFile = newFile;
-                    } else {
-                        // Rare case. For some reason we can't rename the file. Let's drop this
-                        // record instead.
-                        filesToDelete.add(paramsFile);
-                        continue;
-                    }
-                }
-                final String componentNameString = paramsFileName.substring(
-                                0 /* beginIndex */,
-                                paramsFileName.length() - LAUNCH_PARAMS_FILE_SUFFIX.length())
-                        .replace(ESCAPED_COMPONENT_SEPARATOR, ORIGINAL_COMPONENT_SEPARATOR);
-                final ComponentName name = ComponentName.unflattenFromString(
-                        componentNameString);
-                if (name == null) {
-                    Slog.w(TAG, "Unexpected file name: " + paramsFileName);
-                    filesToDelete.add(paramsFile);
-                    continue;
-                }
-
-                if (!packages.contains(name.getPackageName())) {
-                    // Rare case. PersisterQueue doesn't have a chance to remove files for removed
-                    // packages last time.
-                    filesToDelete.add(paramsFile);
-                    continue;
-                }
-
-                try (InputStream in = new FileInputStream(paramsFile)) {
-                    final PersistableLaunchParams params = new PersistableLaunchParams();
-                    final TypedXmlPullParser parser = Xml.resolvePullParser(in);
-                    int event;
-                    while ((event = parser.next()) != XmlPullParser.END_DOCUMENT
-                            && event != XmlPullParser.END_TAG) {
-                        if (event != XmlPullParser.START_TAG) {
-                            continue;
-                        }
-
-                        final String tagName = parser.getName();
-                        if (!TAG_LAUNCH_PARAMS.equals(tagName)) {
-                            Slog.w(TAG, "Unexpected tag name: " + tagName);
-                            continue;
-                        }
-
-                        params.restore(paramsFile, parser);
-                    }
-
-                    map.put(name, params);
-                    addComponentNameToLaunchParamAffinityMapIfNotNull(
-                            name, params.mWindowLayoutAffinity);
-                } catch (Exception e) {
-                    Slog.w(TAG, "Failed to restore launch params for " + name, e);
-                    filesToDelete.add(paramsFile);
-                }
-            }
-
-            synchronized (mSupervisor.mService.getGlobalLock()) {
-                if (!mAborted) {
-                    mLaunchParamsMap.put(mUserId, map);
-                }
-            }
-
-            if (!filesToDelete.isEmpty()) {
-                mPersisterQueue.addItem(new CleanUpComponentQueueItem(filesToDelete), true);
-            }
-        }
-    }
-
     private class LaunchParamsWriteQueueItem
             implements PersisterQueue.WriteQueueItem<LaunchParamsWriteQueueItem> {
         private final int mUserId;
@@ -543,8 +466,7 @@
         }
     }
 
-    private static class CleanUpComponentQueueItem
-            implements PersisterQueue.WriteQueueItem<CleanUpComponentQueueItem> {
+    private class CleanUpComponentQueueItem implements PersisterQueue.WriteQueueItem {
         private final List<File> mComponentFiles;
 
         private CleanUpComponentQueueItem(List<File> componentFiles) {
@@ -561,7 +483,7 @@
         }
     }
 
-    private static class PersistableLaunchParams {
+    private class PersistableLaunchParams {
         private static final String ATTR_WINDOWING_MODE = "windowing_mode";
         private static final String ATTR_DISPLAY_UNIQUE_ID = "display_unique_id";
         private static final String ATTR_BOUNDS = "bounds";
diff --git a/services/core/java/com/android/server/wm/PersisterQueue.java b/services/core/java/com/android/server/wm/PersisterQueue.java
index f66069c..9dc3d6a 100644
--- a/services/core/java/com/android/server/wm/PersisterQueue.java
+++ b/services/core/java/com/android/server/wm/PersisterQueue.java
@@ -49,16 +49,14 @@
     /** Special value for mWriteTime to mean don't wait, just write */
     private static final long FLUSH_QUEUE = -1;
 
-    /**
-     * A {@link QueueItem} that doesn't do anything. Used to trigger
-     * {@link Listener#onPreProcessItem}.
-     */
-    static final QueueItem EMPTY_ITEM = () -> { };
+    /** An {@link WriteQueueItem} that doesn't do anything. Used to trigger {@link
+     * Listener#onPreProcessItem}. */
+    static final WriteQueueItem EMPTY_ITEM = () -> { };
 
     private final long mInterWriteDelayMs;
     private final long mPreTaskDelayMs;
     private final LazyTaskWriterThread mLazyTaskWriterThread;
-    private final ArrayList<QueueItem> mQueue = new ArrayList<>();
+    private final ArrayList<WriteQueueItem> mWriteQueue = new ArrayList<>();
 
     private final ArrayList<Listener> mListeners = new ArrayList<>();
 
@@ -107,10 +105,10 @@
         mLazyTaskWriterThread.join();
     }
 
-    synchronized void addItem(QueueItem item, boolean flush) {
-        mQueue.add(item);
+    synchronized void addItem(WriteQueueItem item, boolean flush) {
+        mWriteQueue.add(item);
 
-        if (flush || mQueue.size() > MAX_WRITE_QUEUE_LENGTH) {
+        if (flush || mWriteQueue.size() > MAX_WRITE_QUEUE_LENGTH) {
             mNextWriteTime = FLUSH_QUEUE;
         } else if (mNextWriteTime == 0) {
             mNextWriteTime = SystemClock.uptimeMillis() + mPreTaskDelayMs;
@@ -118,12 +116,11 @@
         notify();
     }
 
-    synchronized <T extends WriteQueueItem<T>> T findLastItem(Predicate<T> predicate,
-            Class<T> clazz) {
-        for (int i = mQueue.size() - 1; i >= 0; --i) {
-            QueueItem queueItem = mQueue.get(i);
-            if (clazz.isInstance(queueItem)) {
-                T item = clazz.cast(queueItem);
+    synchronized <T extends WriteQueueItem> T findLastItem(Predicate<T> predicate, Class<T> clazz) {
+        for (int i = mWriteQueue.size() - 1; i >= 0; --i) {
+            WriteQueueItem writeQueueItem = mWriteQueue.get(i);
+            if (clazz.isInstance(writeQueueItem)) {
+                T item = clazz.cast(writeQueueItem);
                 if (predicate.test(item)) {
                     return item;
                 }
@@ -137,7 +134,7 @@
      * Updates the last item found in the queue that matches the given item, or adds it to the end
      * of the queue if no such item is found.
      */
-    synchronized <T extends WriteQueueItem<T>> void updateLastOrAddItem(T item, boolean flush) {
+    synchronized <T extends WriteQueueItem> void updateLastOrAddItem(T item, boolean flush) {
         final T itemToUpdate = findLastItem(item::matches, (Class<T>) item.getClass());
         if (itemToUpdate == null) {
             addItem(item, flush);
@@ -151,15 +148,15 @@
     /**
      * Removes all items with which given predicate returns {@code true}.
      */
-    synchronized <T extends QueueItem> void removeItems(Predicate<T> predicate,
+    synchronized <T extends WriteQueueItem> void removeItems(Predicate<T> predicate,
             Class<T> clazz) {
-        for (int i = mQueue.size() - 1; i >= 0; --i) {
-            QueueItem queueItem = mQueue.get(i);
-            if (clazz.isInstance(queueItem)) {
-                T item = clazz.cast(queueItem);
+        for (int i = mWriteQueue.size() - 1; i >= 0; --i) {
+            WriteQueueItem writeQueueItem = mWriteQueue.get(i);
+            if (clazz.isInstance(writeQueueItem)) {
+                T item = clazz.cast(writeQueueItem);
                 if (predicate.test(item)) {
                     if (DEBUG) Slog.d(TAG, "Removing " + item + " from write queue.");
-                    mQueue.remove(i);
+                    mWriteQueue.remove(i);
                 }
             }
         }
@@ -204,7 +201,7 @@
         // See https://b.corp.google.com/issues/64438652#comment7
 
         // If mNextWriteTime, then don't delay between each call to saveToXml().
-        final QueueItem item;
+        final WriteQueueItem item;
         synchronized (this) {
             if (mNextWriteTime != FLUSH_QUEUE) {
                 // The next write we don't have to wait so long.
@@ -215,7 +212,7 @@
                 }
             }
 
-            while (mQueue.isEmpty()) {
+            while (mWriteQueue.isEmpty()) {
                 if (mNextWriteTime != 0) {
                     mNextWriteTime = 0; // idle.
                     notify(); // May need to wake up flush().
@@ -227,18 +224,17 @@
                 }
                 if (DEBUG) Slog.d(TAG, "LazyTaskWriter: waiting indefinitely.");
                 wait();
-                // Invariant: mNextWriteTime is either FLUSH_QUEUE or PRE_TASK_DELAY_MS
+                // Invariant: mNextWriteTime is either FLUSH_QUEUE or PRE_WRITE_DELAY_MS
                 // from now.
             }
-            item = mQueue.remove(0);
+            item = mWriteQueue.remove(0);
 
-            final boolean isWriteItem = item instanceof WriteQueueItem<?>;
             long now = SystemClock.uptimeMillis();
             if (DEBUG) {
                 Slog.d(TAG, "LazyTaskWriter: now=" + now + " mNextWriteTime=" + mNextWriteTime
-                        + " mWriteQueue.size=" + mQueue.size() + " isWriteItem=" + isWriteItem);
+                        + " mWriteQueue.size=" + mWriteQueue.size());
             }
-            while (now < mNextWriteTime && isWriteItem) {
+            while (now < mNextWriteTime) {
                 if (DEBUG) {
                     Slog.d(TAG, "LazyTaskWriter: waiting " + (mNextWriteTime - now));
                 }
@@ -252,18 +248,9 @@
         item.process();
     }
 
-    /**
-     * An item the {@link PersisterQueue} processes. Used for loading tasks. Subclasses of this, but
-     * not {@link WriteQueueItem}, aren't subject to waiting.
-     */
-    interface QueueItem {
+    interface WriteQueueItem<T extends WriteQueueItem<T>> {
         void process();
-    }
 
-    /**
-     * A write item the {@link PersisterQueue} processes. Used for persisting tasks.
-     */
-    interface WriteQueueItem<T extends WriteQueueItem<T>> extends QueueItem {
         default void updateFrom(T item) {}
 
         default boolean matches(T item) {
@@ -301,7 +288,7 @@
                 while (true) {
                     final boolean probablyDone;
                     synchronized (PersisterQueue.this) {
-                        probablyDone = mQueue.isEmpty();
+                        probablyDone = mWriteQueue.isEmpty();
                     }
 
                     for (int i = mListeners.size() - 1; i >= 0; --i) {
diff --git a/services/tests/appfunctions/src/android/app/appfunctions/AppFunctionRuntimeMetadataTest.kt b/services/tests/appfunctions/src/android/app/appfunctions/AppFunctionRuntimeMetadataTest.kt
index dbbb2fe..da3e94f 100644
--- a/services/tests/appfunctions/src/android/app/appfunctions/AppFunctionRuntimeMetadataTest.kt
+++ b/services/tests/appfunctions/src/android/app/appfunctions/AppFunctionRuntimeMetadataTest.kt
@@ -112,4 +112,28 @@
         assertThat(runtimeMetadata.appFunctionStaticMetadataQualifiedId)
             .isEqualTo("android\$apps-db/app_functions#com.pkg/funcId")
     }
+
+    @Test
+    fun setEnabled_true() {
+        val runtimeMetadata =
+            AppFunctionRuntimeMetadata.Builder("com.pkg", "funcId").setEnabled(true).build()
+
+        assertThat(runtimeMetadata.enabled).isTrue()
+    }
+
+    @Test
+    fun setEnabled_false() {
+        val runtimeMetadata =
+            AppFunctionRuntimeMetadata.Builder("com.pkg", "funcId").setEnabled(false).build()
+
+        assertThat(runtimeMetadata.enabled).isFalse()
+    }
+
+    @Test
+    fun setEnabled_null() {
+        val runtimeMetadata =
+            AppFunctionRuntimeMetadata.Builder("com.pkg", "funcId").setEnabled(null).build()
+
+        assertThat(runtimeMetadata.enabled).isNull()
+    }
 }
diff --git a/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt b/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
index bc64e15..c05c381 100644
--- a/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
+++ b/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
@@ -36,6 +36,7 @@
 import com.android.internal.infra.AndroidFuture
 import com.android.server.appfunctions.FutureAppSearchSession.FutureSearchResults
 import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors
 import java.util.concurrent.atomic.AtomicBoolean
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -45,6 +46,7 @@
 class MetadataSyncAdapterTest {
     private val context = InstrumentationRegistry.getInstrumentation().targetContext
     private val appSearchManager = context.getSystemService(AppSearchManager::class.java)
+    private val testExecutor = MoreExecutors.directExecutor()
     private val packageManager = context.packageManager
 
     @Test
@@ -136,7 +138,8 @@
             PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
         runtimeSearchSession.put(putDocumentsRequest).get()
         staticSearchSession.put(putDocumentsRequest).get()
-        val metadataSyncAdapter = MetadataSyncAdapter(packageManager, appSearchManager)
+        val metadataSyncAdapter =
+            MetadataSyncAdapter(testExecutor, packageManager, appSearchManager)
 
         val submitSyncRequest =
             metadataSyncAdapter.trySyncAppFunctionMetadataBlocking(
@@ -177,7 +180,8 @@
         val putDocumentsRequest: PutDocumentsRequest =
             PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
         staticSearchSession.put(putDocumentsRequest).get()
-        val metadataSyncAdapter = MetadataSyncAdapter(packageManager, appSearchManager)
+        val metadataSyncAdapter =
+            MetadataSyncAdapter(testExecutor, packageManager, appSearchManager)
 
         val submitSyncRequest =
             metadataSyncAdapter.trySyncAppFunctionMetadataBlocking(
@@ -232,7 +236,8 @@
         val putDocumentsRequest: PutDocumentsRequest =
             PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
         runtimeSearchSession.put(putDocumentsRequest).get()
-        val metadataSyncAdapter = MetadataSyncAdapter(packageManager, appSearchManager)
+        val metadataSyncAdapter =
+            MetadataSyncAdapter(testExecutor, packageManager, appSearchManager)
 
         val submitSyncRequest =
             metadataSyncAdapter.trySyncAppFunctionMetadataBlocking(
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeClamperTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeClamperTest.java
deleted file mode 100644
index 306b4f8..0000000
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeClamperTest.java
+++ /dev/null
@@ -1,126 +0,0 @@
-/*
- * Copyright (C) 2023 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.server.display.brightness.clamper;
-
-import static com.android.server.display.brightness.clamper.BrightnessWearBedtimeModeClamper.BEDTIME_MODE_OFF;
-import static com.android.server.display.brightness.clamper.BrightnessWearBedtimeModeClamper.BEDTIME_MODE_ON;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.verify;
-
-import android.content.ContentResolver;
-import android.database.ContentObserver;
-import android.provider.Settings;
-import android.testing.TestableContext;
-
-import androidx.annotation.NonNull;
-import androidx.test.platform.app.InstrumentationRegistry;
-
-import com.android.internal.display.BrightnessSynchronizer;
-import com.android.server.testutils.TestHandler;
-
-import org.junit.Before;
-import org.junit.Rule;
-import org.junit.Test;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
-
-public class BrightnessWearBedtimeModeClamperTest {
-
-    private static final float BRIGHTNESS_CAP = 0.3f;
-
-    @Mock
-    private BrightnessClamperController.ClamperChangeListener mMockClamperChangeListener;
-
-    @Rule
-    public final TestableContext mContext = new TestableContext(
-            InstrumentationRegistry.getInstrumentation().getContext());
-
-    private final TestHandler mTestHandler = new TestHandler(null);
-    private final TestInjector mInjector = new TestInjector();
-    private BrightnessWearBedtimeModeClamper mClamper;
-
-    @Before
-    public void setUp() {
-        MockitoAnnotations.initMocks(this);
-        mClamper = new BrightnessWearBedtimeModeClamper(mInjector, mTestHandler, mContext,
-                mMockClamperChangeListener, () -> BRIGHTNESS_CAP);
-        mTestHandler.flush();
-    }
-
-    @Test
-    public void testBrightnessCap() {
-        assertEquals(BRIGHTNESS_CAP, mClamper.getBrightnessCap(), BrightnessSynchronizer.EPSILON);
-    }
-
-    @Test
-    public void testBedtimeModeOn() {
-        setBedtimeModeEnabled(true);
-        assertTrue(mClamper.isActive());
-        verify(mMockClamperChangeListener).onChanged();
-    }
-
-    @Test
-    public void testBedtimeModeOff() {
-        setBedtimeModeEnabled(false);
-        assertFalse(mClamper.isActive());
-        verify(mMockClamperChangeListener).onChanged();
-    }
-
-    @Test
-    public void testType() {
-        assertEquals(BrightnessClamper.Type.WEAR_BEDTIME_MODE, mClamper.getType());
-    }
-
-    @Test
-    public void testOnDisplayChanged() {
-        float newBrightnessCap = 0.61f;
-
-        mClamper.onDisplayChanged(() -> newBrightnessCap);
-        mTestHandler.flush();
-
-        assertEquals(newBrightnessCap, mClamper.getBrightnessCap(), BrightnessSynchronizer.EPSILON);
-        verify(mMockClamperChangeListener).onChanged();
-    }
-
-    private void setBedtimeModeEnabled(boolean enabled) {
-        Settings.Global.putInt(mContext.getContentResolver(), Settings.Global.Wearable.BEDTIME_MODE,
-                enabled ? BEDTIME_MODE_ON : BEDTIME_MODE_OFF);
-        mInjector.notifyBedtimeModeChanged();
-        mTestHandler.flush();
-    }
-
-    private static class TestInjector extends BrightnessWearBedtimeModeClamper.Injector {
-
-        private ContentObserver mObserver;
-
-        @Override
-        void registerBedtimeModeObserver(@NonNull ContentResolver cr,
-                @NonNull ContentObserver observer) {
-            mObserver = observer;
-        }
-
-        private void notifyBedtimeModeChanged() {
-            if (mObserver != null) {
-                mObserver.dispatchChange(/* selfChange= */ false,
-                        Settings.Global.getUriFor(Settings.Global.Wearable.BEDTIME_MODE));
-            }
-        }
-    }
-}
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeModifierTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeModifierTest.java
new file mode 100644
index 0000000..8271a0f
--- /dev/null
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessWearBedtimeModeModifierTest.java
@@ -0,0 +1,180 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.brightness.clamper;
+
+import static com.android.server.display.brightness.clamper.BrightnessWearBedtimeModeModifier.BEDTIME_MODE_OFF;
+import static com.android.server.display.brightness.clamper.BrightnessWearBedtimeModeModifier.BEDTIME_MODE_ON;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.content.ContentResolver;
+import android.database.ContentObserver;
+import android.hardware.display.BrightnessInfo;
+import android.hardware.display.DisplayManagerInternal;
+import android.os.IBinder;
+import android.os.PowerManager;
+import android.provider.Settings;
+import android.testing.TestableContext;
+
+import androidx.annotation.NonNull;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import com.android.internal.display.BrightnessSynchronizer;
+import com.android.server.display.DisplayBrightnessState;
+import com.android.server.display.DisplayDeviceConfig;
+import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.clamper.BrightnessClamperController.ModifiersAggregatedState;
+import com.android.server.testutils.TestHandler;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+public class BrightnessWearBedtimeModeModifierTest {
+    private static final int NO_MODIFIER = 0;
+    private static final float BRIGHTNESS_CAP = 0.3f;
+    private static final String DISPLAY_ID = "displayId";
+
+    @Mock
+    private BrightnessClamperController.ClamperChangeListener mMockClamperChangeListener;
+    @Mock
+    private DisplayManagerInternal.DisplayPowerRequest mMockRequest;
+    @Mock
+    private DisplayDeviceConfig mMockDisplayDeviceConfig;
+    @Mock
+    private IBinder mMockBinder;
+
+    @Rule
+    public final TestableContext mContext = new TestableContext(
+            InstrumentationRegistry.getInstrumentation().getContext());
+
+    private final TestHandler mTestHandler = new TestHandler(null);
+    private final TestInjector mInjector = new TestInjector();
+    private BrightnessWearBedtimeModeModifier mModifier;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mModifier = new BrightnessWearBedtimeModeModifier(mInjector, mTestHandler, mContext,
+                mMockClamperChangeListener, () -> BRIGHTNESS_CAP);
+        mTestHandler.flush();
+    }
+
+    @Test
+    public void testBedtimeModeOff() {
+        setBedtimeModeEnabled(false);
+        assertModifierState(
+                0.5f, true,
+                PowerManager.BRIGHTNESS_MAX, 0.5f,
+                false, true);
+        verify(mMockClamperChangeListener).onChanged();
+    }
+
+    @Test
+    public void testBedtimeModeOn() {
+        setBedtimeModeEnabled(true);
+        assertModifierState(
+                0.5f, true,
+                BRIGHTNESS_CAP, BRIGHTNESS_CAP,
+                true, false);
+        verify(mMockClamperChangeListener).onChanged();
+    }
+
+    @Test
+    public void testOnDisplayChanged() {
+        setBedtimeModeEnabled(true);
+        clearInvocations(mMockClamperChangeListener);
+        float newBrightnessCap = 0.61f;
+        onDisplayChange(newBrightnessCap);
+        mTestHandler.flush();
+
+        assertModifierState(
+                0.5f, true,
+                newBrightnessCap, 0.5f,
+                true, false);
+        verify(mMockClamperChangeListener).onChanged();
+    }
+
+    private void setBedtimeModeEnabled(boolean enabled) {
+        Settings.Global.putInt(mContext.getContentResolver(), Settings.Global.Wearable.BEDTIME_MODE,
+                enabled ? BEDTIME_MODE_ON : BEDTIME_MODE_OFF);
+        mInjector.notifyBedtimeModeChanged();
+        mTestHandler.flush();
+    }
+
+    private void onDisplayChange(float brightnessCap) {
+        when(mMockDisplayDeviceConfig.getBrightnessCapForWearBedtimeMode())
+                .thenReturn(brightnessCap);
+        mModifier.onDisplayChanged(ClamperTestUtilsKt.createDisplayDeviceData(
+                mMockDisplayDeviceConfig, mMockBinder, DISPLAY_ID, DisplayDeviceConfig.DEFAULT_ID));
+    }
+
+    private void assertModifierState(
+            float currentBrightness,
+            boolean currentSlowChange,
+            float maxBrightness, float brightness,
+            boolean isActive,
+            boolean isSlowChange) {
+        ModifiersAggregatedState modifierState = new ModifiersAggregatedState();
+        DisplayBrightnessState.Builder stateBuilder = DisplayBrightnessState.builder();
+        stateBuilder.setBrightness(currentBrightness);
+        stateBuilder.setIsSlowChange(currentSlowChange);
+
+        int maxBrightnessReason = isActive ? BrightnessInfo.BRIGHTNESS_MAX_REASON_WEAR_BEDTIME_MODE
+                : BrightnessInfo.BRIGHTNESS_MAX_REASON_NONE;
+        int modifier = isActive ? BrightnessReason.MODIFIER_THROTTLED : NO_MODIFIER;
+
+        mModifier.applyStateChange(modifierState);
+        assertThat(modifierState.mMaxBrightness).isEqualTo(maxBrightness);
+        assertThat(modifierState.mMaxBrightnessReason).isEqualTo(maxBrightnessReason);
+
+        mModifier.apply(mMockRequest, stateBuilder);
+
+        assertThat(stateBuilder.getMaxBrightness())
+                .isWithin(BrightnessSynchronizer.EPSILON).of(maxBrightness);
+        assertThat(stateBuilder.getBrightness())
+                .isWithin(BrightnessSynchronizer.EPSILON).of(brightness);
+        assertThat(stateBuilder.getBrightnessMaxReason()).isEqualTo(maxBrightnessReason);
+        assertThat(stateBuilder.getBrightnessReason().getModifier()).isEqualTo(modifier);
+        assertThat(stateBuilder.isSlowChange()).isEqualTo(isSlowChange);
+    }
+
+
+    private static class TestInjector extends BrightnessWearBedtimeModeModifier.Injector {
+
+        private ContentObserver mObserver;
+
+        @Override
+        void registerBedtimeModeObserver(@NonNull ContentResolver cr,
+                @NonNull ContentObserver observer) {
+            mObserver = observer;
+        }
+
+        private void notifyBedtimeModeChanged() {
+            if (mObserver != null) {
+                mObserver.dispatchChange(/* selfChange= */ false,
+                        Settings.Global.getUriFor(Settings.Global.Wearable.BEDTIME_MODE));
+            }
+        }
+    }
+}
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/BackgroundUserSoundNotifierTest.java b/services/tests/mockingservicestests/src/com/android/server/pm/BackgroundUserSoundNotifierTest.java
index 3062d51..9ba2724 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/BackgroundUserSoundNotifierTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/BackgroundUserSoundNotifierTest.java
@@ -172,6 +172,7 @@
         mBackgroundUserSoundNotifier.muteAlarmSounds(mSpiedContext);
 
         verify(apc1.getPlayerProxy()).stop();
+        verify(mockAudioPolicy).sendFocusLossAndUpdate(afi);
         verify(apc2.getPlayerProxy(), never()).stop();
     }
 
diff --git a/services/tests/mockingservicestests/src/com/android/server/wallpaper/WallpaperManagerServiceTests.java b/services/tests/mockingservicestests/src/com/android/server/wallpaper/WallpaperManagerServiceTests.java
index 0b762df..9983fb4 100644
--- a/services/tests/mockingservicestests/src/com/android/server/wallpaper/WallpaperManagerServiceTests.java
+++ b/services/tests/mockingservicestests/src/com/android/server/wallpaper/WallpaperManagerServiceTests.java
@@ -32,6 +32,8 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.server.wallpaper.WallpaperUtils.WALLPAPER;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.hamcrest.core.IsNot.not;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
@@ -50,6 +52,7 @@
 
 import android.app.AppGlobals;
 import android.app.AppOpsManager;
+import android.app.Flags;
 import android.app.WallpaperColors;
 import android.app.WallpaperManager;
 import android.content.ComponentName;
@@ -64,7 +67,10 @@
 import android.os.ParcelFileDescriptor;
 import android.os.RemoteException;
 import android.os.SystemClock;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.service.wallpaper.IWallpaperConnection;
 import android.service.wallpaper.IWallpaperEngine;
 import android.service.wallpaper.WallpaperService;
@@ -91,8 +97,10 @@
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
+import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.RuleChain;
 import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -100,6 +108,7 @@
 import org.mockito.quality.Strictness;
 import org.xmlpull.v1.XmlPullParserException;
 
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.FileInputStream;
@@ -125,6 +134,7 @@
     @ClassRule
     public static final TestableContext sContext = new TestableContext(
             InstrumentationRegistry.getInstrumentation().getTargetContext(), null);
+
     private static ComponentName sImageWallpaperComponentName;
     private static ComponentName sDefaultWallpaperComponent;
 
@@ -133,8 +143,11 @@
     @Mock
     private DisplayManager mDisplayManager;
 
+    private final TemporaryFolder mFolder = new TemporaryFolder();
+    private final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
     @Rule
-    public final TemporaryFolder mFolder = new TemporaryFolder();
+    public RuleChain rules = RuleChain.outerRule(mSetFlagsRule).around(mFolder);
+
     private final SparseArray<File> mTempDirs = new SparseArray<>();
     private WallpaperManagerService mService;
     private static IWallpaperConnection.Stub sWallpaperService;
@@ -325,6 +338,7 @@
      * is issued to the wallpaper.
      */
     @Test
+    @Ignore("b/368345733")
     public void testSetCurrentComponent() throws Exception {
         final int testUserId = USER_SYSTEM;
         mService.switchUser(testUserId, null);
@@ -411,26 +425,84 @@
     }
 
     @Test
-    public void testXmlSerializationRoundtrip() {
-        WallpaperData systemWallpaperData = mService.getCurrentWallpaperData(FLAG_SYSTEM, 0);
+    @EnableFlags(Flags.FLAG_REMOVE_NEXT_WALLPAPER_COMPONENT)
+    public void testSaveLoadSettings() {
+        WallpaperData expectedData = mService.getCurrentWallpaperData(FLAG_SYSTEM, 0);
+        expectedData.setComponent(sDefaultWallpaperComponent);
+        expectedData.primaryColors = new WallpaperColors(Color.valueOf(Color.RED),
+                Color.valueOf(Color.BLUE), null);
+        expectedData.mWallpaperDimAmount = 0.5f;
+        expectedData.mUidToDimAmount.put(0, 0.5f);
+        expectedData.mUidToDimAmount.put(1, 0.4f);
+
+        ByteArrayOutputStream ostream = new ByteArrayOutputStream();
         try {
             TypedXmlSerializer serializer = Xml.newBinarySerializer();
-            serializer.setOutput(new ByteArrayOutputStream(), StandardCharsets.UTF_8.name());
-            serializer.startDocument(StandardCharsets.UTF_8.name(), true);
-            mService.mWallpaperDataParser.writeWallpaperAttributes(
-                    serializer, "wp", systemWallpaperData);
+            serializer.setOutput(ostream, StandardCharsets.UTF_8.name());
+            mService.mWallpaperDataParser.saveSettingsToSerializer(serializer, expectedData, null);
+            ostream.close();
+        } catch (IOException e) {
+            fail("exception occurred while writing system wallpaper attributes");
+        }
+
+        WallpaperData actualData = new WallpaperData(0, FLAG_SYSTEM);
+        try {
+            ByteArrayInputStream istream = new ByteArrayInputStream(ostream.toByteArray());
+            TypedXmlPullParser parser = Xml.newBinaryPullParser();
+            parser.setInput(istream, StandardCharsets.UTF_8.name());
+            mService.mWallpaperDataParser.loadSettingsFromSerializer(parser,
+                    actualData, /* userId= */0, /* loadSystem= */ true, /* loadLock= */
+                    false, /* keepDimensionHints= */ true,
+                    new WallpaperDisplayHelper.DisplayData(0));
+        } catch (IOException | XmlPullParserException e) {
+            fail("exception occurred while parsing wallpaper");
+        }
+
+        assertThat(actualData.getComponent()).isEqualTo(expectedData.getComponent());
+        assertThat(actualData.primaryColors).isEqualTo(expectedData.primaryColors);
+        assertThat(actualData.mWallpaperDimAmount).isEqualTo(expectedData.mWallpaperDimAmount);
+        assertThat(actualData.mUidToDimAmount).isNotNull();
+        assertThat(actualData.mUidToDimAmount.size()).isEqualTo(
+                expectedData.mUidToDimAmount.size());
+        for (int i = 0; i < actualData.mUidToDimAmount.size(); i++) {
+            int key = actualData.mUidToDimAmount.keyAt(0);
+            assertThat(actualData.mUidToDimAmount.get(key)).isEqualTo(
+                    expectedData.mUidToDimAmount.get(key));
+        }
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_REMOVE_NEXT_WALLPAPER_COMPONENT)
+    public void testSaveLoadSettings_legacyNextComponent() {
+        WallpaperData systemWallpaperData = mService.getCurrentWallpaperData(FLAG_SYSTEM, 0);
+        systemWallpaperData.setComponent(sDefaultWallpaperComponent);
+        ByteArrayOutputStream ostream = new ByteArrayOutputStream();
+        try {
+            TypedXmlSerializer serializer = Xml.newBinarySerializer();
+            serializer.setOutput(ostream, StandardCharsets.UTF_8.name());
+            mService.mWallpaperDataParser.saveSettingsToSerializer(serializer, systemWallpaperData,
+                    null);
+            ostream.close();
         } catch (IOException e) {
             fail("exception occurred while writing system wallpaper attributes");
         }
 
         WallpaperData shouldMatchSystem = new WallpaperData(0, FLAG_SYSTEM);
         try {
+            ByteArrayInputStream istream = new ByteArrayInputStream(ostream.toByteArray());
             TypedXmlPullParser parser = Xml.newBinaryPullParser();
-            mService.mWallpaperDataParser.parseWallpaperAttributes(parser, shouldMatchSystem, true);
-        } catch (XmlPullParserException e) {
+            parser.setInput(istream, StandardCharsets.UTF_8.name());
+            mService.mWallpaperDataParser.loadSettingsFromSerializer(parser,
+                    shouldMatchSystem, /* userId= */0, /* loadSystem= */ true, /* loadLock= */
+                    false, /* keepDimensionHints= */ true,
+                    new WallpaperDisplayHelper.DisplayData(0));
+        } catch (IOException | XmlPullParserException e) {
             fail("exception occurred while parsing wallpaper");
         }
-        assertEquals(systemWallpaperData.primaryColors, shouldMatchSystem.primaryColors);
+
+        assertThat(shouldMatchSystem.nextWallpaperComponent).isEqualTo(
+                systemWallpaperData.getComponent());
+        assertThat(shouldMatchSystem.primaryColors).isEqualTo(systemWallpaperData.primaryColors);
     }
 
     @Test
diff --git a/services/tests/servicestests/src/com/android/server/audio/MediaFocusControlTest.java b/services/tests/servicestests/src/com/android/server/audio/MediaFocusControlTest.java
new file mode 100644
index 0000000..34878c8
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/audio/MediaFocusControlTest.java
@@ -0,0 +1,128 @@
+/*
+ * Copyright 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.audio;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.media.AudioAttributes;
+import android.media.AudioFocusInfo;
+import android.media.AudioManager;
+import android.os.Binder;
+import android.os.IBinder;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.filters.MediumTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.List;
+
+@MediumTest
+@RunWith(AndroidJUnit4.class)
+public class MediaFocusControlTest {
+    private static final String TAG = "MediaFocusControlTest";
+
+    private Context mContext;
+    private MediaFocusControl mMediaFocusControl;
+    private final IBinder mICallBack = new Binder();
+
+
+    private static class NoopPlayerFocusEnforcer implements PlayerFocusEnforcer {
+        public boolean duckPlayers(@NonNull FocusRequester winner, @NonNull FocusRequester loser,
+                boolean forceDuck) {
+            return true;
+        }
+
+        public void restoreVShapedPlayers(@NonNull FocusRequester winner) {
+        }
+
+        public void mutePlayersForCall(int[] usagesToMute) {
+        }
+
+        public void unmutePlayersForCall() {
+        }
+
+        public boolean fadeOutPlayers(@NonNull FocusRequester winner,
+                @NonNull FocusRequester loser) {
+            return true;
+        }
+
+        public void forgetUid(int uid) {
+        }
+
+        public long getFadeOutDurationMillis(@NonNull AudioAttributes aa) {
+            return 100;
+        }
+
+        public long getFadeInDelayForOffendersMillis(@NonNull AudioAttributes aa) {
+            return 100;
+        }
+
+        public boolean shouldEnforceFade() {
+            return false;
+        }
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        mContext = InstrumentationRegistry.getTargetContext();
+        mMediaFocusControl = new MediaFocusControl(mContext, new NoopPlayerFocusEnforcer());
+    }
+
+    private static final AudioAttributes MEDIA_ATTRIBUTES = new AudioAttributes.Builder()
+            .setUsage(AudioAttributes.USAGE_MEDIA).build();
+    private static final AudioAttributes ALARM_ATTRIBUTES = new AudioAttributes.Builder()
+            .setUsage(AudioAttributes.USAGE_ALARM).build();
+    private static final int MEDIA_UID = 10300;
+    private static final int ALARM_UID = 10301;
+
+    /**
+     * Test {@link MediaFocusControl#sendFocusLossAndUpdate(AudioFocusInfo)}
+     */
+    @Test
+    public void testSendFocusLossAndUpdate() throws Exception {
+        // simulate a media app requesting focus, followed by an alarm
+        mMediaFocusControl.requestAudioFocus(MEDIA_ATTRIBUTES, AudioManager.AUDIOFOCUS_GAIN,
+                mICallBack, null /*focusDispatcher*/, "clientMedia", "packMedia",
+                AudioManager.AUDIOFOCUS_FLAG_TEST /*flags*/, 35 /*sdk*/, false/*forceDuck*/,
+                MEDIA_UID, true /*permissionOverridesCheck*/);
+        final AudioFocusInfo alarm = new AudioFocusInfo(ALARM_ATTRIBUTES, ALARM_UID,
+                "clientAlarm", "packAlarm",
+                AudioManager.AUDIOFOCUS_GAIN_TRANSIENT_MAY_DUCK, 0/*lossReceived*/,
+                AudioManager.AUDIOFOCUS_FLAG_TEST /*flags*/, 35 /*sdk*/);
+        mMediaFocusControl.requestAudioFocus(alarm.getAttributes(), alarm.getGainRequest(),
+                mICallBack, null /*focusDispatcher*/, alarm.getClientId(), alarm.getPackageName(),
+                alarm.getFlags(), alarm.getSdkTarget(), false/*forceDuck*/,
+                alarm.getClientUid(), true /*permissionOverridesCheck*/);
+        // verify stack is in expected state
+        List<AudioFocusInfo> stack = mMediaFocusControl.getFocusStack();
+        Assert.assertEquals("focus stack should have 2 entries", 2, stack.size());
+        Assert.assertEquals("focus loser should have received LOSS_TRANSIENT_CAN_DUCK",
+                AudioManager.AUDIOFOCUS_LOSS_TRANSIENT_CAN_DUCK, stack.get(0).getLossReceived());
+
+        // make alarm app lose focus and check stack
+        mMediaFocusControl.sendFocusLossAndUpdate(alarm);
+        stack = mMediaFocusControl.getFocusStack();
+        Assert.assertEquals("focus stack should have 1 entry after sendFocusLossAndUpdate",
+                1, stack.size());
+        Assert.assertEquals("new top of stack should be media app",
+                MEDIA_UID, stack.get(0).getClientUid());
+    }
+}
diff --git a/services/tests/wmtests/src/com/android/server/policy/SingleKeyGestureTests.java b/services/tests/wmtests/src/com/android/server/policy/SingleKeyGestureTests.java
index 7ea5010..ff8b6d3 100644
--- a/services/tests/wmtests/src/com/android/server/policy/SingleKeyGestureTests.java
+++ b/services/tests/wmtests/src/com/android/server/policy/SingleKeyGestureTests.java
@@ -141,7 +141,8 @@
                     }
 
                     @Override
-                    void onKeyUp(long eventTime, int multiPressCount, int displayId) {
+                    void onKeyUp(long eventTime, int multiPressCount, int displayId, int deviceId,
+                            int metaState) {
                         mKeyUpQueue.add(new KeyUpData(KEYCODE_POWER, multiPressCount));
                     }
                 });
@@ -177,7 +178,8 @@
                     }
 
                     @Override
-                    void onKeyUp(long eventTime, int multiPressCount, int displayId) {
+                    void onKeyUp(long eventTime, int multiPressCount, int displayId, int deviceId,
+                            int metaState) {
                         mKeyUpQueue.add(new KeyUpData(KEYCODE_BACK, multiPressCount));
                     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java b/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java
index 62d3949..1be61c3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/LaunchParamsPersisterTests.java
@@ -293,7 +293,6 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         target.getLaunchParams(mTestTask, null, mResult);
 
@@ -312,7 +311,6 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         mTaskWithDifferentComponent.mWindowLayoutAffinity = TEST_WINDOW_LAYOUT_AFFINITY;
         target.getLaunchParams(mTaskWithDifferentComponent, null, mResult);
@@ -341,7 +339,6 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         target.getLaunchParams(mTaskWithDifferentComponent, null, mResult);
 
@@ -411,7 +408,6 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         target.getLaunchParams(mTestTask, null, mResult);
 
@@ -429,7 +425,6 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         target.getLaunchParams(mTestTask, null, mResult);
 
@@ -458,7 +453,6 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         target.getLaunchParams(mTestTask, null, mResult);
 
@@ -476,7 +470,6 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         target.getLaunchParams(mTestTask, null, mResult);
 
@@ -495,52 +488,12 @@
                 mUserFolderGetter);
         target.onSystemReady();
         target.onUnlockUser(TEST_USER_ID);
-        mPersisterQueue.flush();
 
         target.getLaunchParams(mTestTask, null, mResult);
 
         assertTrue("Result should be empty.", mResult.isEmpty());
     }
 
-    @Test
-    public void testAbortsLoadingWhenUserCleansUpBeforeLoadingFinishes() {
-        mTarget.saveTask(mTestTask);
-        mPersisterQueue.flush();
-
-        final LaunchParamsPersister target = new LaunchParamsPersister(mPersisterQueue, mSupervisor,
-                mUserFolderGetter);
-        target.onSystemReady();
-        target.onUnlockUser(TEST_USER_ID);
-        assertEquals(1, mPersisterQueue.mQueue.size());
-        PersisterQueue.QueueItem item = mPersisterQueue.mQueue.get(0);
-
-        target.onCleanupUser(TEST_USER_ID);
-        mPersisterQueue.flush();
-
-        // Explicitly run the loading item to mimic the situation where the item already started.
-        item.process();
-
-        target.getLaunchParams(mTestTask, null, mResult);
-        assertTrue("Result should be empty.", mResult.isEmpty());
-    }
-
-    @Test
-    public void testGetLaunchParamsNotBlockedByAbortedLoading() {
-        mTarget.saveTask(mTestTask);
-        mPersisterQueue.flush();
-
-        final LaunchParamsPersister target = new LaunchParamsPersister(mPersisterQueue, mSupervisor,
-                mUserFolderGetter);
-        target.onSystemReady();
-        target.onUnlockUser(TEST_USER_ID);
-        target.onCleanupUser(TEST_USER_ID);
-
-        // As long as the call in the next line returns, we know it's not waiting for the loading to
-        // finish because we run items synchronously in this test.
-        target.getLaunchParams(mTestTask, null, mResult);
-        assertTrue("Result should be empty.", mResult.isEmpty());
-    }
-
     private static boolean deleteRecursively(File file) {
         boolean result = true;
         if (file.isDirectory()) {
@@ -555,17 +508,17 @@
 
     /**
      * Test double to {@link PersisterQueue}. This is not thread-safe and caller should always use
-     * {@link #flush()} to execute items in it.
+     * {@link #flush()} to execute write items in it.
      */
     static class TestPersisterQueue extends PersisterQueue {
-        private List<QueueItem> mQueue = new ArrayList<>();
+        private List<WriteQueueItem> mWriteQueue = new ArrayList<>();
         private List<Listener> mListeners = new ArrayList<>();
 
         @Override
         void flush() {
-            while (!mQueue.isEmpty()) {
-                final QueueItem item = mQueue.remove(0);
-                final boolean queueEmpty = mQueue.isEmpty();
+            while (!mWriteQueue.isEmpty()) {
+                final WriteQueueItem item = mWriteQueue.remove(0);
+                final boolean queueEmpty = mWriteQueue.isEmpty();
                 for (Listener listener : mListeners) {
                     listener.onPreProcessItem(queueEmpty);
                 }
@@ -584,18 +537,18 @@
         }
 
         @Override
-        synchronized void addItem(QueueItem item, boolean flush) {
-            mQueue.add(item);
+        void addItem(WriteQueueItem item, boolean flush) {
+            mWriteQueue.add(item);
             if (flush) {
                 flush();
             }
         }
 
         @Override
-        synchronized <T extends WriteQueueItem<T>> T findLastItem(Predicate<T> predicate,
+        synchronized <T extends WriteQueueItem> T findLastItem(Predicate<T> predicate,
                 Class<T> clazz) {
-            for (int i = mQueue.size() - 1; i >= 0; --i) {
-                QueueItem writeQueueItem = mQueue.get(i);
+            for (int i = mWriteQueue.size() - 1; i >= 0; --i) {
+                WriteQueueItem writeQueueItem = mWriteQueue.get(i);
                 if (clazz.isInstance(writeQueueItem)) {
                     T item = clazz.cast(writeQueueItem);
                     if (predicate.test(item)) {
@@ -608,14 +561,14 @@
         }
 
         @Override
-        synchronized <T extends QueueItem> void removeItems(Predicate<T> predicate,
+        synchronized <T extends WriteQueueItem> void removeItems(Predicate<T> predicate,
                 Class<T> clazz) {
-            for (int i = mQueue.size() - 1; i >= 0; --i) {
-                QueueItem writeQueueItem = mQueue.get(i);
+            for (int i = mWriteQueue.size() - 1; i >= 0; --i) {
+                WriteQueueItem writeQueueItem = mWriteQueue.get(i);
                 if (clazz.isInstance(writeQueueItem)) {
                     T item = clazz.cast(writeQueueItem);
                     if (predicate.test(item)) {
-                        mQueue.remove(i);
+                        mWriteQueue.remove(i);
                     }
                 }
             }
diff --git a/services/tests/wmtests/src/com/android/server/wm/PersisterQueueTests.java b/services/tests/wmtests/src/com/android/server/wm/PersisterQueueTests.java
index ce0e6f8..3e87f1f 100644
--- a/services/tests/wmtests/src/com/android/server/wm/PersisterQueueTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/PersisterQueueTests.java
@@ -90,26 +90,8 @@
         mFactory.setExpectedProcessedItemNumber(1);
         mListener.setExpectedOnPreProcessItemCallbackTimes(1);
 
-        mTarget.addItem(mFactory.createItem(), false);
-        assertTrue("Target didn't process item enough times.",
-                mFactory.waitForAllExpectedItemsProcessed(TIMEOUT_ALLOWANCE));
-        assertEquals("Target didn't process item.", 1, mFactory.getTotalProcessedItemCount());
-
-        assertTrue("Target didn't call callback enough times.",
-                mListener.waitForAllExpectedCallbackDone(TIMEOUT_ALLOWANCE));
-        // Once before processing this item, once after that.
-        assertEquals(2, mListener.mProbablyDoneResults.size());
-        // The last one must be called with probably done being true.
-        assertTrue("The last probablyDone must be true.", mListener.mProbablyDoneResults.get(1));
-    }
-
-    @Test
-    public void testProcessOneWriteItem() throws Exception {
-        mFactory.setExpectedProcessedItemNumber(1);
-        mListener.setExpectedOnPreProcessItemCallbackTimes(1);
-
         final long dispatchTime = SystemClock.uptimeMillis();
-        mTarget.addItem(mFactory.createWriteItem(), false);
+        mTarget.addItem(mFactory.createItem(), false);
         assertTrue("Target didn't process item enough times.",
                 mFactory.waitForAllExpectedItemsProcessed(PRE_TASK_DELAY_MS + TIMEOUT_ALLOWANCE));
         assertEquals("Target didn't process item.", 1, mFactory.getTotalProcessedItemCount());
@@ -127,12 +109,12 @@
     }
 
     @Test
-    public void testProcessOneWriteItem_Flush() throws Exception {
+    public void testProcessOneItem_Flush() throws Exception {
         mFactory.setExpectedProcessedItemNumber(1);
         mListener.setExpectedOnPreProcessItemCallbackTimes(1);
 
         final long dispatchTime = SystemClock.uptimeMillis();
-        mTarget.addItem(mFactory.createWriteItem(), true);
+        mTarget.addItem(mFactory.createItem(), true);
         assertTrue("Target didn't process item enough times.",
                 mFactory.waitForAllExpectedItemsProcessed(TIMEOUT_ALLOWANCE));
         assertEquals("Target didn't process item.", 1, mFactory.getTotalProcessedItemCount());
@@ -156,8 +138,8 @@
         mListener.setExpectedOnPreProcessItemCallbackTimes(2);
 
         final long dispatchTime = SystemClock.uptimeMillis();
-        mTarget.addItem(mFactory.createWriteItem(), false);
-        mTarget.addItem(mFactory.createWriteItem(), false);
+        mTarget.addItem(mFactory.createItem(), false);
+        mTarget.addItem(mFactory.createItem(), false);
         assertTrue("Target didn't call callback enough times.",
                 mFactory.waitForAllExpectedItemsProcessed(PRE_TASK_DELAY_MS + INTER_WRITE_DELAY_MS
                         + TIMEOUT_ALLOWANCE));
@@ -183,7 +165,7 @@
         mFactory.setExpectedProcessedItemNumber(1);
         mListener.setExpectedOnPreProcessItemCallbackTimes(1);
         long dispatchTime = SystemClock.uptimeMillis();
-        mTarget.addItem(mFactory.createWriteItem(), false);
+        mTarget.addItem(mFactory.createItem(), false);
         assertTrue("Target didn't process item enough times.",
                 mFactory.waitForAllExpectedItemsProcessed(PRE_TASK_DELAY_MS + TIMEOUT_ALLOWANCE));
         long processDuration = SystemClock.uptimeMillis() - dispatchTime;
@@ -202,7 +184,7 @@
         // Synchronize on the instance to make sure we schedule the item after it starts to wait for
         // task indefinitely.
         synchronized (mTarget) {
-            mTarget.addItem(mFactory.createWriteItem(), false);
+            mTarget.addItem(mFactory.createItem(), false);
         }
         assertTrue("Target didn't process item enough times.",
                 mFactory.waitForAllExpectedItemsProcessed(PRE_TASK_DELAY_MS + TIMEOUT_ALLOWANCE));
@@ -224,9 +206,9 @@
     @Test
     public void testFindLastItemNotReturnDifferentType() {
         synchronized (mTarget) {
-            mTarget.addItem(mFactory.createWriteItem(), false);
-            assertNull(mTarget.findLastItem(TestWriteItem::shouldKeepOnFilter,
-                    FilterableTestWriteItem.class));
+            mTarget.addItem(mFactory.createItem(), false);
+            assertNull(mTarget.findLastItem(TestItem::shouldKeepOnFilter,
+                    FilterableTestItem.class));
         }
     }
 
@@ -234,18 +216,18 @@
     public void testFindLastItemNotReturnMismatchItem() {
         synchronized (mTarget) {
             mTarget.addItem(mFactory.createFilterableItem(false), false);
-            assertNull(mTarget.findLastItem(TestWriteItem::shouldKeepOnFilter,
-                    FilterableTestWriteItem.class));
+            assertNull(mTarget.findLastItem(TestItem::shouldKeepOnFilter,
+                    FilterableTestItem.class));
         }
     }
 
     @Test
     public void testFindLastItemReturnMatchedItem() {
         synchronized (mTarget) {
-            final FilterableTestWriteItem item = mFactory.createFilterableItem(true);
+            final FilterableTestItem item = mFactory.createFilterableItem(true);
             mTarget.addItem(item, false);
-            assertSame(item, mTarget.findLastItem(TestWriteItem::shouldKeepOnFilter,
-                    FilterableTestWriteItem.class));
+            assertSame(item, mTarget.findLastItem(TestItem::shouldKeepOnFilter,
+                    FilterableTestItem.class));
         }
     }
 
@@ -253,8 +235,8 @@
     public void testRemoveItemsNotRemoveDifferentType() throws Exception {
         mListener.setExpectedOnPreProcessItemCallbackTimes(1);
         synchronized (mTarget) {
-            mTarget.addItem(mFactory.createWriteItem(), false);
-            mTarget.removeItems(TestWriteItem::shouldKeepOnFilter, FilterableTestWriteItem.class);
+            mTarget.addItem(mFactory.createItem(), false);
+            mTarget.removeItems(TestItem::shouldKeepOnFilter, FilterableTestItem.class);
         }
         assertTrue("Target didn't call callback enough times.",
                 mListener.waitForAllExpectedCallbackDone(PRE_TASK_DELAY_MS + TIMEOUT_ALLOWANCE));
@@ -266,7 +248,7 @@
         mListener.setExpectedOnPreProcessItemCallbackTimes(1);
         synchronized (mTarget) {
             mTarget.addItem(mFactory.createFilterableItem(false), false);
-            mTarget.removeItems(TestWriteItem::shouldKeepOnFilter, FilterableTestWriteItem.class);
+            mTarget.removeItems(TestItem::shouldKeepOnFilter, FilterableTestItem.class);
         }
         assertTrue("Target didn't call callback enough times.",
                 mListener.waitForAllExpectedCallbackDone(PRE_TASK_DELAY_MS + TIMEOUT_ALLOWANCE));
@@ -276,8 +258,8 @@
     @Test
     public void testUpdateLastOrAddItemUpdatesMatchedItem() throws Exception {
         mListener.setExpectedOnPreProcessItemCallbackTimes(1);
-        final FilterableTestWriteItem scheduledItem = mFactory.createFilterableItem(true);
-        final FilterableTestWriteItem expected = mFactory.createFilterableItem(true);
+        final FilterableTestItem scheduledItem = mFactory.createFilterableItem(true);
+        final FilterableTestItem expected = mFactory.createFilterableItem(true);
         synchronized (mTarget) {
             mTarget.addItem(scheduledItem, false);
             mTarget.updateLastOrAddItem(expected, false);
@@ -292,8 +274,8 @@
     @Test
     public void testUpdateLastOrAddItemUpdatesAddItemWhenNoMatch() throws Exception {
         mListener.setExpectedOnPreProcessItemCallbackTimes(2);
-        final FilterableTestWriteItem scheduledItem = mFactory.createFilterableItem(false);
-        final FilterableTestWriteItem expected = mFactory.createFilterableItem(true);
+        final FilterableTestItem scheduledItem = mFactory.createFilterableItem(false);
+        final FilterableTestItem expected = mFactory.createFilterableItem(true);
         synchronized (mTarget) {
             mTarget.addItem(scheduledItem, false);
             mTarget.updateLastOrAddItem(expected, false);
@@ -310,9 +292,9 @@
     public void testRemoveItemsRemoveMatchedItem() throws Exception {
         mListener.setExpectedOnPreProcessItemCallbackTimes(1);
         synchronized (mTarget) {
-            mTarget.addItem(mFactory.createWriteItem(), false);
+            mTarget.addItem(mFactory.createItem(), false);
             mTarget.addItem(mFactory.createFilterableItem(true), false);
-            mTarget.removeItems(TestWriteItem::shouldKeepOnFilter, FilterableTestWriteItem.class);
+            mTarget.removeItems(TestItem::shouldKeepOnFilter, FilterableTestItem.class);
         }
         assertTrue("Target didn't call callback enough times.",
                 mListener.waitForAllExpectedCallbackDone(PRE_TASK_DELAY_MS + TIMEOUT_ALLOWANCE));
@@ -322,8 +304,8 @@
     @Test
     public void testFlushWaitSynchronously() {
         final long dispatchTime = SystemClock.uptimeMillis();
-        mTarget.addItem(mFactory.createWriteItem(), false);
-        mTarget.addItem(mFactory.createWriteItem(), false);
+        mTarget.addItem(mFactory.createItem(), false);
+        mTarget.addItem(mFactory.createItem(), false);
         mTarget.flush();
         assertEquals("Flush should wait until all items are processed before return.",
                 2, mFactory.getTotalProcessedItemCount());
@@ -353,18 +335,15 @@
             return new TestItem(mItemCount, mLatch);
         }
 
-        TestWriteItem createWriteItem() {
-            return new TestWriteItem(mItemCount, mLatch);
-        }
-
-        FilterableTestWriteItem createFilterableItem(boolean shouldKeepOnFilter) {
-            return new FilterableTestWriteItem(shouldKeepOnFilter, mItemCount, mLatch);
+        FilterableTestItem createFilterableItem(boolean shouldKeepOnFilter) {
+            return new FilterableTestItem(shouldKeepOnFilter, mItemCount, mLatch);
         }
     }
 
-    private static class TestItem implements PersisterQueue.QueueItem {
-        private final AtomicInteger mItemCount;
-        private final CountDownLatch mLatch;
+    private static class TestItem<T extends TestItem<T>>
+            implements PersisterQueue.WriteQueueItem<T> {
+        private AtomicInteger mItemCount;
+        private CountDownLatch mLatch;
 
         TestItem(AtomicInteger itemCount, CountDownLatch latch) {
             mItemCount = itemCount;
@@ -380,37 +359,30 @@
                 mLatch.countDown();
             }
         }
-    }
-
-    private static class TestWriteItem<T extends TestWriteItem<T>>
-            extends TestItem implements PersisterQueue.WriteQueueItem<T> {
-        TestWriteItem(AtomicInteger itemCount, CountDownLatch latch) {
-            super(itemCount, latch);
-        }
 
         boolean shouldKeepOnFilter() {
             return true;
         }
     }
 
-    private static class FilterableTestWriteItem extends TestWriteItem<FilterableTestWriteItem> {
+    private static class FilterableTestItem extends TestItem<FilterableTestItem> {
         private boolean mShouldKeepOnFilter;
 
-        private FilterableTestWriteItem mUpdateFromItem;
+        private FilterableTestItem mUpdateFromItem;
 
-        private FilterableTestWriteItem(boolean shouldKeepOnFilter, AtomicInteger mItemCount,
+        private FilterableTestItem(boolean shouldKeepOnFilter, AtomicInteger mItemCount,
                 CountDownLatch mLatch) {
             super(mItemCount, mLatch);
             mShouldKeepOnFilter = shouldKeepOnFilter;
         }
 
         @Override
-        public boolean matches(FilterableTestWriteItem item) {
+        public boolean matches(FilterableTestItem item) {
             return item.mShouldKeepOnFilter;
         }
 
         @Override
-        public void updateFrom(FilterableTestWriteItem item) {
+        public void updateFrom(FilterableTestItem item) {
             mUpdateFromItem = item;
         }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java b/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
index 4104999..12b7445 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
@@ -16,6 +16,8 @@
 
 package com.android.server.wm;
 
+import static android.tools.traces.Utils.busyWaitForDataSourceRegistration;
+
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verifyZeroInteractions;
@@ -28,7 +30,6 @@
 import static java.io.File.createTempFile;
 import static java.nio.file.Files.createTempDirectory;
 
-import android.os.ParcelFileDescriptor;
 import android.platform.test.annotations.Presubmit;
 import android.tools.ScenarioBuilder;
 import android.tools.traces.io.ResultWriter;
@@ -36,9 +37,6 @@
 import android.view.Choreographer;
 
 import androidx.test.filters.SmallTest;
-import androidx.test.platform.app.InstrumentationRegistry;
-
-import com.google.protobuf.InvalidProtocolBufferException;
 
 import org.junit.After;
 import org.junit.Before;
@@ -46,12 +44,9 @@
 import org.junit.Test;
 import org.mockito.Mockito;
 
-import perfetto.protos.PerfettoConfig.TracingServiceState;
 import perfetto.protos.PerfettoConfig.WindowManagerConfig.LogFrequency;
 
-import java.io.FileInputStream;
 import java.io.IOException;
-import java.util.Optional;
 
 /**
  * Test class for {@link WindowTracingPerfetto}.
@@ -74,7 +69,7 @@
         sChoreographer = Mockito.mock(Choreographer.class);
         sWindowTracing = new WindowTracingPerfetto(sWmMock, sChoreographer,
                 new WindowManagerGlobalLock(), TEST_DATA_SOURCE_NAME);
-        waitDataSourceIsAvailable();
+        busyWaitForDataSourceRegistration(TEST_DATA_SOURCE_NAME);
     }
 
     @Before
@@ -156,67 +151,4 @@
 
         mTraceMonitor.stop(writer);
     }
-
-    private static void waitDataSourceIsAvailable() {
-        final int timeoutMs = 10000;
-        final int busyWaitIntervalMs = 100;
-
-        int elapsedMs = 0;
-
-        while (!isDataSourceAvailable()) {
-            try {
-                Thread.sleep(busyWaitIntervalMs);
-                elapsedMs += busyWaitIntervalMs;
-                if (elapsedMs >= timeoutMs) {
-                    throw new RuntimeException("Data source didn't become available."
-                            + " Waited for: " + timeoutMs + " ms");
-                }
-            } catch (InterruptedException e) {
-                throw new RuntimeException(e);
-            }
-        }
-    }
-
-    private static boolean isDataSourceAvailable() {
-        byte[] proto = executeShellCommand("perfetto --query-raw");
-
-        try {
-            TracingServiceState state = TracingServiceState.parseFrom(proto);
-
-            Optional<Integer> producerId = Optional.empty();
-
-            for (TracingServiceState.Producer producer : state.getProducersList()) {
-                if (producer.getPid() == android.os.Process.myPid()) {
-                    producerId = Optional.of(producer.getId());
-                    break;
-                }
-            }
-
-            if (producerId.isEmpty()) {
-                return false;
-            }
-
-            for (TracingServiceState.DataSource ds : state.getDataSourcesList()) {
-                if (ds.getDsDescriptor().getName().equals(TEST_DATA_SOURCE_NAME)
-                        && ds.getProducerId() == producerId.get()) {
-                    return true;
-                }
-            }
-        } catch (InvalidProtocolBufferException e) {
-            throw new RuntimeException(e);
-        }
-
-        return false;
-    }
-
-    private static byte[] executeShellCommand(String command) {
-        try {
-            ParcelFileDescriptor fd = InstrumentationRegistry.getInstrumentation().getUiAutomation()
-                    .executeShellCommand(command);
-            FileInputStream is = new ParcelFileDescriptor.AutoCloseInputStream(fd);
-            return is.readAllBytes();
-        } catch (IOException e) {
-            throw new RuntimeException(e);
-        }
-    }
 }
diff --git a/telephony/java/android/telephony/BarringInfo.java b/telephony/java/android/telephony/BarringInfo.java
index e20e4d2..e42b41f 100644
--- a/telephony/java/android/telephony/BarringInfo.java
+++ b/telephony/java/android/telephony/BarringInfo.java
@@ -159,7 +159,7 @@
 
         /**
          * @return the conditional barring factor as a percentage 0-100, which is the probability of
-         *         a random device being barred for the service type.
+         *         a random device being allowed for a conditionally barred service.
          */
         public int getConditionalBarringFactor() {
             return mConditionalBarringFactor;
diff --git a/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java b/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
index cfb2645..6f3deab 100644
--- a/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
+++ b/tests/Tracing/src/com/android/internal/protolog/PerfettoProtoLogImplTest.java
@@ -16,7 +16,7 @@
 
 package com.android.internal.protolog;
 
-import static android.tools.traces.Utils.executeShellCommand;
+import static android.tools.traces.Utils.busyWaitForDataSourceRegistration;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertThrows;
@@ -50,7 +50,6 @@
 import com.android.internal.protolog.common.LogLevel;
 
 import com.google.common.truth.Truth;
-import com.google.protobuf.InvalidProtocolBufferException;
 
 import org.junit.After;
 import org.junit.Before;
@@ -60,14 +59,12 @@
 import org.junit.runners.JUnit4;
 import org.mockito.Mockito;
 
-import perfetto.protos.PerfettoConfig.TracingServiceState;
 import perfetto.protos.Protolog;
 import perfetto.protos.ProtologCommon;
 
 import java.io.File;
 import java.io.IOException;
 import java.util.List;
-import java.util.Optional;
 import java.util.Random;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -184,7 +181,7 @@
                     TestProtoLogGroup.values(), dataSourceBuilder, sProtoLogConfigurationService);
         }
 
-        waitDataSourceIsAvailable();
+        busyWaitForDataSourceRegistration(TEST_PROTOLOG_DATASOURCE_NAME);
     }
 
     @Before
@@ -870,54 +867,6 @@
                 .isEqualTo("This message should also be logged 567");
     }
 
-    private static void waitDataSourceIsAvailable() {
-        final int timeoutMs = 10000;
-        final int busyWaitIntervalMs = 100;
-
-        int elapsedMs = 0;
-
-        while (!isDataSourceAvailable()) {
-            SystemClock.sleep(busyWaitIntervalMs);
-            elapsedMs += busyWaitIntervalMs;
-            if (elapsedMs >= timeoutMs) {
-                throw new RuntimeException("Data source didn't become available."
-                        + " Waited for: " + timeoutMs + " ms");
-            }
-        }
-    }
-
-    private static boolean isDataSourceAvailable() {
-        byte[] proto = executeShellCommand("perfetto --query-raw");
-
-        try {
-            TracingServiceState state = TracingServiceState.parseFrom(proto);
-
-            Optional<Integer> producerId = Optional.empty();
-
-            for (TracingServiceState.Producer producer : state.getProducersList()) {
-                if (producer.getPid() == android.os.Process.myPid()) {
-                    producerId = Optional.of(producer.getId());
-                    break;
-                }
-            }
-
-            if (producerId.isEmpty()) {
-                return false;
-            }
-
-            for (TracingServiceState.DataSource ds : state.getDataSourcesList()) {
-                if (ds.getDsDescriptor().getName().equals(TEST_PROTOLOG_DATASOURCE_NAME)
-                        && ds.getProducerId() == producerId.get()) {
-                    return true;
-                }
-            }
-        } catch (InvalidProtocolBufferException e) {
-            throw new RuntimeException(e);
-        }
-
-        return false;
-    }
-
     private enum TestProtoLogGroup implements IProtoLogGroup {
         TEST_GROUP(true, true, false, "TEST_TAG");