Merge "Improving documentation in ProcessingSignal and updateProcessingState" into main
diff --git a/core/java/android/companion/AssociationInfo.java b/core/java/android/companion/AssociationInfo.java
index 843158c..b4b96e2 100644
--- a/core/java/android/companion/AssociationInfo.java
+++ b/core/java/android/companion/AssociationInfo.java
@@ -252,6 +252,14 @@
     }
 
     /**
+     * @return true if the association is not revoked nor pending
+     * @hide
+     */
+    public boolean isActive() {
+        return !mRevoked && !mPending;
+    }
+
+    /**
      * @return the last time self reported disconnected for selfManaged only.
      * @hide
      */
diff --git a/core/java/android/hardware/display/DisplayManager.java b/core/java/android/hardware/display/DisplayManager.java
index eb26a76..4894fb1 100644
--- a/core/java/android/hardware/display/DisplayManager.java
+++ b/core/java/android/hardware/display/DisplayManager.java
@@ -1653,6 +1653,22 @@
     }
 
     /**
+     * Allows internal application to restrict display modes to specified modeIds
+     *
+     * @param displayId display that restrictions will be applied to
+     * @param modeIds allowed mode ids
+     *
+     * @hide
+     */
+    @RequiresPermission("android.permission.RESTRICT_DISPLAY_MODES")
+    public void requestDisplayModes(int displayId, @Nullable int[] modeIds) {
+        if (modeIds != null && modeIds.length == 0) {
+            throw new IllegalArgumentException("requestDisplayModes: modesIds can't be empty");
+        }
+        mGlobal.requestDisplayModes(displayId, modeIds);
+    }
+
+    /**
      * Listens for changes in available display devices.
      */
     public interface DisplayListener {
diff --git a/core/java/android/hardware/display/DisplayManagerGlobal.java b/core/java/android/hardware/display/DisplayManagerGlobal.java
index 75f0ceb..3d7b714 100644
--- a/core/java/android/hardware/display/DisplayManagerGlobal.java
+++ b/core/java/android/hardware/display/DisplayManagerGlobal.java
@@ -38,6 +38,7 @@
 import android.hardware.graphics.common.DisplayDecorationSupport;
 import android.media.projection.IMediaProjection;
 import android.media.projection.MediaProjection;
+import android.os.Binder;
 import android.os.Handler;
 import android.os.HandlerExecutor;
 import android.os.IBinder;
@@ -138,6 +139,8 @@
 
     private int mWifiDisplayScanNestCount;
 
+    private final Binder mToken = new Binder();
+
     @VisibleForTesting
     public DisplayManagerGlobal(IDisplayManager dm) {
         mDm = dm;
@@ -1182,6 +1185,20 @@
         }
     }
 
+    /**
+     * Sets allowed display mode ids
+     *
+     * @hide
+     */
+    @RequiresPermission("android.permission.RESTRICT_DISPLAY_MODES")
+    public void requestDisplayModes(int displayId, @Nullable int[] modeIds) {
+        try {
+            mDm.requestDisplayModes(mToken, displayId, modeIds);
+        } catch (RemoteException ex) {
+            throw ex.rethrowFromSystemServer();
+        }
+    }
+
     private final class DisplayManagerCallback extends IDisplayManagerCallback.Stub {
         @Override
         public void onDisplayEvent(int displayId, @DisplayEvent int event) {
diff --git a/core/java/android/hardware/display/IDisplayManager.aidl b/core/java/android/hardware/display/IDisplayManager.aidl
index 83de4e4..70efc6f 100644
--- a/core/java/android/hardware/display/IDisplayManager.aidl
+++ b/core/java/android/hardware/display/IDisplayManager.aidl
@@ -235,4 +235,8 @@
     // Disable a connected display that is enabled.
     @EnforcePermission("MANAGE_DISPLAYS")
     void disableConnectedDisplay(int displayId);
+
+    // Restricts display modes to specified modeIds.
+    @EnforcePermission("RESTRICT_DISPLAY_MODES")
+    void requestDisplayModes(in IBinder token, int displayId, in @nullable int[] modeIds);
 }
diff --git a/core/java/com/android/internal/statusbar/IStatusBar.aidl b/core/java/com/android/internal/statusbar/IStatusBar.aidl
index a2efbd2..a22232a 100644
--- a/core/java/com/android/internal/statusbar/IStatusBar.aidl
+++ b/core/java/com/android/internal/statusbar/IStatusBar.aidl
@@ -376,6 +376,12 @@
     void moveFocusedTaskToStageSplit(int displayId, boolean leftOrTop);
 
     /**
+     * Set the split screen focus to the left / top app or the right / bottom app based on
+     * {@param leftOrTop}.
+     */
+    void setSplitscreenFocus(boolean leftOrTop);
+
+    /**
      * Shows the media output switcher dialog.
      *
      * @param packageName of the session for which the output switcher is shown.
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index 0869af5..7e46818 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -8150,6 +8150,14 @@
     <permission android:name="android.permission.EMERGENCY_INSTALL_PACKAGES"
                 android:protectionLevel="signature|privileged"/>
 
+    <!-- Allows internal applications to restrict display modes
+       <p>Not for use by third-party applications.
+       <p>Protection level: signature
+       @hide
+    -->
+    <permission android:name="android.permission.RESTRICT_DISPLAY_MODES"
+        android:protectionLevel="signature" />
+
     <!-- Attribution for Geofencing service. -->
     <attribution android:tag="GeofencingService" android:label="@string/geofencing_service"/>
     <!-- Attribution for Country Detector. -->
diff --git a/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
index 48ba526..fc1fbb5 100644
--- a/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
@@ -120,6 +120,10 @@
             final var statsToken = ImeTracker.Token.empty();
             mImeConsumer.onWindowFocusGained(true);
             mController.show(WindowInsets.Type.ime(), true /* fromIme */, statsToken);
+            // Called once through the show flow.
+            verify(mController).applyAnimation(
+                    eq(WindowInsets.Type.ime()), eq(true) /* show */, eq(true) /* fromIme */,
+                    eq(statsToken));
 
             // set control and verify visibility is applied.
             InsetsSourceControl control = new InsetsSourceControl(ID_IME,
@@ -127,11 +131,11 @@
             mController.onControlsChanged(new InsetsSourceControl[] { control });
             // IME show animation should be triggered when control becomes available.
             verify(mController).applyAnimation(
-                    eq(WindowInsets.Type.ime()), eq(true) /* show */, eq(true) /* fromIme */,
-                    eq(statsToken));
+                    eq(WindowInsets.Type.ime()), eq(true) /* show */, eq(false) /* fromIme */,
+                    and(not(eq(statsToken)), notNull()));
             verify(mController, never()).applyAnimation(
-                    eq(WindowInsets.Type.ime()), eq(false) /* show */, eq(true) /* fromIme */,
-                    eq(statsToken));
+                    eq(WindowInsets.Type.ime()), eq(false) /* show */, eq(false) /* fromIme */,
+                    and(not(eq(statsToken)), notNull()));
         });
     }
 
@@ -159,6 +163,10 @@
             final var statsToken = ImeTracker.Token.empty();
             if (imeVisible) {
                 mController.show(WindowInsets.Type.ime(), true /* fromIme */, statsToken);
+                // Called once through the show flow.
+                verify(mController).applyAnimation(eq(WindowInsets.Type.ime()),
+                        eq(true) /* show */, eq(true) /* fromIme */,
+                        eq(false) /* skipAnim */, eq(statsToken));
             }
 
             // set control and verify visibility is applied.
@@ -184,13 +192,17 @@
             if (!hasViewFocus) {
                 final var statsTokenNext = ImeTracker.Token.empty();
                 mController.show(WindowInsets.Type.ime(), true /* fromIme */, statsTokenNext);
+                // Called once through the show flow.
+                verify(mController).applyAnimation(eq(WindowInsets.Type.ime()),
+                        eq(true) /* show */, eq(true) /* fromIme */,
+                        eq(false) /* skipAnim */, eq(statsTokenNext));
                 mController.onControlsChanged(new InsetsSourceControl[]{ control });
                 // Verify IME show animation should be triggered when control becomes available and
                 // the animation will be skipped by getAndClearSkipAnimationOnce invoked.
                 verify(control).getAndClearSkipAnimationOnce();
                 verify(mController).applyAnimation(eq(WindowInsets.Type.ime()),
-                        eq(true) /* show */, eq(true) /* fromIme */,
-                        eq(false) /* skipAnim */, eq(statsTokenNext));
+                        eq(true) /* show */, eq(false) /* fromIme */,
+                        eq(true) /* skipAnim */, and(not(eq(statsToken)), notNull()));
             }
         });
     }
diff --git a/libs/WindowManager/Shell/res/values/dimen.xml b/libs/WindowManager/Shell/res/values/dimen.xml
index 7dd3961..00fb298 100644
--- a/libs/WindowManager/Shell/res/values/dimen.xml
+++ b/libs/WindowManager/Shell/res/values/dimen.xml
@@ -272,6 +272,10 @@
     <dimen name="bubble_bar_expanded_view_corner_radius">16dp</dimen>
     <!-- Corner radius for expanded view while it is being dragged -->
     <dimen name="bubble_bar_expanded_view_corner_radius_dragged">28dp</dimen>
+    <!-- Width of the box around bottom center of the screen where drag only leads to dismiss -->
+    <dimen name="bubble_bar_dismiss_zone_width">192dp</dimen>
+    <!-- Height of the box around bottom center of the screen where drag only leads to dismiss -->
+    <dimen name="bubble_bar_dismiss_zone_height">242dp</dimen>
 
     <!-- Bottom and end margin for compat buttons. -->
     <dimen name="compat_button_margin">24dp</dimen>
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt
index 056598b..b5b8a81 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt
@@ -17,8 +17,10 @@
 package com.android.wm.shell.bubbles.bar
 
 import android.annotation.SuppressLint
+import android.graphics.RectF
 import android.view.MotionEvent
 import android.view.View
+import com.android.wm.shell.R
 import com.android.wm.shell.bubbles.BubblePositioner
 import com.android.wm.shell.common.bubbles.BubbleBarLocation
 import com.android.wm.shell.common.bubbles.DismissView
@@ -43,6 +45,8 @@
     private val magnetizedExpandedView: MagnetizedObject<BubbleBarExpandedView> =
         MagnetizedObject.magnetizeView(expandedView)
     private val magnetizedDismissTarget: MagnetizedObject.MagneticTarget
+    private val dismissZoneHeight: Int
+    private val dismissZoneWidth: Int
 
     init {
         magnetizedExpandedView.magnetListener = MagnetListener()
@@ -74,6 +78,11 @@
             }
             return@setOnTouchListener dragMotionEventHandler.onTouch(view, event) || magnetConsumed
         }
+
+        dismissZoneHeight =
+            dismissView.resources.getDimensionPixelSize(R.dimen.bubble_bar_dismiss_zone_height)
+        dismissZoneWidth =
+            dismissView.resources.getDimensionPixelSize(R.dimen.bubble_bar_dismiss_zone_width)
     }
 
     /** Listener to receive callback about dragging events */
@@ -97,12 +106,23 @@
         private var isMoving = false
         private var screenCenterX: Int = -1
         private var isOnLeft = false
+        private val dismissZone = RectF()
 
         override fun onDown(v: View, ev: MotionEvent): Boolean {
             // While animating, don't allow new touch events
             if (expandedView.isAnimating) return false
-            screenCenterX = bubblePositioner.screenRect.centerX()
             isOnLeft = bubblePositioner.isBubbleBarOnLeft
+
+            val screenRect = bubblePositioner.screenRect
+            screenCenterX = screenRect.centerX()
+            val screenBottom = screenRect.bottom
+
+            dismissZone.set(
+                screenCenterX - dismissZoneWidth / 2f,
+                (screenBottom - dismissZoneHeight).toFloat(),
+                screenCenterX + dismissZoneHeight / 2f,
+                screenBottom.toFloat()
+            )
             return true
         }
 
@@ -122,6 +142,11 @@
             expandedView.translationY = expandedViewInitialTranslationY + dy
             dismissView.show()
 
+            // Check if we are in the zone around dismiss view where drag can only lead to dismiss
+            if (dismissZone.contains(ev.rawX, ev.rawY)) {
+                return
+            }
+
             if (isOnLeft && ev.rawX > screenCenterX) {
                 isOnLeft = false
                 dragListener.onLocationChanged(BubbleBarLocation.RIGHT)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
index f7ffdd1..b99e943 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
@@ -105,6 +105,9 @@
     /** Called when requested to go to fullscreen from the current active split app. */
     void goToFullscreenFromSplit();
 
+    /** Called when splitscreen focused app is changed. */
+    void setSplitscreenFocus(boolean leftOrTop);
+
     /** Get a string representation of a stage type */
     static String stageTypeToString(@StageType int stage) {
         switch (stage) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 218abe2..5fbb152 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -485,6 +485,12 @@
         }
     }
 
+    public void setSplitscreenFocus(boolean leftOrTop) {
+        if (mStageCoordinator.isSplitActive()) {
+            mStageCoordinator.grantFocusToPosition(leftOrTop);
+        }
+    }
+
     /** Move the specified task to fullscreen, regardless of focus state. */
     public void moveTaskToFullscreen(int taskId, int exitReason) {
         mStageCoordinator.moveTaskToFullscreen(taskId, exitReason);
@@ -1152,6 +1158,12 @@
         public void goToFullscreenFromSplit() {
             mMainExecutor.execute(SplitScreenController.this::goToFullscreenFromSplit);
         }
+
+        @Override
+        public void setSplitscreenFocus(boolean leftOrTop) {
+            mMainExecutor.execute(
+                    () -> SplitScreenController.this.setSplitscreenFocus(leftOrTop));
+        }
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 8ab6cf7..661faf5 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -1604,6 +1604,11 @@
         }
     }
 
+    protected void grantFocusToPosition(boolean leftOrTop) {
+        grantFocusToStage(mSideStagePosition == SPLIT_POSITION_BOTTOM_OR_RIGHT
+                ? getMainStagePosition() : getSideStagePosition());
+    }
+
     private void clearRequestIfPresented() {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "clearRequestIfPresented");
         if (mSideStageListener.mVisible && mSideStageListener.mHasChildren
diff --git a/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java b/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java
index 3266c12..b151a53 100644
--- a/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java
+++ b/packages/SettingsProvider/src/com/android/providers/settings/SettingsState.java
@@ -60,10 +60,12 @@
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
 
+import java.io.BufferedWriter;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
+import java.io.FileWriter;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.nio.file.Files;
@@ -161,6 +163,11 @@
     private static final String APEX_DIR = "/apex";
     private static final String APEX_ACONFIG_PATH_SUFFIX = "/etc/aconfig_flags.pb";
 
+    private static final String STORAGE_MIGRATION_FLAG =
+            "core_experiments_team_internal/com.android.providers.settings.storage_test_mission_1";
+    private static final String STORAGE_MIGRATION_LOG =
+            "/metadata/aconfig/flags/storage_migration.log";
+
     /**
      * This tag is applied to all aconfig default value-loaded flags.
      */
@@ -1439,6 +1446,20 @@
                     }
                 }
 
+                if (name != null && name.equals(STORAGE_MIGRATION_FLAG) && value.equals("true")) {
+                    File file = new File(STORAGE_MIGRATION_LOG);
+                    if (!file.exists()) {
+                        try (BufferedWriter writer =
+                                new BufferedWriter(new FileWriter(STORAGE_MIGRATION_LOG))) {
+                            final long timestamp = System.currentTimeMillis();
+                            String entry = String.format("%d | Log init", timestamp);
+                            writer.write(entry);
+                        } catch (IOException e) {
+                            Slog.e(LOG_TAG, "failed to write storage migration file", e);
+                        }
+                    }
+                }
+
                 mSettings.put(name, new Setting(name, value, defaultValue, packageName, tag,
                         fromSystem, id, isPreservedInRestore));
 
diff --git a/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig b/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig
index 2e14e9b..c572bdb 100644
--- a/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig
+++ b/packages/SettingsProvider/src/com/android/providers/settings/device_config_service.aconfig
@@ -33,3 +33,10 @@
     bug: "327383546"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "storage_test_mission_1"
+    namespace: "core_experiments_team_internal"
+    description: "If this flag is detected as true on boot, writes a logfile to track storage migration correctness."
+    bug: "328444881"
+}
diff --git a/packages/SystemUI/AndroidManifest.xml b/packages/SystemUI/AndroidManifest.xml
index 3c18f17..f123201 100644
--- a/packages/SystemUI/AndroidManifest.xml
+++ b/packages/SystemUI/AndroidManifest.xml
@@ -434,7 +434,7 @@
             </intent-filter>
         </receiver>
 
-        <activity android:name=".screenshot.LongScreenshotActivity"
+        <activity android:name=".screenshot.scroll.LongScreenshotActivity"
                   android:theme="@style/LongScreenshotActivity"
                   android:process=":screenshot"
                   android:exported="false"
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
index 63ec54f..82083f9 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
@@ -605,6 +605,8 @@
 
     override val isInitiatedByUserInput = true
 
+    override var bouncingScene: SceneKey? = null
+
     /** The current offset caused by the drag gesture. */
     var dragOffset by mutableFloatStateOf(0f)
 
@@ -694,14 +696,31 @@
     ): OffsetAnimation {
         return startOffsetAnimation {
             val animatable = Animatable(dragOffset, OffsetVisibilityThreshold)
+            val isTargetGreater = targetOffset > animatable.value
             val job =
                 coroutineScope
                     .launch {
-                        animatable.animateTo(
-                            targetValue = targetOffset,
-                            animationSpec = swipeSpec,
-                            initialVelocity = initialVelocity,
-                        )
+                        try {
+                            animatable.animateTo(
+                                targetValue = targetOffset,
+                                animationSpec = swipeSpec,
+                                initialVelocity = initialVelocity,
+                            ) {
+                                if (bouncingScene == null) {
+                                    val isBouncing =
+                                        if (isTargetGreater) {
+                                            value > targetOffset
+                                        } else {
+                                            value < targetOffset
+                                        }
+                                    if (isBouncing) {
+                                        bouncingScene = targetScene
+                                    }
+                                }
+                            }
+                        } finally {
+                            bouncingScene = null
+                        }
                     }
                     // Make sure that we settle to target scene at the end of the animation or if
                     // the animation is cancelled.
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
index c7186da..f1177a8 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
@@ -588,7 +588,8 @@
             // TODO(b/290184746): Make sure that we don't overflow transformations associated to a
             // range.
             val directionSign = if (transition.isUpOrLeft) -1 else 1
-            val overscrollProgress = transition.progress.let { if (it > 1f) it - 1f else it }
+            val isToScene = overscroll.scene == transition.toScene
+            val overscrollProgress = transition.progress.let { if (isToScene) it - 1f else it }
             val progress = directionSign * overscrollProgress
             val rangeProgress = propertySpec.range?.progress(progress) ?: progress
 
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
index e6f5d58..617a8ea 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
@@ -255,6 +255,12 @@
          */
         val overscrollScope: OverscrollScope
 
+        /**
+         * The scene around which the transition is currently bouncing. When not `null`, this
+         * transition is currently oscillating around this scene and will soon settle to that scene.
+         */
+        val bouncingScene: SceneKey?
+
         companion object {
             const val DistanceUnspecified = 0f
         }
@@ -287,9 +293,10 @@
             val transition = currentTransition ?: return null
             if (transition !is TransitionState.HasOverscrollProperties) return null
             val progress = transition.progress
+            val bouncingScene = transition.bouncingScene
             return when {
-                progress < 0f -> fromOverscrollSpec
-                progress > 1f -> toOverscrollSpec
+                progress < 0f || bouncingScene == transition.fromScene -> fromOverscrollSpec
+                progress > 1f || bouncingScene == transition.toScene -> toOverscrollSpec
                 else -> null
             }
         }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
index 26e01fe..0804761 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
@@ -16,6 +16,8 @@
 
 package com.android.compose.animation.scene
 
+import androidx.compose.animation.core.Spring
+import androidx.compose.animation.core.spring
 import androidx.compose.animation.core.tween
 import androidx.compose.foundation.ExperimentalFoundationApi
 import androidx.compose.foundation.gestures.Orientation
@@ -752,4 +754,55 @@
         assertThat(state.currentOverscrollSpec).isNotNull()
         fooElement.assertTopPositionInRootIsEqualTo(layoutHeight * 1.5f)
     }
+
+    @Test
+    fun elementTransitionWithDistanceDuringOverscrollBouncing() {
+        val layoutWidth = 200.dp
+        val layoutHeight = 400.dp
+        val state =
+            setupOverscrollScenario(
+                layoutWidth = layoutWidth,
+                layoutHeight = layoutHeight,
+                sceneTransitions = {
+                    defaultSwipeSpec =
+                        spring(
+                            dampingRatio = Spring.DampingRatioMediumBouncy,
+                            stiffness = Spring.StiffnessLow,
+                        )
+
+                    overscroll(TestScenes.SceneB, Orientation.Vertical) {
+                        // On overscroll 100% -> Foo should translate by layoutHeight
+                        translate(TestElements.Foo, y = { absoluteDistance })
+                    }
+                },
+                firstScroll = 1f, // 100% scroll
+            )
+
+        val fooElement = rule.onNodeWithTag(TestElements.Foo.testTag, useUnmergedTree = true)
+        fooElement.assertTopPositionInRootIsEqualTo(0.dp)
+
+        rule.onRoot().performTouchInput {
+            // Scroll another 50%
+            moveBy(Offset(0f, layoutHeight.toPx() * 0.5f), delayMillis = 1_000)
+        }
+
+        val transition = state.currentTransition
+        assertThat(transition).isNotNull()
+        transition as TransitionState.HasOverscrollProperties
+
+        // Scroll 150% (100% scroll + 50% overscroll)
+        assertThat(transition.progress).isEqualTo(1.5f)
+        assertThat(state.currentOverscrollSpec).isNotNull()
+        fooElement.assertTopPositionInRootIsEqualTo(layoutHeight * (transition.progress - 1f))
+
+        // finger raised
+        rule.onRoot().performTouchInput { up() }
+
+        // The target value is 1f, but the spring (defaultSwipeSpec) allows you to go to a lower
+        // value.
+        rule.waitUntil(timeoutMillis = 10_000) { transition.progress < 1f }
+
+        assertThat(state.currentOverscrollSpec).isNotNull()
+        assertThat(transition.bouncingScene).isEqualTo(transition.toScene)
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
index 73a66c6..a32fe22 100644
--- a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/Transition.kt
@@ -27,6 +27,7 @@
     isInitiatedByUserInput: Boolean = false,
     isUserInputOngoing: Boolean = false,
     isUpOrLeft: Boolean = false,
+    bouncingScene: SceneKey? = null,
     orientation: Orientation = Orientation.Horizontal,
 ): TransitionState.Transition {
     return object : TransitionState.Transition(from, to), TransitionState.HasOverscrollProperties {
@@ -37,6 +38,7 @@
         override val isInitiatedByUserInput: Boolean = isInitiatedByUserInput
         override val isUserInputOngoing: Boolean = isUserInputOngoing
         override val isUpOrLeft: Boolean = isUpOrLeft
+        override val bouncingScene: SceneKey? = bouncingScene
         override val orientation: Orientation = orientation
         override val overscrollScope: OverscrollScope =
             object : OverscrollScope {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
index 24c651f3..a9541d9 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorTest.kt
@@ -61,6 +61,7 @@
     private val testScope = kosmos.testScope
     private val repository by lazy { kosmos.fakeKeyguardRepository }
     private val sceneInteractor by lazy { kosmos.sceneInteractor }
+    private val fromGoneTransitionInteractor by lazy { kosmos.fromGoneTransitionInteractor }
     private val commandQueue by lazy { FakeCommandQueue() }
     private val bouncerRepository = FakeKeyguardBouncerRepository()
     private val shadeRepository = FakeShadeRepository()
@@ -79,6 +80,7 @@
             shadeRepository = shadeRepository,
             keyguardTransitionInteractor = kosmos.keyguardTransitionInteractor,
             sceneInteractorProvider = { sceneInteractor },
+            fromGoneTransitionInteractor = { fromGoneTransitionInteractor },
         )
     }
 
diff --git a/packages/SystemUI/res/layout/app_clips_screenshot.xml b/packages/SystemUI/res/layout/app_clips_screenshot.xml
index cb638ee..bcc7bca 100644
--- a/packages/SystemUI/res/layout/app_clips_screenshot.xml
+++ b/packages/SystemUI/res/layout/app_clips_screenshot.xml
@@ -69,7 +69,7 @@
         tools:minHeight="100dp"
         tools:minWidth="100dp" />
 
-    <com.android.systemui.screenshot.CropView
+    <com.android.systemui.screenshot.scroll.CropView
         android:id="@+id/crop_view"
         android:layout_width="0px"
         android:layout_height="0px"
diff --git a/packages/SystemUI/res/layout/long_screenshot.xml b/packages/SystemUI/res/layout/long_screenshot.xml
index 8a19c2e..4d207da 100644
--- a/packages/SystemUI/res/layout/long_screenshot.xml
+++ b/packages/SystemUI/res/layout/long_screenshot.xml
@@ -100,7 +100,7 @@
         app:layout_constraintStart_toStartOf="parent"
         android:transitionName="screenshot_preview_image"/>
 
-    <com.android.systemui.screenshot.CropView
+    <com.android.systemui.screenshot.scroll.CropView
         android:id="@+id/crop_view"
         android:layout_width="0px"
         android:layout_height="0px"
@@ -122,7 +122,7 @@
         tools:minHeight="100dp"
         tools:minWidth="100dp" />
 
-    <com.android.systemui.screenshot.MagnifierView
+    <com.android.systemui.screenshot.scroll.MagnifierView
         android:id="@+id/magnifier"
         android:visibility="invisible"
         android:layout_width="200dp"
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index e3a5e15..774bbe5 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -3358,7 +3358,7 @@
     <string name="microphone_blocked_dream_overlay_content_description">Microphone blocked</string>
 
     <!-- Content description for priority mode icon on dream [CHAR LIMIT=NONE] -->
-    <string name="priority_mode_dream_overlay_content_description">Priority mode on</string>
+    <string name="priority_mode_dream_overlay_content_description">Do not disturb</string>
 
     <!-- Content description for when assistant attention is active [CHAR LIMIT=NONE] -->
     <string name="assistant_attention_content_description">User presence is detected</string>
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java b/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java
index 9afd5ed..d2df276 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/DefaultActivityBinder.java
@@ -24,9 +24,9 @@
 import com.android.systemui.keyguard.WorkLockActivity;
 import com.android.systemui.people.PeopleSpaceActivity;
 import com.android.systemui.people.widget.LaunchConversationActivity;
-import com.android.systemui.screenshot.LongScreenshotActivity;
 import com.android.systemui.screenshot.appclips.AppClipsActivity;
 import com.android.systemui.screenshot.appclips.AppClipsTrampolineActivity;
+import com.android.systemui.screenshot.scroll.LongScreenshotActivity;
 import com.android.systemui.sensorprivacy.SensorUseStartedActivity;
 import com.android.systemui.settings.brightness.BrightnessDialog;
 import com.android.systemui.telephony.ui.activity.SwitchToManagedProfileForCallActivity;
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 43a8b40..3b34750 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -40,6 +40,7 @@
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_AFTER_USER_LOCKDOWN;
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.STRONG_AUTH_REQUIRED_FOR_UNATTENDED_UPDATE;
 import static com.android.systemui.DejankUtils.whitelistIpcs;
+import static com.android.systemui.Flags.migrateClocksToBlueprint;
 import static com.android.systemui.Flags.notifyPowerManagerUserActivityBackground;
 import static com.android.systemui.Flags.refactorGetCurrentUser;
 import static com.android.systemui.keyguard.ui.viewmodel.LockscreenToDreamingTransitionViewModel.DREAMING_ANIMATION_DURATION_MS;
@@ -3401,6 +3402,11 @@
                 mSurfaceBehindRemoteAnimationFinishedCallback = null;
             }
         }
+
+        // Ensure that keyguard becomes visible if the going away animation is canceled
+        if (showKeyguard && !KeyguardWmStateRefactor.isEnabled() && migrateClocksToBlueprint()) {
+            mKeyguardInteractor.showKeyguard();
+        }
     }
 
     private void adjustStatusBarLocked() {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromGoneTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromGoneTransitionInteractor.kt
index d5a9bd1..4a3232e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromGoneTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromGoneTransitionInteractor.kt
@@ -71,6 +71,10 @@
         listenForGoneToDreamingLockscreenHosted()
     }
 
+    fun showKeyguard() {
+        scope.launch { startTransitionTo(KeyguardState.LOCKSCREEN) }
+    }
+
     // Primarily for when the user chooses to lock down the device
     private fun listenForGoneToLockscreenOrHub() {
         if (KeyguardWmStateRefactor.isEnabled) {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
index 7734973..283f160 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractor.kt
@@ -83,6 +83,7 @@
     shadeRepository: ShadeRepository,
     keyguardTransitionInteractor: KeyguardTransitionInteractor,
     sceneInteractorProvider: Provider<SceneInteractor>,
+    private val fromGoneTransitionInteractor: Provider<FromGoneTransitionInteractor>,
 ) {
     // TODO(b/296118689): move to a repository
     private val _sharedNotificationContainerBounds = MutableStateFlow(NotificationContainerBounds())
@@ -383,6 +384,11 @@
         repository.topClippingBounds.value = top
     }
 
+    /** Temporary shim, until [KeyguardWmStateRefactor] is enabled */
+    fun showKeyguard() {
+        fromGoneTransitionInteractor.get().showKeyguard()
+    }
+
     companion object {
         private const val TAG = "KeyguardInteractor"
     }
diff --git a/packages/SystemUI/src/com/android/systemui/people/widget/PeopleSpaceWidgetManager.java b/packages/SystemUI/src/com/android/systemui/people/widget/PeopleSpaceWidgetManager.java
index 1d820a1..5f0635b 100644
--- a/packages/SystemUI/src/com/android/systemui/people/widget/PeopleSpaceWidgetManager.java
+++ b/packages/SystemUI/src/com/android/systemui/people/widget/PeopleSpaceWidgetManager.java
@@ -21,6 +21,9 @@
 import static android.app.NotificationManager.INTERRUPTION_FILTER_ALL;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_NONE;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
+import static android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_HOME_SCREEN;
+import static android.appwidget.AppWidgetProviderInfo.WIDGET_CATEGORY_KEYGUARD;
+import static android.appwidget.flags.Flags.generatedPreviews;
 import static android.content.Intent.ACTION_BOOT_COMPLETED;
 import static android.content.Intent.ACTION_PACKAGE_ADDED;
 import static android.content.Intent.ACTION_PACKAGE_REMOVED;
@@ -80,12 +83,15 @@
 import android.service.notification.ZenModeConfig;
 import android.text.TextUtils;
 import android.util.Log;
+import android.util.SparseBooleanArray;
 import android.widget.RemoteViews;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.UiEventLogger;
 import com.android.internal.logging.UiEventLoggerImpl;
+import com.android.keyguard.KeyguardUpdateMonitor;
+import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.systemui.Dumpable;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.SysUISingleton;
@@ -96,6 +102,8 @@
 import com.android.systemui.people.PeopleSpaceUtils;
 import com.android.systemui.people.PeopleTileViewHelper;
 import com.android.systemui.people.SharedPreferencesHelper;
+import com.android.systemui.res.R;
+import com.android.systemui.settings.UserTracker;
 import com.android.systemui.statusbar.NotificationListener;
 import com.android.systemui.statusbar.NotificationListener.NotificationHandler;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
@@ -160,13 +168,27 @@
     @GuardedBy("mLock")
     public static Map<Integer, PeopleSpaceTile> mTiles = new HashMap<>();
 
+    @NonNull private final UserTracker mUserTracker;
+    @NonNull private final SparseBooleanArray mUpdatedPreviews = new SparseBooleanArray();
+    @NonNull private final KeyguardUpdateMonitorCallback mKeyguardUpdateMonitorCallback =
+            new KeyguardUpdateMonitorCallback() {
+                @Override
+                public void onUserUnlocked() {
+                    if (DEBUG) {
+                        Log.d(TAG, "onUserUnlocked " + mUserTracker.getUserId());
+                    }
+                    updateGeneratedPreviewForUser(mUserTracker.getUserHandle());
+                }
+            };
+
     @Inject
     public PeopleSpaceWidgetManager(Context context, LauncherApps launcherApps,
             CommonNotifCollection notifCollection,
             PackageManager packageManager, Optional<Bubbles> bubblesOptional,
             UserManager userManager, NotificationManager notificationManager,
             BroadcastDispatcher broadcastDispatcher, @Background Executor bgExecutor,
-            DumpManager dumpManager) {
+            DumpManager dumpManager, @NonNull UserTracker userTracker,
+            @NonNull KeyguardUpdateMonitor keyguardUpdateMonitor) {
         if (DEBUG) Log.d(TAG, "constructor");
         mContext = context;
         mAppWidgetManager = AppWidgetManager.getInstance(context);
@@ -187,6 +209,8 @@
         mBroadcastDispatcher = broadcastDispatcher;
         mBgExecutor = bgExecutor;
         dumpManager.registerNormalDumpable(TAG, this);
+        mUserTracker = userTracker;
+        keyguardUpdateMonitor.registerCallback(mKeyguardUpdateMonitorCallback);
     }
 
     /** Initializes {@PeopleSpaceWidgetManager}. */
@@ -246,7 +270,7 @@
             CommonNotifCollection notifCollection, PackageManager packageManager,
             Optional<Bubbles> bubblesOptional, UserManager userManager, BackupManager backupManager,
             INotificationManager iNotificationManager, NotificationManager notificationManager,
-            @Background Executor executor) {
+            @Background Executor executor, UserTracker userTracker) {
         mContext = context;
         mAppWidgetManager = appWidgetManager;
         mIPeopleManager = iPeopleManager;
@@ -262,6 +286,7 @@
         mManager = this;
         mSharedPrefs = PreferenceManager.getDefaultSharedPreferences(context);
         mBgExecutor = executor;
+        mUserTracker = userTracker;
     }
 
     /**
@@ -1407,4 +1432,24 @@
 
         Trace.traceEnd(Trace.TRACE_TAG_APP);
     }
+
+    @VisibleForTesting
+    void updateGeneratedPreviewForUser(UserHandle user) {
+        if (!generatedPreviews() || mUpdatedPreviews.get(user.getIdentifier())
+                || !mUserManager.isUserUnlocked(user)) {
+            return;
+        }
+        if (DEBUG) {
+            Log.d(TAG, "Updating People Space widget preview for user " + user.getIdentifier());
+        }
+        boolean success = mAppWidgetManager.setWidgetPreview(
+                new ComponentName(mContext, PeopleSpaceWidgetProvider.class),
+                WIDGET_CATEGORY_HOME_SCREEN | WIDGET_CATEGORY_KEYGUARD,
+                new RemoteViews(mContext.getPackageName(),
+                        R.layout.people_space_placeholder_layout));
+        if (DEBUG && !success) {
+            Log.d(TAG, "Failed to update generated preview for user " + user.getIdentifier());
+        }
+        mUpdatedPreviews.put(user.getIdentifier(), success);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/TileSpec.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/TileSpec.kt
index aed08f8..4a8e33a 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/TileSpec.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/shared/TileSpec.kt
@@ -37,7 +37,11 @@
     data class PlatformTileSpec
     internal constructor(
         override val spec: String,
-    ) : TileSpec(spec)
+    ) : TileSpec(spec) {
+        override fun toString(): String {
+            return "P($spec)"
+        }
+    }
 
     /**
      * Container for the spec of a tile provided by an app.
@@ -50,7 +54,7 @@
         val componentName: ComponentName,
     ) : TileSpec(spec) {
         override fun toString(): String {
-            return "CustomTileSpec(${componentName.toShortString()})"
+            return "C(${componentName.flattenToShortString()})"
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java b/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java
index c4287ca..864f29a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ImageExporter.java
@@ -126,7 +126,7 @@
     /**
      * Writes the given Bitmap to outputFile.
      */
-    ListenableFuture<File> exportToRawFile(Executor executor, Bitmap bitmap,
+    public ListenableFuture<File> exportToRawFile(Executor executor, Bitmap bitmap,
             final File outputFile) {
         return CallbackToFutureAdapter.getFuture(
                 (completer) -> {
@@ -196,7 +196,7 @@
      * @param bitmap   the bitmap to export
      * @return a listenable future result
      */
-    ListenableFuture<Result> export(Executor executor, UUID requestId, Bitmap bitmap,
+    public ListenableFuture<Result> export(Executor executor, UUID requestId, Bitmap bitmap,
             ZonedDateTime captureTime, UserHandle owner, int displayId) {
         return export(executor, new Task(mResolver, requestId, bitmap, captureTime, mCompressFormat,
                 mQuality, /* publish */ true, owner, mFlags,
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotViewProxy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotViewProxy.kt
index 1f6d212..a1481f6 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotViewProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/LegacyScreenshotViewProxy.kt
@@ -34,6 +34,7 @@
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.res.R
+import com.android.systemui.screenshot.scroll.ScrollCaptureController
 import com.android.systemui.screenshot.LogConfig.DEBUG_DISMISS
 import com.android.systemui.screenshot.ScreenshotEvent.SCREENSHOT_DISMISSED_OTHER
 import dagger.assisted.Assisted
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LogConfig.java b/packages/SystemUI/src/com/android/systemui/screenshot/LogConfig.java
index 6050c2b..440cf1c 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LogConfig.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/LogConfig.java
@@ -16,8 +16,9 @@
 
 package com.android.systemui.screenshot;
 
+/** Stores debug log configuration for screenshots. */
 @SuppressWarnings("PointlessBooleanExpression")
-class LogConfig {
+public class LogConfig {
 
     /** Log ALL the things... */
     private static final boolean DEBUG_ALL = false;
@@ -29,36 +30,37 @@
     private static final boolean TAG_WITH_CLASS_NAME = false;
 
     /** Action creation and user selection: Share, Save, Edit, Delete, Smart action, etc */
-    static final boolean DEBUG_ACTIONS = DEBUG_ALL || false;
+    public static final boolean DEBUG_ACTIONS = DEBUG_ALL || false;
 
     /** Debug info about animations such as start, complete and cancel */
-    static final boolean DEBUG_ANIM = DEBUG_ALL || false;
+    public static final boolean DEBUG_ANIM = DEBUG_ALL || false;
 
     /** Whenever Uri is supplied to consumer, or onComplete runnable is run() */
-    static final boolean DEBUG_CALLBACK = DEBUG_ALL || false;
+    public static final boolean DEBUG_CALLBACK = DEBUG_ALL || false;
 
     /** Logs information about dismissing the screenshot tool */
-    static final boolean DEBUG_DISMISS = DEBUG_ALL || false;
+    public static final boolean DEBUG_DISMISS = DEBUG_ALL || false;
 
     /** Touch or key event driven action or side effects */
-    static final boolean DEBUG_INPUT = DEBUG_ALL || false;
+    public static final boolean DEBUG_INPUT = DEBUG_ALL || false;
 
     /** Scroll capture usage */
-    static final boolean DEBUG_SCROLL = DEBUG_ALL || false;
+    public static final boolean DEBUG_SCROLL = DEBUG_ALL || false;
 
     /** Service lifecycle events and callbacks */
-    static final boolean DEBUG_SERVICE = DEBUG_ALL || false;
+    public static final boolean DEBUG_SERVICE = DEBUG_ALL || false;
 
     /** Storage related actions, Bitmap.compress, ContentManager, etc */
-    static final boolean DEBUG_STORAGE = DEBUG_ALL || false;
+    public static final boolean DEBUG_STORAGE = DEBUG_ALL || false;
 
     /** High level logical UI actions: timeout, onConfigChanged, insets, show actions, reset  */
-    static final boolean DEBUG_UI = DEBUG_ALL || false;
+    public static final boolean DEBUG_UI = DEBUG_ALL || false;
 
     /** Interactions with Window and WindowManager */
-    static final boolean DEBUG_WINDOW = DEBUG_ALL || false;
+    public static final boolean DEBUG_WINDOW = DEBUG_ALL || false;
 
-    static String logTag(Class<?> cls) {
+    /** Get the appropriate class name */
+    public static String logTag(Class<?> cls) {
         return TAG_WITH_CLASS_NAME ? cls.getSimpleName() : TAG_SS;
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index 198a29c..c8e13bb 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -87,6 +87,10 @@
 import com.android.systemui.flags.Flags;
 import com.android.systemui.res.R;
 import com.android.systemui.screenshot.TakeScreenshotService.RequestCallback;
+import com.android.systemui.screenshot.scroll.LongScreenshotActivity;
+import com.android.systemui.screenshot.scroll.LongScreenshotData;
+import com.android.systemui.screenshot.scroll.ScrollCaptureClient;
+import com.android.systemui.screenshot.scroll.ScrollCaptureController;
 import com.android.systemui.util.Assert;
 
 import com.google.common.util.concurrent.ListenableFuture;
@@ -200,7 +204,7 @@
         void onActionsReady(ScreenshotController.QuickShareData quickShareData);
     }
 
-    interface TransitionDestination {
+    public interface TransitionDestination {
         /**
          * Allows the long screenshot activity to call back with a destination location (the bounds
          * on screen of the destination for the transitioning view) and a Runnable to be run once
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
index 1c5a8a1..cb2dba0 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotView.java
@@ -92,6 +92,7 @@
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
 import com.android.systemui.res.R;
+import com.android.systemui.screenshot.scroll.ScrollCaptureController;
 import com.android.systemui.shared.system.InputChannelCompat;
 import com.android.systemui.shared.system.InputMonitorCompat;
 import com.android.systemui.shared.system.QuickStepContract;
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotViewProxy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotViewProxy.kt
index 182b889..6be32a9 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotViewProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotViewProxy.kt
@@ -25,6 +25,7 @@
 import android.view.View
 import android.view.ViewGroup
 import android.view.WindowInsets
+import com.android.systemui.screenshot.scroll.ScrollCaptureController
 
 /** Abstraction of the surface between ScreenshotController and ScreenshotView */
 interface ScreenshotViewProxy {
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
index aa23d6b..d87d85b 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
@@ -52,7 +52,7 @@
 import com.android.internal.logging.UiEventLogger.UiEventEnum;
 import com.android.settingslib.Utils;
 import com.android.systemui.res.R;
-import com.android.systemui.screenshot.CropView;
+import com.android.systemui.screenshot.scroll.CropView;
 import com.android.systemui.settings.UserTracker;
 
 import javax.inject.Inject;
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/CropView.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/CropView.java
similarity index 98%
rename from packages/SystemUI/src/com/android/systemui/screenshot/CropView.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/CropView.java
index 2f411ea..5e561cf 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/CropView.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/CropView.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.animation.ValueAnimator;
 import android.content.Context;
@@ -265,7 +265,7 @@
                 Log.w(TAG, "No boundary selected");
                 break;
         }
-        Log.i(TAG,  "Updated mCrop: " + mCrop);
+        Log.i(TAG, "Updated mCrop: " + mCrop);
 
         invalidate();
     }
@@ -385,7 +385,7 @@
 
     /**
      * @param action either ACTION_DOWN, ACTION_UP or ACTION_MOVE.
-     * @param x coordinate of the relevant pointer.
+     * @param x      x-coordinate of the relevant pointer.
      */
     private void updateListener(int action, float x) {
         if (mCropInteractionListener != null && isVertical(mCurrentDraggingBoundary)) {
@@ -643,11 +643,13 @@
     /**
      * Listen for crop motion events and state.
      */
-    public interface CropInteractionListener {
+    interface CropInteractionListener {
         void onCropDragStarted(CropBoundary boundary, float boundaryPosition,
                 int boundaryPositionPx, float horizontalCenter, float x);
+
         void onCropDragMoved(CropBoundary boundary, float boundaryPosition,
                 int boundaryPositionPx, float horizontalCenter, float x);
+
         void onCropDragComplete();
     }
 
@@ -675,8 +677,7 @@
             out.writeParcelable(mCrop, 0);
         }
 
-        public static final Parcelable.Creator<SavedState> CREATOR
-                = new Parcelable.Creator<SavedState>() {
+        public static final Parcelable.Creator<SavedState> CREATOR = new Parcelable.Creator<>() {
             public SavedState createFromParcel(Parcel in) {
                 return new SavedState(in);
             }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageLoader.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageLoader.java
similarity index 81%
rename from packages/SystemUI/src/com/android/systemui/screenshot/ImageLoader.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageLoader.java
index 7ee7c31..df86d69 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageLoader.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageLoader.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.annotation.Nullable;
 import android.content.ContentResolver;
@@ -35,20 +35,20 @@
 import javax.inject.Inject;
 
 /** Loads images. */
-public class ImageLoader {
+class ImageLoader {
     private final ContentResolver mResolver;
 
     static class Result {
-        @Nullable Uri uri;
-        @Nullable File fileName;
-        @Nullable Bitmap bitmap;
+        @Nullable Uri mUri;
+        @Nullable File mFilename;
+        @Nullable Bitmap mBitmap;
 
         @Override
         public String toString() {
             final StringBuilder sb = new StringBuilder("Result{");
-            sb.append("uri=").append(uri);
-            sb.append(", fileName=").append(fileName);
-            sb.append(", bitmap=").append(bitmap);
+            sb.append("uri=").append(mUri);
+            sb.append(", fileName=").append(mFilename);
+            sb.append(", bitmap=").append(mBitmap);
             sb.append('}');
             return sb.toString();
         }
@@ -69,11 +69,10 @@
         return CallbackToFutureAdapter.getFuture(completer -> {
             Result result = new Result();
             try (InputStream in = mResolver.openInputStream(uri)) {
-                result.uri = uri;
-                result.bitmap = BitmapFactory.decodeStream(in);
+                result.mUri = uri;
+                result.mBitmap = BitmapFactory.decodeStream(in);
                 completer.set(result);
-            }
-            catch (IOException e) {
+            } catch (IOException e) {
                 completer.setException(e);
             }
             return "BitmapFactory#decodeStream";
@@ -91,8 +90,8 @@
         return CallbackToFutureAdapter.getFuture(completer -> {
             try (InputStream in = new BufferedInputStream(new FileInputStream(file))) {
                 Result result = new Result();
-                result.fileName = file;
-                result.bitmap = BitmapFactory.decodeStream(in);
+                result.mFilename = file;
+                result.mBitmap = BitmapFactory.decodeStream(in);
                 completer.set(result);
             } catch (IOException e) {
                 completer.setException(e);
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageTile.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageTile.java
similarity index 96%
rename from packages/SystemUI/src/com/android/systemui/screenshot/ImageTile.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageTile.java
index a95c91b..c9c297e 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageTile.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageTile.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import static android.graphics.ColorSpace.Named.SRGB;
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ImageTileSet.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageTileSet.java
similarity index 97%
rename from packages/SystemUI/src/com/android/systemui/screenshot/ImageTileSet.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageTileSet.java
index 356f67e..76a72f7 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ImageTileSet.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ImageTileSet.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.annotation.AnyThread;
 import android.graphics.Bitmap;
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotActivity.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotActivity.java
similarity index 96%
rename from packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotActivity.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotActivity.java
index 00d480a..1e1a577 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotActivity.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.app.Activity;
 import android.app.ActivityOptions;
@@ -47,11 +47,14 @@
 import com.android.internal.view.OneShotPreDrawListener;
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
-import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.res.R;
-import com.android.systemui.screenshot.CropView.CropBoundary;
-import com.android.systemui.screenshot.ScrollCaptureController.LongScreenshot;
-import com.android.systemui.settings.UserTracker;
+import com.android.systemui.screenshot.ActionIntentCreator;
+import com.android.systemui.screenshot.ActionIntentExecutor;
+import com.android.systemui.screenshot.ImageExporter;
+import com.android.systemui.screenshot.LogConfig;
+import com.android.systemui.screenshot.ScreenshotEvent;
+import com.android.systemui.screenshot.scroll.CropView.CropBoundary;
+import com.android.systemui.screenshot.scroll.ScrollCaptureController.LongScreenshot;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
@@ -81,8 +84,6 @@
     private final ImageExporter mImageExporter;
     private final LongScreenshotData mLongScreenshotHolder;
     private final ActionIntentExecutor mActionExecutor;
-    private final FeatureFlags mFeatureFlags;
-    private final UserTracker mUserTracker;
 
     private ImageView mPreview;
     private ImageView mTransitionView;
@@ -113,16 +114,13 @@
     @Inject
     public LongScreenshotActivity(UiEventLogger uiEventLogger, ImageExporter imageExporter,
             @Main Executor mainExecutor, @Background Executor bgExecutor,
-            LongScreenshotData longScreenshotHolder, ActionIntentExecutor actionExecutor,
-            FeatureFlags featureFlags, UserTracker userTracker) {
+            LongScreenshotData longScreenshotHolder, ActionIntentExecutor actionExecutor) {
         mUiEventLogger = uiEventLogger;
         mUiExecutor = mainExecutor;
         mBackgroundExecutor = bgExecutor;
         mImageExporter = imageExporter;
         mLongScreenshotHolder = longScreenshotHolder;
         mActionExecutor = actionExecutor;
-        mFeatureFlags = featureFlags;
-        mUserTracker = userTracker;
     }
 
 
@@ -265,13 +263,13 @@
     private void onCachedImageLoaded(ImageLoader.Result imageResult) {
         mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_ACTIVITY_CACHED_IMAGE_LOADED);
 
-        BitmapDrawable drawable = new BitmapDrawable(getResources(), imageResult.bitmap);
+        BitmapDrawable drawable = new BitmapDrawable(getResources(), imageResult.mBitmap);
         mPreview.setImageDrawable(drawable);
         mPreview.setAlpha(1f);
-        mMagnifierView.setDrawable(drawable, imageResult.bitmap.getWidth(),
-                imageResult.bitmap.getHeight());
+        mMagnifierView.setDrawable(drawable, imageResult.mBitmap.getWidth(),
+                imageResult.mBitmap.getHeight());
         mCropView.setVisibility(View.VISIBLE);
-        mSavedImagePath = imageResult.fileName;
+        mSavedImagePath = imageResult.mFilename;
 
         setButtonsEnabled(true);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotData.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotData.java
similarity index 93%
rename from packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotData.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotData.java
index f549faf..ebac5bf 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/LongScreenshotData.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/LongScreenshotData.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.screenshot.ScreenshotController;
 
 import java.util.concurrent.atomic.AtomicReference;
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/MagnifierView.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/MagnifierView.java
similarity index 94%
rename from packages/SystemUI/src/com/android/systemui/screenshot/MagnifierView.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/MagnifierView.java
index 0c543cd..0a1a747 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/MagnifierView.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/MagnifierView.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -64,16 +64,16 @@
     private ViewPropertyAnimator mTranslationAnimator;
     private final Animator.AnimatorListener mTranslationAnimatorListener =
             new AnimatorListenerAdapter() {
-        @Override
-        public void onAnimationCancel(Animator animation) {
-            mTranslationAnimator = null;
-        }
+                @Override
+                public void onAnimationCancel(Animator animation) {
+                    mTranslationAnimator = null;
+                }
 
-        @Override
-        public void onAnimationEnd(Animator animation) {
-            mTranslationAnimator = null;
-        }
-    };
+                @Override
+                public void onAnimationEnd(Animator animation) {
+                    mTranslationAnimator = null;
+                }
+            };
 
     public MagnifierView(Context context, @Nullable AttributeSet attrs) {
         this(context, attrs, 0);
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureClient.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ScrollCaptureClient.java
similarity index 98%
rename from packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureClient.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/ScrollCaptureClient.java
index e93f737..0e43343 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureClient.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ScrollCaptureClient.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import static com.android.systemui.screenshot.LogConfig.DEBUG_SCROLL;
 
@@ -46,6 +46,7 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.systemui.dagger.qualifiers.Background;
+import com.android.systemui.screenshot.LogConfig;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureController.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ScrollCaptureController.java
similarity index 96%
rename from packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureController.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/ScrollCaptureController.java
index 8a2678c..f4c77da 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScrollCaptureController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/ScrollCaptureController.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.content.Context;
 import android.graphics.Bitmap;
@@ -30,8 +30,10 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.UiEventLogger;
 import com.android.systemui.dagger.qualifiers.Background;
-import com.android.systemui.screenshot.ScrollCaptureClient.CaptureResult;
-import com.android.systemui.screenshot.ScrollCaptureClient.Session;
+import com.android.systemui.screenshot.LogConfig;
+import com.android.systemui.screenshot.ScreenshotEvent;
+import com.android.systemui.screenshot.scroll.ScrollCaptureClient.CaptureResult;
+import com.android.systemui.screenshot.scroll.ScrollCaptureClient.Session;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
@@ -85,7 +87,7 @@
             mImageTileSet = imageTileSet;
         }
 
-        /** Returns a bitmap containing the combinded result. */
+        /** Returns a bitmap containing the combined result. */
         public Bitmap toBitmap() {
             return mImageTileSet.toBitmap();
         }
@@ -167,7 +169,7 @@
      *                 {@link ScrollCaptureResponse#isConnected() connected}.
      * @return a future ImageTile set containing the result
      */
-    ListenableFuture<LongScreenshot> run(ScrollCaptureResponse response) {
+    public ListenableFuture<LongScreenshot> run(ScrollCaptureResponse response) {
         mCancelled = false;
         return CallbackToFutureAdapter.getFuture(completer -> {
             mCaptureCompleter = completer;
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TiledImageDrawable.java b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/TiledImageDrawable.java
similarity index 97%
rename from packages/SystemUI/src/com/android/systemui/screenshot/TiledImageDrawable.java
rename to packages/SystemUI/src/com/android/systemui/screenshot/scroll/TiledImageDrawable.java
index 71df369..00455bc 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TiledImageDrawable.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/scroll/TiledImageDrawable.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.annotation.Nullable;
 import android.graphics.Canvas;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java b/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
index 4275fc6..4406813 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
@@ -178,6 +178,7 @@
     private static final int MSG_IMMERSIVE_CHANGED = 78 << MSG_SHIFT;
     private static final int MSG_SET_QS_TILES = 79 << MSG_SHIFT;
     private static final int MSG_ENTER_DESKTOP = 80 << MSG_SHIFT;
+    private static final int MSG_SET_SPLITSCREEN_FOCUS = 81 << MSG_SHIFT;
     public static final int FLAG_EXCLUDE_NONE = 0;
     public static final int FLAG_EXCLUDE_SEARCH_PANEL = 1 << 0;
     public static final int FLAG_EXCLUDE_RECENTS_PANEL = 1 << 1;
@@ -508,6 +509,11 @@
         default void moveFocusedTaskToStageSplit(int displayId, boolean leftOrTop) {}
 
         /**
+         * @see IStatusBar#setSplitscreenFocus
+         */
+        default void setSplitscreenFocus(boolean leftOrTop) {}
+
+        /**
          * @see IStatusBar#showMediaOutputSwitcher
          */
         default void showMediaOutputSwitcher(String packageName) {}
@@ -1349,6 +1355,12 @@
     }
 
     @Override
+    public void setSplitscreenFocus(boolean leftOrTop) {
+        synchronized (mLock) {
+            mHandler.obtainMessage(MSG_SET_SPLITSCREEN_FOCUS, leftOrTop).sendToTarget();
+        }
+    }
+    @Override
     public void showMediaOutputSwitcher(String packageName) {
         int callingUid = Binder.getCallingUid();
         if (callingUid != 0 && callingUid != Process.SYSTEM_UID) {
@@ -1919,6 +1931,11 @@
                     }
                     break;
                 }
+                case MSG_SET_SPLITSCREEN_FOCUS:
+                    for (int i = 0; i < mCallbacks.size(); i++) {
+                        mCallbacks.get(i).setSplitscreenFocus((Boolean) msg.obj);
+                    }
+                    break;
                 case MSG_SHOW_MEDIA_OUTPUT_SWITCHER:
                     args = (SomeArgs) msg.obj;
                     String clientPackageName = (String) args.arg1;
diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
index e18bc04..c294bd8 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
@@ -260,6 +260,11 @@
             public void moveFocusedTaskToFullscreen(int displayId) {
                 splitScreen.goToFullscreenFromSplit();
             }
+
+            @Override
+            public void setSplitscreenFocus(boolean leftOrTop) {
+                splitScreen.setSplitscreenFocus(leftOrTop);
+            }
         });
         splitScreen.registerSplitAnimationListener(new SplitScreen.SplitInvocationListener() {
             @Override
diff --git a/packages/SystemUI/tests/AndroidManifest.xml b/packages/SystemUI/tests/AndroidManifest.xml
index 5882b56..572a6c1 100644
--- a/packages/SystemUI/tests/AndroidManifest.xml
+++ b/packages/SystemUI/tests/AndroidManifest.xml
@@ -113,7 +113,7 @@
             android:excludeFromRecents="true"
             />
 
-        <activity android:name="com.android.systemui.screenshot.ScrollViewActivity"
+        <activity android:name="com.android.systemui.screenshot.scroll.ScrollViewActivity"
                   android:exported="false" />
 
         <activity android:name="com.android.systemui.screenshot.RecyclerViewActivity"
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewControllerBaseTest.java
index 90587d7..f1dfdf4 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardStatusViewControllerBaseTest.java
@@ -20,16 +20,18 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import static kotlinx.coroutines.flow.FlowKt.emptyFlow;
+
 import android.view.View;
 import android.view.ViewTreeObserver;
 import android.widget.FrameLayout;
 
-import com.android.internal.jank.InteractionJankMonitor;
 import com.android.keyguard.logging.KeyguardLogger;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractorFactory;
+import com.android.systemui.kosmos.KosmosJavaAdapter;
 import com.android.systemui.power.data.repository.FakePowerRepository;
 import com.android.systemui.power.domain.interactor.PowerInteractorFactory;
 import com.android.systemui.res.R;
@@ -46,6 +48,7 @@
 
 public class KeyguardStatusViewControllerBaseTest extends SysuiTestCase {
 
+    private KosmosJavaAdapter mKosmos;
     @Mock protected KeyguardStatusView mKeyguardStatusView;
 
     @Mock protected KeyguardSliceViewController mKeyguardSliceViewController;
@@ -57,7 +60,6 @@
     @Mock protected ScreenOffAnimationController mScreenOffAnimationController;
     @Mock protected KeyguardLogger mKeyguardLogger;
     @Mock protected KeyguardStatusViewController mControllerMock;
-    @Mock protected InteractionJankMonitor mInteractionJankMonitor;
     @Mock protected ViewTreeObserver mViewTreeObserver;
     @Mock protected DumpManager mDumpManager;
     protected FakeKeyguardRepository mFakeKeyguardRepository;
@@ -71,6 +73,7 @@
 
     @Before
     public void setup() {
+        mKosmos = new KosmosJavaAdapter(this);
         MockitoAnnotations.initMocks(this);
 
         KeyguardInteractorFactory.WithDependencies deps = KeyguardInteractorFactory.create();
@@ -87,7 +90,7 @@
                 mDozeParameters,
                 mScreenOffAnimationController,
                 mKeyguardLogger,
-                mInteractionJankMonitor,
+                mKosmos.getInteractionJankMonitor(),
                 deps.getKeyguardInteractor(),
                 mDumpManager,
                 PowerInteractorFactory.create(
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
index 336a97e..fde45d3 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
@@ -135,6 +135,7 @@
 import com.android.systemui.deviceentry.shared.model.FaceDetectionStatus;
 import com.android.systemui.deviceentry.shared.model.FailedFaceAuthenticationStatus;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.kosmos.KosmosJavaAdapter;
 import com.android.systemui.log.SessionTracker;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.settings.UserTracker;
@@ -195,7 +196,7 @@
             DATA_ROAMING_DISABLE, null, null, null, null, false, null, "", false, TEST_GROUP_UUID,
             TEST_CARRIER_ID, PROFILE_CLASS_PROVISIONING);
     private static final int FINGERPRINT_SENSOR_ID = 1;
-
+    private KosmosJavaAdapter mKosmos;
     @Mock
     private UserTracker mUserTracker;
     @Mock
@@ -240,7 +241,6 @@
     private AuthController mAuthController;
     @Mock
     private TelephonyListenerManager mTelephonyListenerManager;
-    @Mock
     private InteractionJankMonitor mInteractionJankMonitor;
     @Mock
     private LatencyTracker mLatencyTracker;
@@ -300,6 +300,8 @@
 
     @Before
     public void setup() throws RemoteException {
+        mKosmos = new KosmosJavaAdapter(this);
+        mInteractionJankMonitor = mKosmos.getInteractionJankMonitor();
         MockitoAnnotations.initMocks(this);
         when(mSessionTracker.getSessionId(SESSION_KEYGUARD)).thenReturn(mKeyguardInstanceId);
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/animation/DialogTransitionAnimatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/animation/DialogTransitionAnimatorTest.kt
index 96ce3ab..b73e4e6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/animation/DialogTransitionAnimatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/animation/DialogTransitionAnimatorTest.kt
@@ -18,6 +18,8 @@
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.internal.policy.DecorView
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.jank.interactionJankMonitor
+import com.android.systemui.kosmos.Kosmos
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert.assertEquals
 import junit.framework.Assert.assertFalse
@@ -31,7 +33,6 @@
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
 import org.mockito.Mockito.any
 import org.mockito.Mockito.verify
 import org.mockito.junit.MockitoJUnit
@@ -43,7 +44,7 @@
     private lateinit var mDialogTransitionAnimator: DialogTransitionAnimator
     private val attachedViews = mutableSetOf<View>()
 
-    @Mock lateinit var interactionJankMonitor: InteractionJankMonitor
+    val interactionJankMonitor = Kosmos().interactionJankMonitor
     @get:Rule val rule = MockitoJUnit.rule()
 
     @Before
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
index 1849245..272b488 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
@@ -75,7 +75,6 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.foldables.FoldGracePeriodProvider;
-import com.android.internal.jank.InteractionJankMonitor;
 import com.android.internal.logging.InstanceId;
 import com.android.internal.logging.UiEventLogger;
 import com.android.internal.widget.LockPatternUtils;
@@ -181,7 +180,6 @@
     private @Mock NotificationShadeDepthController mNotificationShadeDepthController;
     private @Mock KeyguardUnlockAnimationController mKeyguardUnlockAnimationController;
     private @Mock ScreenOffAnimationController mScreenOffAnimationController;
-    private @Mock InteractionJankMonitor mInteractionJankMonitor;
     private @Mock ScreenOnCoordinator mScreenOnCoordinator;
     private @Mock KeyguardTransitions mKeyguardTransitions;
     private @Mock ShadeController mShadeController;
@@ -235,8 +233,6 @@
         when(mLockPatternUtils.getDevicePolicyManager()).thenReturn(mDevicePolicyManager);
         when(mPowerManager.newWakeLock(anyInt(), any())).thenReturn(mock(WakeLock.class));
         when(mPowerManager.isInteractive()).thenReturn(true);
-        when(mInteractionJankMonitor.begin(any(), anyInt())).thenReturn(true);
-        when(mInteractionJankMonitor.end(anyInt())).thenReturn(true);
         mContext.addMockSystemService(Context.ALARM_SERVICE, mAlarmManager);
         final ViewRootImpl testViewRoot = mock(ViewRootImpl.class);
         when(testViewRoot.getView()).thenReturn(mock(View.class));
@@ -1245,7 +1241,7 @@
                 () -> mNotificationShadeDepthController,
                 mScreenOnCoordinator,
                 mKeyguardTransitions,
-                mInteractionJankMonitor,
+                mKosmos.getInteractionJankMonitor(),
                 mDreamOverlayStateController,
                 mJavaAdapter,
                 mWallpaperRepository,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/people/widget/PeopleSpaceWidgetManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/people/widget/PeopleSpaceWidgetManagerTest.java
index a63b221..d1d9efc 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/people/widget/PeopleSpaceWidgetManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/people/widget/PeopleSpaceWidgetManagerTest.java
@@ -101,11 +101,12 @@
 import androidx.preference.PreferenceManager;
 import androidx.test.filters.SmallTest;
 
-import com.android.systemui.res.R;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.people.PeopleBackupFollowUpJob;
 import com.android.systemui.people.PeopleSpaceUtils;
 import com.android.systemui.people.SharedPreferencesHelper;
+import com.android.systemui.res.R;
+import com.android.systemui.settings.FakeUserTracker;
 import com.android.systemui.statusbar.NotificationListener;
 import com.android.systemui.statusbar.NotificationListener.NotificationHandler;
 import com.android.systemui.statusbar.SbnBuilder;
@@ -265,6 +266,8 @@
 
     private final FakeExecutor mFakeExecutor = new FakeExecutor(mClock);
 
+    private final FakeUserTracker mUserTracker = new FakeUserTracker();
+
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
@@ -272,7 +275,7 @@
         mManager = new PeopleSpaceWidgetManager(mContext, mAppWidgetManager, mIPeopleManager,
                 mPeopleManager, mLauncherApps, mNotifCollection, mPackageManager,
                 Optional.of(mBubbles), mUserManager, mBackupManager, mINotificationManager,
-                mNotificationManager, mFakeExecutor);
+                mNotificationManager, mFakeExecutor, mUserTracker);
         mManager.attach(mListenerService);
 
         verify(mListenerService).addNotificationHandler(mListenerCaptor.capture());
@@ -1562,6 +1565,43 @@
                 String.valueOf(WIDGET_ID_WITH_KEY_IN_OPTIONS));
     }
 
+    @Test
+    public void testUpdateGeneratedPreview_flagDisabled() {
+        mSetFlagsRule.disableFlags(android.appwidget.flags.Flags.FLAG_GENERATED_PREVIEWS);
+        mManager.updateGeneratedPreviewForUser(mUserTracker.getUserHandle());
+        verify(mAppWidgetManager, times(0)).setWidgetPreview(any(), anyInt(), any());
+    }
+
+    @Test
+    public void testUpdateGeneratedPreview_userLocked() {
+        mSetFlagsRule.enableFlags(android.appwidget.flags.Flags.FLAG_GENERATED_PREVIEWS);
+        when(mUserManager.isUserUnlocked(mUserTracker.getUserHandle())).thenReturn(false);
+
+        mManager.updateGeneratedPreviewForUser(mUserTracker.getUserHandle());
+        verify(mAppWidgetManager, times(0)).setWidgetPreview(any(), anyInt(), any());
+    }
+
+    @Test
+    public void testUpdateGeneratedPreview_userUnlocked() {
+        mSetFlagsRule.enableFlags(android.appwidget.flags.Flags.FLAG_GENERATED_PREVIEWS);
+        when(mUserManager.isUserUnlocked(mUserTracker.getUserHandle())).thenReturn(true);
+        when(mAppWidgetManager.setWidgetPreview(any(), anyInt(), any())).thenReturn(true);
+
+        mManager.updateGeneratedPreviewForUser(mUserTracker.getUserHandle());
+        verify(mAppWidgetManager, times(1)).setWidgetPreview(any(), anyInt(), any());
+    }
+
+    @Test
+    public void testUpdateGeneratedPreview_doesNotSetTwice() {
+        mSetFlagsRule.enableFlags(android.appwidget.flags.Flags.FLAG_GENERATED_PREVIEWS);
+        when(mUserManager.isUserUnlocked(mUserTracker.getUserHandle())).thenReturn(true);
+        when(mAppWidgetManager.setWidgetPreview(any(), anyInt(), any())).thenReturn(true);
+
+        mManager.updateGeneratedPreviewForUser(mUserTracker.getUserHandle());
+        mManager.updateGeneratedPreviewForUser(mUserTracker.getUserHandle());
+        verify(mAppWidgetManager, times(1)).setWidgetPreview(any(), anyInt(), any());
+    }
+
     private void setFinalField(String fieldName, int value) {
         try {
             Field field = NotificationManager.Policy.class.getDeclaredField(fieldName);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeSessionTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/FakeSessionTest.java
similarity index 98%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeSessionTest.java
rename to packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/FakeSessionTest.java
index 4c8a4b0..aad46139 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/FakeSessionTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/FakeSessionTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import static com.google.common.util.concurrent.Futures.getUnchecked;
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureClientTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureClientTest.java
similarity index 93%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureClientTest.java
rename to packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureClientTest.java
index 670a130..1023260 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureClientTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureClientTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 
 import static org.junit.Assert.assertEquals;
@@ -37,8 +37,8 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.screenshot.ScrollCaptureClient.CaptureResult;
-import com.android.systemui.screenshot.ScrollCaptureClient.Session;
+import com.android.systemui.screenshot.scroll.ScrollCaptureClient.CaptureResult;
+import com.android.systemui.screenshot.scroll.ScrollCaptureClient.Session;
 
 import com.google.common.util.concurrent.ListenableFuture;
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureControllerTest.java
similarity index 98%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureControllerTest.java
rename to packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureControllerTest.java
index 6f081c7..f39f543 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureControllerTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import static com.google.common.util.concurrent.Futures.getUnchecked;
 import static com.google.common.util.concurrent.Futures.immediateFuture;
@@ -36,7 +36,7 @@
 
 import com.android.internal.logging.testing.UiEventLoggerFake;
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.screenshot.ScrollCaptureClient.Session;
+import com.android.systemui.screenshot.scroll.ScrollCaptureClient.Session;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureFrameworkSmokeTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureFrameworkSmokeTest.java
similarity index 96%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureFrameworkSmokeTest.java
rename to packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureFrameworkSmokeTest.java
index de97bc3..5699cfc 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollCaptureFrameworkSmokeTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollCaptureFrameworkSmokeTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollViewActivity.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollViewActivity.java
similarity index 93%
rename from packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollViewActivity.java
rename to packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollViewActivity.java
index 4c84df2..04aba11 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/ScrollViewActivity.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/scroll/ScrollViewActivity.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.app.Activity;
 import android.os.Bundle;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
index c31c625..1ee26db 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationShadeWindowControllerImplTest.java
@@ -196,7 +196,8 @@
                 new ConfigurationInteractor(configurationRepository),
                 shadeRepository,
                 keyguardTransitionInteractor,
-                () -> sceneInteractor);
+                () -> sceneInteractor,
+                () -> mKosmos.getFromGoneTransitionInteractor());
         CommunalInteractor communalInteractor = mKosmos.getCommunalInteractor();
 
         mFromLockscreenTransitionInteractor = mKosmos.getFromLockscreenTransitionInteractor();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
index a077164..b9451ba 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/QuickSettingsControllerImplBaseTest.java
@@ -17,7 +17,6 @@
 package com.android.systemui.shade;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -223,7 +222,8 @@
                 new ConfigurationInteractor(configurationRepository),
                 mShadeRepository,
                 keyguardTransitionInteractor,
-                () -> sceneInteractor);
+                () -> sceneInteractor,
+                () -> mKosmos.getFromGoneTransitionInteractor());
 
         mFromLockscreenTransitionInteractor = mKosmos.getFromLockscreenTransitionInteractor();
         mFromPrimaryBouncerTransitionInteractor =
@@ -289,10 +289,6 @@
 
         when(mNotificationRemoteInputManager.isRemoteInputActive())
                 .thenReturn(false);
-        when(mInteractionJankMonitor.begin(any(), anyInt()))
-                .thenReturn(true);
-        when(mInteractionJankMonitor.end(anyInt()))
-                .thenReturn(true);
 
         when(mPanelView.getParent()).thenReturn(mPanelViewParent);
         when(mQs.getHeader()).thenReturn(mQsHeader);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt
index 103dcb7..dcd000a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarStateControllerImplTest.kt
@@ -45,6 +45,7 @@
 import com.android.systemui.keyguard.domain.interactor.FromLockscreenTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.FromPrimaryBouncerTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
+import com.android.systemui.keyguard.domain.interactor.fromGoneTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.fromLockscreenTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.fromPrimaryBouncerTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
@@ -153,6 +154,7 @@
                 shadeRepository,
                 keyguardTransitionInteractor,
                 { kosmos.sceneInteractor },
+                { kosmos.fromGoneTransitionInteractor },
             )
 
         whenever(deviceEntryUdfpsInteractor.isUdfpsSupported).thenReturn(MutableStateFlow(false))
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
index 1748cff..d9e9c59 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/KeyguardStatusBarViewControllerTest.java
@@ -177,7 +177,8 @@
                 new ConfigurationInteractor(new FakeConfigurationRepository()),
                 new FakeShadeRepository(),
                 keyguardTransitionInteractor,
-                () -> mKosmos.getSceneInteractor());
+                () -> mKosmos.getSceneInteractor(),
+                () -> mKosmos.getFromGoneTransitionInteractor());
         mViewModel =
                 new KeyguardStatusBarViewModel(
                         mTestScope.getBackgroundScope(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
index 76913e8..e4b9f10 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
+import com.android.systemui.keyguard.domain.interactor.fromGoneTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.kosmos.testScope
@@ -63,9 +64,9 @@
             ConfigurationInteractor(FakeConfigurationRepository()),
             FakeShadeRepository(),
             kosmos.keyguardTransitionInteractor,
-        ) {
-            kosmos.sceneInteractor
-        }
+            { kosmos.sceneInteractor },
+            { kosmos.fromGoneTransitionInteractor },
+        )
     private val keyguardStatusBarInteractor =
         KeyguardStatusBarInteractor(
             FakeKeyguardStatusBarRepository(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index 19f31d5..ec27f48 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -441,7 +441,8 @@
                 new ConfigurationInteractor(configurationRepository),
                 shadeRepository,
                 keyguardTransitionInteractor,
-                () -> sceneInteractor);
+                () -> sceneInteractor,
+                () -> mKosmos.getFromGoneTransitionInteractor());
 
         mFromLockscreenTransitionInteractor = mKosmos.getFromLockscreenTransitionInteractor();
         mFromPrimaryBouncerTransitionInteractor =
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/jank/InteractionJankMonitorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/jank/InteractionJankMonitorKosmos.kt
index 5c5016d..e2b5869 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/jank/InteractionJankMonitorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/jank/InteractionJankMonitorKosmos.kt
@@ -16,9 +16,24 @@
 
 package com.android.systemui.jank
 
+import android.os.HandlerThread
 import com.android.internal.jank.InteractionJankMonitor
+import com.android.internal.jank.InteractionJankMonitor.Configuration.Builder
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.any
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Mockito.doReturn
+import org.mockito.Mockito.spy
 
-val Kosmos.interactionJankMonitor by Fixture<InteractionJankMonitor> { mock() }
+val Kosmos.interactionJankMonitor by
+    Fixture<InteractionJankMonitor> {
+        spy(InteractionJankMonitor(HandlerThread("InteractionJankMonitor-Kosmos"))).apply {
+            doReturn(true).`when`(this).shouldMonitor()
+            doReturn(true).`when`(this).begin(any(), anyInt())
+            doReturn(true).`when`(this).begin(any<Builder>())
+            doReturn(true).`when`(this).end(anyInt())
+            doReturn(true).`when`(this).cancel(anyInt())
+            doReturn(true).`when`(this).cancel(anyInt(), anyInt())
+        }
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
index 3893a9b7..00cdc33 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorFactory.kt
@@ -51,6 +51,7 @@
         configurationRepository: FakeConfigurationRepository = FakeConfigurationRepository(),
         shadeRepository: FakeShadeRepository = FakeShadeRepository(),
         sceneInteractor: SceneInteractor = mock(),
+        fromGoneTransitionInteractor: FromGoneTransitionInteractor = mock(),
         powerInteractor: PowerInteractor = PowerInteractorFactory.create().powerInteractor,
     ): WithDependencies {
         // Mock this until the class is replaced by kosmos
@@ -77,6 +78,7 @@
                 sceneInteractorProvider = { sceneInteractor },
                 keyguardTransitionInteractor = keyguardTransitionInteractor,
                 powerInteractor = powerInteractor,
+                fromGoneTransitionInteractor = { fromGoneTransitionInteractor },
             ),
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
index 5140a9f..d61bc9f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardInteractorKosmos.kt
@@ -26,7 +26,7 @@
 import com.android.systemui.shade.data.repository.shadeRepository
 import com.android.systemui.statusbar.commandQueue
 
-val Kosmos.keyguardInteractor by
+val Kosmos.keyguardInteractor: KeyguardInteractor by
     Kosmos.Fixture {
         KeyguardInteractor(
             repository = keyguardRepository,
@@ -38,5 +38,6 @@
             shadeRepository = shadeRepository,
             keyguardTransitionInteractor = keyguardTransitionInteractor,
             sceneInteractorProvider = { sceneInteractor },
+            fromGoneTransitionInteractor = { fromGoneTransitionInteractor },
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
index 3fc5af1..e861892 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
@@ -34,6 +34,7 @@
 import com.android.systemui.jank.interactionJankMonitor
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.domain.interactor.fromGoneTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.fromLockscreenTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.fromPrimaryBouncerTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
@@ -92,6 +93,7 @@
     val fromPrimaryBouncerTransitionInteractor by lazy {
         kosmos.fromPrimaryBouncerTransitionInteractor
     }
+    val fromGoneTransitionInteractor by lazy { kosmos.fromGoneTransitionInteractor }
     val globalActionsInteractor by lazy { kosmos.globalActionsInteractor }
     val sceneDataSource by lazy { kosmos.sceneDataSource }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/FakeScrollCaptureConnection.java b/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/scroll/FakeScrollCaptureConnection.java
similarity index 97%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/FakeScrollCaptureConnection.java
rename to packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/scroll/FakeScrollCaptureConnection.java
index 63f7c97..ea59c0a 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/FakeScrollCaptureConnection.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/scroll/FakeScrollCaptureConnection.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import android.content.pm.ActivityInfo;
 import android.graphics.Canvas;
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/FakeSession.java b/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/scroll/FakeSession.java
similarity index 97%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/FakeSession.java
rename to packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/scroll/FakeSession.java
index 478658e..3b7b158 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/FakeSession.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/screenshot/scroll/FakeSession.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2021 The Android Open Source Project
+ * 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.screenshot;
+package com.android.systemui.screenshot.scroll;
 
 import static android.util.MathUtils.constrain;
 
@@ -32,6 +32,8 @@
 import android.media.Image;
 import android.util.Log;
 
+import com.android.systemui.screenshot.scroll.ScrollCaptureClient;
+
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 
diff --git a/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java b/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
index f2409fb..5e52e06 100644
--- a/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
+++ b/services/companion/java/com/android/server/companion/BackupRestoreProcessor.java
@@ -18,7 +18,8 @@
 
 import static android.os.UserHandle.getCallingUserId;
 
-import static com.android.server.companion.CompanionDeviceManagerService.PerUserAssociationSet;
+import static com.android.server.companion.association.AssociationDiskStore.readAssociationsFromPayload;
+import static com.android.server.companion.utils.RolesUtils.addRoleHolderForAssociation;
 
 import android.annotation.NonNull;
 import android.annotation.SuppressLint;
@@ -26,62 +27,50 @@
 import android.companion.AssociationInfo;
 import android.companion.Flags;
 import android.companion.datatransfer.SystemDataTransferRequest;
+import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManagerInternal;
-import android.util.ArraySet;
-import android.util.Log;
 import android.util.Slog;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.CollectionUtils;
 import com.android.server.companion.association.AssociationDiskStore;
 import com.android.server.companion.association.AssociationRequestsProcessor;
 import com.android.server.companion.association.AssociationStore;
+import com.android.server.companion.association.Associations;
 import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
 
 import java.nio.ByteBuffer;
-import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
-import java.util.Set;
 import java.util.function.Predicate;
 
 @SuppressLint("LongLogTag")
 class BackupRestoreProcessor {
-    static final String TAG = "CDM_BackupRestoreProcessor";
+    private static final String TAG = "CDM_BackupRestoreProcessor";
     private static final int BACKUP_AND_RESTORE_VERSION = 0;
 
+    private final Context mContext;
     @NonNull
-    private final CompanionDeviceManagerService mService;
-    @NonNull
-    private final PackageManagerInternal mPackageManager;
+    private final PackageManagerInternal mPackageManagerInternal;
     @NonNull
     private final AssociationStore mAssociationStore;
     @NonNull
-    private final AssociationDiskStore mPersistentStore;
+    private final AssociationDiskStore mAssociationDiskStore;
     @NonNull
     private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
     @NonNull
     private final AssociationRequestsProcessor mAssociationRequestsProcessor;
 
-    /**
-     * A structure that consists of a set of restored associations that are pending corresponding
-     * companion app to be installed.
-     */
-    @GuardedBy("mAssociationsPendingAppInstall")
-    private final PerUserAssociationSet mAssociationsPendingAppInstall =
-            new PerUserAssociationSet();
-
-    BackupRestoreProcessor(@NonNull CompanionDeviceManagerService service,
+    BackupRestoreProcessor(@NonNull Context context,
+                           @NonNull PackageManagerInternal packageManagerInternal,
                            @NonNull AssociationStore associationStore,
-                           @NonNull AssociationDiskStore persistentStore,
+                           @NonNull AssociationDiskStore associationDiskStore,
                            @NonNull SystemDataTransferRequestStore systemDataTransferRequestStore,
                            @NonNull AssociationRequestsProcessor associationRequestsProcessor) {
-        mService = service;
-        mPackageManager = service.mPackageManagerInternal;
+        mContext = context;
+        mPackageManagerInternal = packageManagerInternal;
         mAssociationStore = associationStore;
-        mPersistentStore = persistentStore;
+        mAssociationDiskStore = associationDiskStore;
         mSystemDataTransferRequestStore = systemDataTransferRequestStore;
         mAssociationRequestsProcessor = associationRequestsProcessor;
     }
@@ -93,9 +82,9 @@
      * | (4) SystemDataTransferRequest length | SystemDataTransferRequest XML (without userId)|
      */
     byte[] getBackupPayload(int userId) {
-        // Persist state first to generate an up-to-date XML file
-        mService.persistStateForUser(userId);
-        byte[] associationsPayload = mPersistentStore.getBackupPayload(userId);
+        Slog.i(TAG, "getBackupPayload() userId=[" + userId + "].");
+
+        byte[] associationsPayload = mAssociationDiskStore.getBackupPayload(userId);
         int associationsPayloadLength = associationsPayload.length;
 
         // System data transfer requests are persisted up-to-date already
@@ -119,6 +108,9 @@
      * Create new associations and system data transfer request consents using backed up payload.
      */
     void applyRestoredPayload(byte[] payload, int userId) {
+        Slog.i(TAG, "applyRestoredPayload() userId=[" + userId + "], payload size=["
+                + payload.length + "].");
+
         ByteBuffer buffer = ByteBuffer.wrap(payload);
 
         // Make sure that payload version matches current version to ensure proper deserialization
@@ -131,9 +123,8 @@
         // Read the bytes containing backed-up associations
         byte[] associationsPayload = new byte[buffer.getInt()];
         buffer.get(associationsPayload);
-        final Set<AssociationInfo> restoredAssociations = new HashSet<>();
-        mPersistentStore.readStateFromPayload(associationsPayload, userId,
-                restoredAssociations, new HashMap<>());
+        final Associations restoredAssociations = readAssociationsFromPayload(
+                associationsPayload, userId);
 
         // Read the bytes containing backed-up system data transfer requests user consent
         byte[] requestsPayload = new byte[buffer.getInt()];
@@ -142,13 +133,13 @@
                 mSystemDataTransferRequestStore.readRequestsFromPayload(requestsPayload, userId);
 
         // Get a list of installed packages ahead of time.
-        List<ApplicationInfo> installedApps = mPackageManager.getInstalledApplications(
+        List<ApplicationInfo> installedApps = mPackageManagerInternal.getInstalledApplications(
                 0, userId, getCallingUserId());
 
         // Restored device may have a different user ID than the backed-up user's user-ID. Since
         // association ID is dependent on the user ID, restored associations must account for
         // this potential difference on their association IDs.
-        for (AssociationInfo restored : restoredAssociations) {
+        for (AssociationInfo restored : restoredAssociations.getAssociations()) {
             // Don't restore a revoked association. Since they weren't added to the device being
             // restored in the first place, there is no need to worry about revoking a role that
             // was never granted either.
@@ -168,10 +159,9 @@
 
             // Create a new association reassigned to this user and a valid association ID
             final String packageName = restored.getPackageName();
-            final int newId = mService.getNewAssociationIdForPackage(userId, packageName);
-            AssociationInfo newAssociation =
-                    new AssociationInfo.Builder(newId, userId, packageName, restored)
-                            .build();
+            final int newId = mAssociationStore.getNextId(userId);
+            AssociationInfo newAssociation = new AssociationInfo.Builder(newId, userId, packageName,
+                    restored).build();
 
             // Check if the companion app for this association is already installed, then do one
             // of the following:
@@ -179,13 +169,15 @@
             // the role attached to this association to the app.
             // (2) If the app isn't yet installed, then add this association to the list of pending
             // associations to be added when the package is installed in the future.
-            boolean isPackageInstalled = installedApps.stream()
-                    .anyMatch(app -> packageName.equals(app.packageName));
+            boolean isPackageInstalled = installedApps.stream().anyMatch(
+                    app -> packageName.equals(app.packageName));
             if (isPackageInstalled) {
                 mAssociationRequestsProcessor.maybeGrantRoleAndStoreAssociation(newAssociation,
                         null, null);
             } else {
-                addToPendingAppInstall(newAssociation);
+                newAssociation = (new AssociationInfo.Builder(newAssociation)).setPending(true)
+                        .build();
+                mAssociationStore.addAssociation(newAssociation);
             }
 
             // Re-map restored system data transfer requests to newly created associations
@@ -195,32 +187,27 @@
                 mSystemDataTransferRequestStore.writeRequest(userId, newRequest);
             }
         }
-
-        // Persist restored state.
-        mService.persistStateForUser(userId);
     }
 
-    void addToPendingAppInstall(@NonNull AssociationInfo association) {
-        association = (new AssociationInfo.Builder(association))
-                .setPending(true)
-                .build();
-
-        synchronized (mAssociationsPendingAppInstall) {
-            mAssociationsPendingAppInstall.forUser(association.getUserId()).add(association);
+    public void restorePendingAssociations(int userId, String packageName) {
+        List<AssociationInfo> pendingAssociations = mAssociationStore.getPendingAssociations(userId,
+                packageName);
+        if (!pendingAssociations.isEmpty()) {
+            Slog.i(TAG, "Found pending associations for package=[" + packageName
+                    + "]. Restoring...");
         }
-    }
-
-    void removeFromPendingAppInstall(@NonNull AssociationInfo association) {
-        synchronized (mAssociationsPendingAppInstall) {
-            mAssociationsPendingAppInstall.forUser(association.getUserId()).remove(association);
-        }
-    }
-
-    @NonNull
-    Set<AssociationInfo> getAssociationsPendingAppInstallForUser(@UserIdInt int userId) {
-        synchronized (mAssociationsPendingAppInstall) {
-            // Return a copy.
-            return new ArraySet<>(mAssociationsPendingAppInstall.forUser(userId));
+        for (AssociationInfo association : pendingAssociations) {
+            AssociationInfo newAssociation = new AssociationInfo.Builder(association)
+                    .setPending(false)
+                    .build();
+            addRoleHolderForAssociation(mContext, newAssociation, success -> {
+                if (success) {
+                    mAssociationStore.updateAssociation(newAssociation);
+                    Slog.i(TAG, "Association=[" + association + "] is restored.");
+                } else {
+                    Slog.e(TAG, "Failed to restore association=[" + association + "].");
+                }
+            });
         }
     }
 
@@ -231,7 +218,7 @@
     private boolean handleCollision(@UserIdInt int userId,
             AssociationInfo restored,
             List<SystemDataTransferRequest> restoredRequests) {
-        List<AssociationInfo> localAssociations = mAssociationStore.getAssociationsForPackage(
+        List<AssociationInfo> localAssociations = mAssociationStore.getActiveAssociationsByPackage(
                 restored.getUserId(), restored.getPackageName());
         Predicate<AssociationInfo> isSameDevice = associationInfo -> {
             boolean matchesMacAddress = Objects.equals(
@@ -248,7 +235,7 @@
             return false;
         }
 
-        Log.d(TAG, "Conflict detected with association id=" + local.getId()
+        Slog.d(TAG, "Conflict detected with association id=" + local.getId()
                 + " while restoring CDM backup. Keeping local association.");
 
         List<SystemDataTransferRequest> localRequests = mSystemDataTransferRequestStore
@@ -266,8 +253,8 @@
                 continue;
             }
 
-            Log.d(TAG, "Restoring " + restoredRequest.getClass().getSimpleName()
-                    + " to an existing association id=" + local.getId() + ".");
+            Slog.d(TAG, "Restoring " + restoredRequest.getClass().getSimpleName()
+                    + " to an existing association id=[" + local.getId() + "].");
 
             SystemDataTransferRequest newRequest =
                     restoredRequest.copyWithNewId(local.getId());
diff --git a/services/companion/java/com/android/server/companion/CompanionApplicationController.java b/services/companion/java/com/android/server/companion/CompanionApplicationController.java
index c801489..0a41485 100644
--- a/services/companion/java/com/android/server/companion/CompanionApplicationController.java
+++ b/services/companion/java/com/android/server/companion/CompanionApplicationController.java
@@ -397,7 +397,7 @@
         // First, disable hint mode for Auto profile and mark not BOUND for primary service ONLY.
         if (isPrimary) {
             final List<AssociationInfo> associations =
-                    mAssociationStore.getAssociationsForPackage(userId, packageName);
+                    mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
 
             for (AssociationInfo association : associations) {
                 final String deviceProfile = association.getDeviceProfile();
@@ -442,7 +442,7 @@
                 mObservableUuidStore.getObservableUuidsForPackage(userId, packageName);
 
         for (AssociationInfo ai :
-                mAssociationStore.getAssociationsForPackage(userId, packageName)) {
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName)) {
             final int associationId = ai.getId();
             stillAssociated = true;
             if (ai.isSelfManaged()) {
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
index 3846e98..73ebbc7 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceManagerService.java
@@ -37,12 +37,9 @@
 import static com.android.internal.util.CollectionUtils.any;
 import static com.android.internal.util.Preconditions.checkState;
 import static com.android.internal.util.function.pooled.PooledLambda.obtainMessage;
-import static com.android.server.companion.association.AssociationStore.CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED;
-import static com.android.server.companion.utils.AssociationUtils.getFirstAssociationIdForUser;
-import static com.android.server.companion.utils.AssociationUtils.getLastAssociationIdForUser;
-import static com.android.server.companion.utils.PackageUtils.isRestrictedSettingsAllowed;
 import static com.android.server.companion.utils.PackageUtils.enforceUsesCompanionDeviceFeature;
 import static com.android.server.companion.utils.PackageUtils.getPackageInfo;
+import static com.android.server.companion.utils.PackageUtils.isRestrictedSettingsAllowed;
 import static com.android.server.companion.utils.PermissionsUtils.checkCallerCanManageCompanionDevice;
 import static com.android.server.companion.utils.PermissionsUtils.enforceCallerCanManageAssociationsForPackage;
 import static com.android.server.companion.utils.PermissionsUtils.enforceCallerCanObservingDevicePresenceByUuid;
@@ -82,20 +79,16 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManagerInternal;
-import android.content.pm.UserInfo;
 import android.hardware.power.Mode;
 import android.net.MacAddress;
 import android.net.NetworkPolicyManager;
 import android.os.Binder;
 import android.os.Environment;
-import android.os.Handler;
-import android.os.Message;
 import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.ParcelUuid;
+import android.os.PowerExemptionManager;
 import android.os.PowerManagerInternal;
-import android.os.PowerWhitelistManager;
-import android.os.RemoteCallbackList;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.SystemProperties;
@@ -105,13 +98,9 @@
 import android.util.ExceptionUtils;
 import android.util.Log;
 import android.util.Slog;
-import android.util.SparseArray;
-import android.util.SparseBooleanArray;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.app.IAppOpsService;
 import com.android.internal.content.PackageMonitor;
-import com.android.internal.infra.PerUser;
 import com.android.internal.notification.NotificationAccessConfirmationActivityContract;
 import com.android.internal.os.BackgroundThread;
 import com.android.internal.util.ArrayUtils;
@@ -121,8 +110,8 @@
 import com.android.server.SystemService;
 import com.android.server.companion.association.AssociationDiskStore;
 import com.android.server.companion.association.AssociationRequestsProcessor;
-import com.android.server.companion.association.AssociationRevokeProcessor;
 import com.android.server.companion.association.AssociationStore;
+import com.android.server.companion.association.DisassociationProcessor;
 import com.android.server.companion.association.InactiveAssociationsRemovalService;
 import com.android.server.companion.datatransfer.SystemDataTransferProcessor;
 import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
@@ -139,7 +128,6 @@
 import java.io.File;
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
@@ -164,80 +152,51 @@
     private static final long ASSOCIATION_REMOVAL_TIME_WINDOW_DEFAULT = DAYS.toMillis(90);
     private static final int MAX_CN_LENGTH = 500;
 
-    private final ActivityManager mActivityManager;
-    private AssociationDiskStore mAssociationDiskStore;
-    private final PersistUserStateHandler mUserPersistenceHandler;
-
-    private final AssociationStore mAssociationStore;
-    private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
-    private AssociationRequestsProcessor mAssociationRequestsProcessor;
-    private SystemDataTransferProcessor mSystemDataTransferProcessor;
-    private BackupRestoreProcessor mBackupRestoreProcessor;
-    private CompanionDevicePresenceMonitor mDevicePresenceMonitor;
-    private CompanionApplicationController mCompanionAppController;
-    private CompanionTransportManager mTransportManager;
-    private AssociationRevokeProcessor mAssociationRevokeProcessor;
-
     private final ActivityTaskManagerInternal mAtmInternal;
     private final ActivityManagerInternal mAmInternal;
     private final IAppOpsService mAppOpsManager;
-    private final PowerWhitelistManager mPowerWhitelistManager;
-    private final UserManager mUserManager;
-    public final PackageManagerInternal mPackageManagerInternal;
+    private final PowerExemptionManager mPowerExemptionManager;
+    private final PackageManagerInternal mPackageManagerInternal;
     private final PowerManagerInternal mPowerManagerInternal;
 
-    /**
-     * A structure that consists of two nested maps, and effectively maps (userId + packageName) to
-     * a list of IDs that have been previously assigned to associations for that package.
-     * We maintain this structure so that we never re-use association IDs for the same package
-     * (until it's uninstalled).
-     */
-    @GuardedBy("mPreviouslyUsedIds")
-    private final SparseArray<Map<String, Set<Integer>>> mPreviouslyUsedIds = new SparseArray<>();
-
-    private final RemoteCallbackList<IOnAssociationsChangedListener> mListeners =
-            new RemoteCallbackList<>();
-
-    private CrossDeviceSyncController mCrossDeviceSyncController;
-
-    private ObservableUuidStore mObservableUuidStore;
+    private final AssociationStore mAssociationStore;
+    private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
+    private final ObservableUuidStore mObservableUuidStore;
+    private final AssociationRequestsProcessor mAssociationRequestsProcessor;
+    private final SystemDataTransferProcessor mSystemDataTransferProcessor;
+    private final BackupRestoreProcessor mBackupRestoreProcessor;
+    private final CompanionDevicePresenceMonitor mDevicePresenceMonitor;
+    private final CompanionApplicationController mCompanionAppController;
+    private final CompanionTransportManager mTransportManager;
+    private final DisassociationProcessor mDisassociationProcessor;
+    private final CrossDeviceSyncController mCrossDeviceSyncController;
 
     public CompanionDeviceManagerService(Context context) {
         super(context);
 
-        mActivityManager = context.getSystemService(ActivityManager.class);
-        mPowerWhitelistManager = context.getSystemService(PowerWhitelistManager.class);
+        final ActivityManager activityManager = context.getSystemService(ActivityManager.class);
+        mPowerExemptionManager = context.getSystemService(PowerExemptionManager.class);
         mAppOpsManager = IAppOpsService.Stub.asInterface(
                 ServiceManager.getService(Context.APP_OPS_SERVICE));
         mAtmInternal = LocalServices.getService(ActivityTaskManagerInternal.class);
         mAmInternal = LocalServices.getService(ActivityManagerInternal.class);
         mPackageManagerInternal = LocalServices.getService(PackageManagerInternal.class);
-        mUserManager = context.getSystemService(UserManager.class);
-
-        mUserPersistenceHandler = new PersistUserStateHandler();
-        mAssociationStore = new AssociationStore();
-        mSystemDataTransferRequestStore = new SystemDataTransferRequestStore();
-
+        final UserManager userManager = context.getSystemService(UserManager.class);
         mPowerManagerInternal = LocalServices.getService(PowerManagerInternal.class);
+
+        final AssociationDiskStore associationDiskStore = new AssociationDiskStore();
+        mAssociationStore = new AssociationStore(userManager, associationDiskStore);
+        mSystemDataTransferRequestStore = new SystemDataTransferRequestStore();
         mObservableUuidStore = new ObservableUuidStore();
-    }
 
-    @Override
-    public void onStart() {
-        final Context context = getContext();
+        // Init processors
+        mAssociationRequestsProcessor = new AssociationRequestsProcessor(context,
+                mPackageManagerInternal, mAssociationStore);
+        mBackupRestoreProcessor = new BackupRestoreProcessor(context, mPackageManagerInternal,
+                mAssociationStore, associationDiskStore, mSystemDataTransferRequestStore,
+                mAssociationRequestsProcessor);
 
-        mAssociationDiskStore = new AssociationDiskStore();
-        mAssociationRequestsProcessor = new AssociationRequestsProcessor(
-                /* cdmService */ this, mAssociationStore);
-        mBackupRestoreProcessor = new BackupRestoreProcessor(
-                /* cdmService */ this, mAssociationStore, mAssociationDiskStore,
-                mSystemDataTransferRequestStore, mAssociationRequestsProcessor);
-
-        mObservableUuidStore.getObservableUuidsForUser(getContext().getUserId());
-
-        mAssociationStore.registerListener(mAssociationStoreChangeListener);
-
-        mDevicePresenceMonitor = new CompanionDevicePresenceMonitor(mUserManager,
+        mDevicePresenceMonitor = new CompanionDevicePresenceMonitor(userManager,
                 mAssociationStore, mObservableUuidStore, mDevicePresenceCallback);
 
         mCompanionAppController = new CompanionApplicationController(
@@ -246,11 +205,9 @@
 
         mTransportManager = new CompanionTransportManager(context, mAssociationStore);
 
-        mAssociationRevokeProcessor = new AssociationRevokeProcessor(this, mAssociationStore,
-                mPackageManagerInternal, mDevicePresenceMonitor, mCompanionAppController,
-                mSystemDataTransferRequestStore, mTransportManager);
-
-        loadAssociationsFromDisk();
+        mDisassociationProcessor = new DisassociationProcessor(context, activityManager,
+                mAssociationStore, mPackageManagerInternal, mDevicePresenceMonitor,
+                mCompanionAppController, mSystemDataTransferRequestStore, mTransportManager);
 
         mSystemDataTransferProcessor = new SystemDataTransferProcessor(this,
                 mPackageManagerInternal, mAssociationStore,
@@ -258,6 +215,16 @@
 
         // TODO(b/279663946): move context sync to a dedicated system service
         mCrossDeviceSyncController = new CrossDeviceSyncController(getContext(), mTransportManager);
+    }
+
+    @Override
+    public void onStart() {
+        // Init association stores
+        mAssociationStore.refreshCache();
+        mAssociationStore.registerLocalListener(mAssociationStoreChangeListener);
+
+        // Init UUID store
+        mObservableUuidStore.getObservableUuidsForUser(getContext().getUserId());
 
         // Publish "binder" service.
         final CompanionDeviceManagerImpl impl = new CompanionDeviceManagerImpl();
@@ -267,50 +234,6 @@
         LocalServices.addService(CompanionDeviceManagerServiceInternal.class, new LocalService());
     }
 
-    void loadAssociationsFromDisk() {
-        final Set<AssociationInfo> allAssociations = new ArraySet<>();
-        synchronized (mPreviouslyUsedIds) {
-            List<Integer> userIds = new ArrayList<>();
-            for (UserInfo user : mUserManager.getAliveUsers()) {
-                userIds.add(user.id);
-            }
-            // The data is stored in DE directories, so we can read the data for all users now
-            // (which would not be possible if the data was stored to CE directories).
-            mAssociationDiskStore.readStateForUsers(userIds, allAssociations, mPreviouslyUsedIds);
-        }
-
-        final Set<AssociationInfo> activeAssociations =
-                new ArraySet<>(/* capacity */ allAssociations.size());
-        // A set contains the userIds that need to persist state after remove the app
-        // from the list of role holders.
-        final Set<Integer> usersToPersistStateFor = new ArraySet<>();
-
-        for (AssociationInfo association : allAssociations) {
-            if (association.isPending()) {
-                mBackupRestoreProcessor.addToPendingAppInstall(association);
-            } else if (!association.isRevoked()) {
-                activeAssociations.add(association);
-            } else if (mAssociationRevokeProcessor.maybeRemoveRoleHolderForAssociation(
-                    association)) {
-                // Nothing more to do here, but we'll need to persist all the associations to the
-                // disk afterwards.
-                usersToPersistStateFor.add(association.getUserId());
-            } else {
-                mAssociationRevokeProcessor.addToPendingRoleHolderRemoval(association);
-            }
-        }
-
-        mAssociationStore.setAssociationsToCache(activeAssociations);
-
-        // IMPORTANT: only do this AFTER mAssociationStore.setAssociations(), because
-        // persistStateForUser() queries AssociationStore.
-        // (If persistStateForUser() is invoked before mAssociationStore.setAssociations() it
-        // would effectively just clear-out all the persisted associations).
-        for (int userId : usersToPersistStateFor) {
-            persistStateForUser(userId);
-        }
-    }
-
     @Override
     public void onBootPhase(int phase) {
         final Context context = getContext();
@@ -329,8 +252,10 @@
 
     @Override
     public void onUserUnlocking(@NonNull TargetUser user) {
+        Slog.d(TAG, "onUserUnlocking...");
         final int userId = user.getUserIdentifier();
-        final List<AssociationInfo> associations = mAssociationStore.getAssociationsForUser(userId);
+        final List<AssociationInfo> associations = mAssociationStore.getActiveAssociationsByUser(
+                userId);
 
         if (associations.isEmpty()) return;
 
@@ -359,7 +284,8 @@
                         ? Collections.emptyList() : Arrays.asList(bluetoothDeviceUuids);
 
                 for (AssociationInfo ai :
-                        mAssociationStore.getAssociationsByAddress(bluetoothDevice.getAddress())) {
+                        mAssociationStore.getActiveAssociationsByAddress(
+                                bluetoothDevice.getAddress())) {
                     Slog.i(TAG, "onUserUnlocked, device id( " + ai.getId() + " ) is connected");
                     mDevicePresenceMonitor.onBluetoothCompanionDeviceConnected(ai.getId());
                 }
@@ -379,7 +305,7 @@
     @NonNull
     AssociationInfo getAssociationWithCallerChecks(
             @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress) {
-        AssociationInfo association = mAssociationStore.getAssociationsForPackageWithAddress(
+        AssociationInfo association = mAssociationStore.getFirstAssociationByAddress(
                 userId, packageName, macAddress);
         association = sanitizeWithCallerChecks(getContext(), association);
         if (association != null) {
@@ -533,7 +459,7 @@
      */
     private boolean shouldBindPackage(@UserIdInt int userId, @NonNull String packageName) {
         final List<AssociationInfo> packageAssociations =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
         final List<ObservableUuid> observableUuids =
                 mObservableUuidStore.getObservableUuidsForPackage(userId, packageName);
 
@@ -551,77 +477,6 @@
         return false;
     }
 
-    private void onAssociationChangedInternal(
-            @AssociationStore.ChangeType int changeType, AssociationInfo association) {
-        final int id = association.getId();
-        final int userId = association.getUserId();
-        final String packageName = association.getPackageName();
-
-        if (changeType == AssociationStore.CHANGE_TYPE_REMOVED) {
-            markIdAsPreviouslyUsedForPackage(id, userId, packageName);
-        }
-
-        final List<AssociationInfo> updatedAssociations =
-                mAssociationStore.getAssociationsForUser(userId);
-
-        mUserPersistenceHandler.postPersistUserState(userId);
-
-        // Notify listeners if ADDED, REMOVED or UPDATED_ADDRESS_CHANGED.
-        // Do NOT notify when UPDATED_ADDRESS_UNCHANGED, which means a minor tweak in association's
-        // configs, which "listeners" won't (and shouldn't) be able to see.
-        if (changeType != CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED) {
-            notifyListeners(userId, updatedAssociations);
-        }
-        updateAtm(userId, updatedAssociations);
-    }
-
-    void persistStateForUser(@UserIdInt int userId) {
-        // We want to store both active associations and the revoked (removed) association that we
-        // are keeping around for the final clean-up (delayed role holder removal).
-        final List<AssociationInfo> allAssociations;
-        // Start with the active associations - these we can get from the AssociationStore.
-        allAssociations = new ArrayList<>(
-                mAssociationStore.getAssociationsForUser(userId));
-        // ... and add the revoked (removed) association, that are yet to be permanently removed.
-        allAssociations.addAll(
-                mAssociationRevokeProcessor.getPendingRoleHolderRemovalAssociationsForUser(userId));
-        // ... and add the restored associations that are pending missing package installation.
-        allAssociations.addAll(mBackupRestoreProcessor
-                .getAssociationsPendingAppInstallForUser(userId));
-
-        final Map<String, Set<Integer>> usedIdsForUser = getPreviouslyUsedIdsForUser(userId);
-
-        mAssociationDiskStore.persistStateForUser(userId, allAssociations, usedIdsForUser);
-    }
-
-    private void notifyListeners(
-            @UserIdInt int userId, @NonNull List<AssociationInfo> associations) {
-        mListeners.broadcast((listener, callbackUserId) -> {
-            int listenerUserId = (int) callbackUserId;
-            if (listenerUserId == userId || listenerUserId == UserHandle.USER_ALL) {
-                try {
-                    listener.onAssociationsChanged(associations);
-                } catch (RemoteException ignored) {
-                }
-            }
-        });
-    }
-
-    private void markIdAsPreviouslyUsedForPackage(
-            int associationId, @UserIdInt int userId, @NonNull String packageName) {
-        synchronized (mPreviouslyUsedIds) {
-            Map<String, Set<Integer>> usedIdsForUser = mPreviouslyUsedIds.get(userId);
-            if (usedIdsForUser == null) {
-                usedIdsForUser = new HashMap<>();
-                mPreviouslyUsedIds.put(userId, usedIdsForUser);
-            }
-
-            final Set<Integer> usedIdsForPackage =
-                    usedIdsForUser.computeIfAbsent(packageName, it -> new HashSet<>());
-            usedIdsForPackage.add(associationId);
-        }
-    }
-
     private void onPackageRemoveOrDataClearedInternal(
             @UserIdInt int userId, @NonNull String packageName) {
         if (DEBUG) {
@@ -629,19 +484,20 @@
                     + packageName);
         }
 
-        // Clear associations.
+        // Clear all associations for the package.
         final List<AssociationInfo> associationsForPackage =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
+                mAssociationStore.getAssociationsByPackage(userId, packageName);
+        if (!associationsForPackage.isEmpty()) {
+            Slog.i(TAG, "Package removed or data cleared for user=[" + userId + "], package=["
+                    + packageName + "]. Cleaning up CDM data...");
+        }
+        for (AssociationInfo association : associationsForPackage) {
+            mDisassociationProcessor.disassociate(association.getId());
+        }
+
+        // Clear observable UUIDs for the package.
         final List<ObservableUuid> uuidsTobeObserved =
                 mObservableUuidStore.getObservableUuidsForPackage(userId, packageName);
-        for (AssociationInfo association : associationsForPackage) {
-            mAssociationStore.removeAssociation(association.getId());
-        }
-        // Clear role holders
-        for (AssociationInfo association : associationsForPackage) {
-            mAssociationRevokeProcessor.maybeRemoveRoleHolderForAssociation(association);
-        }
-        // Clear the uuids to be observed.
         for (ObservableUuid uuid : uuidsTobeObserved) {
             mObservableUuidStore.removeObservableUuid(userId, uuid.getUuid(), packageName);
         }
@@ -652,31 +508,13 @@
     private void onPackageModifiedInternal(@UserIdInt int userId, @NonNull String packageName) {
         if (DEBUG) Log.i(TAG, "onPackageModified() u" + userId + "/" + packageName);
 
-        final List<AssociationInfo> associationsForPackage =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
-        for (AssociationInfo association : associationsForPackage) {
-            updateSpecialAccessPermissionForAssociatedPackage(association.getUserId(),
-                    association.getPackageName());
-        }
+        updateSpecialAccessPermissionForAssociatedPackage(userId, packageName);
 
         mCompanionAppController.onPackagesChanged(userId);
     }
 
     private void onPackageAddedInternal(@UserIdInt int userId, @NonNull String packageName) {
-        if (DEBUG) Log.i(TAG, "onPackageAddedInternal() u" + userId + "/" + packageName);
-
-        Set<AssociationInfo> associationsPendingAppInstall = mBackupRestoreProcessor
-                .getAssociationsPendingAppInstallForUser(userId);
-        for (AssociationInfo association : associationsPendingAppInstall) {
-            if (!packageName.equals(association.getPackageName())) continue;
-
-            AssociationInfo newAssociation = new AssociationInfo.Builder(association)
-                    .setPending(false)
-                    .build();
-            mAssociationRequestsProcessor.maybeGrantRoleAndStoreAssociation(newAssociation,
-                    null, null);
-            mBackupRestoreProcessor.removeFromPendingAppInstall(association);
-        }
+        mBackupRestoreProcessor.restorePendingAssociations(userId, packageName);
     }
 
     // Revoke associations if the selfManaged companion device does not connect for 3 months.
@@ -698,7 +536,7 @@
             final int id = association.getId();
 
             Slog.i(TAG, "Removing inactive self-managed association id=" + id);
-            mAssociationRevokeProcessor.disassociateInternal(id);
+            mDisassociationProcessor.disassociate(id);
         }
     }
 
@@ -750,7 +588,7 @@
                 enforceUsesCompanionDeviceFeature(getContext(), userId, packageName);
             }
 
-            return mAssociationStore.getAssociationsForPackage(userId, packageName);
+            return mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
         }
 
         @Override
@@ -761,9 +599,9 @@
             enforceCallerIsSystemOrCanInteractWithUserId(getContext(), userId);
 
             if (userId == UserHandle.USER_ALL) {
-                return List.copyOf(mAssociationStore.getAssociations());
+                return mAssociationStore.getActiveAssociations();
             }
-            return mAssociationStore.getAssociationsForUser(userId);
+            return mAssociationStore.getActiveAssociationsByUser(userId);
         }
 
         @Override
@@ -773,7 +611,8 @@
             addOnAssociationsChangedListener_enforcePermission();
 
             enforceCallerIsSystemOrCanInteractWithUserId(getContext(), userId);
-            mListeners.register(listener, userId);
+
+            mAssociationStore.registerRemoteListener(listener, userId);
         }
 
         @Override
@@ -784,7 +623,7 @@
 
             enforceCallerIsSystemOrCanInteractWithUserId(getContext(), userId);
 
-            mListeners.unregister(listener);
+            mAssociationStore.unregisterRemoteListener(listener);
         }
 
         @Override
@@ -843,16 +682,16 @@
 
             final AssociationInfo association =
                     getAssociationWithCallerChecks(userId, packageName, deviceMacAddress);
-            mAssociationRevokeProcessor.disassociateInternal(association.getId());
+            mDisassociationProcessor.disassociate(association.getId());
         }
 
         @Override
         public void disassociate(int associationId) {
-            Log.i(TAG, "disassociate() associationId=" + associationId);
+            Slog.i(TAG, "disassociate() associationId=" + associationId);
 
             final AssociationInfo association =
                     getAssociationWithCallerChecks(associationId);
-            mAssociationRevokeProcessor.disassociateInternal(association.getId());
+            mDisassociationProcessor.disassociate(association.getId());
         }
 
         @Override
@@ -867,8 +706,7 @@
                 throw new IllegalArgumentException("Component name is too long.");
             }
 
-            final long identity = Binder.clearCallingIdentity();
-            try {
+            return Binder.withCleanCallingIdentity(() -> {
                 if (!isRestrictedSettingsAllowed(getContext(), callingPackage, callingUid)) {
                     Slog.e(TAG, "Side loaded app must enable restricted "
                             + "setting before request the notification access");
@@ -882,9 +720,7 @@
                                 | PendingIntent.FLAG_CANCEL_CURRENT,
                         null /* options */,
                         new UserHandle(userId));
-            } finally {
-                Binder.restoreCallingIdentity(identity);
-            }
+            });
         }
 
         /**
@@ -912,7 +748,7 @@
                 return true;
             }
 
-            return any(mAssociationStore.getAssociationsForPackage(userId, packageName),
+            return any(mAssociationStore.getActiveAssociationsByPackage(userId, packageName),
                     a -> a.isLinkedTo(macAddress));
         }
 
@@ -1166,7 +1002,7 @@
             final int userId = getCallingUserId();
             enforceCallerIsSystemOr(userId, packageName);
 
-            AssociationInfo association = mAssociationStore.getAssociationsForPackageWithAddress(
+            AssociationInfo association = mAssociationStore.getFirstAssociationByAddress(
                     userId, packageName, deviceAddress);
 
             if (association == null) {
@@ -1239,14 +1075,15 @@
 
             enforceUsesCompanionDeviceFeature(getContext(), userId, callingPackage);
             checkState(!ArrayUtils.isEmpty(
-                            mAssociationStore.getAssociationsForPackage(userId, callingPackage)),
+                            mAssociationStore.getActiveAssociationsByPackage(userId,
+                                    callingPackage)),
                     "App must have an association before calling this API");
         }
 
         @Override
         public boolean canPairWithoutPrompt(String packageName, String macAddress, int userId) {
             final AssociationInfo association =
-                    mAssociationStore.getAssociationsForPackageWithAddress(
+                    mAssociationStore.getFirstAssociationByAddress(
                             userId, packageName, macAddress);
             if (association == null) {
                 return false;
@@ -1269,13 +1106,11 @@
 
         @Override
         public byte[] getBackupPayload(int userId) {
-            Log.i(TAG, "getBackupPayload() userId=" + userId);
             return mBackupRestoreProcessor.getBackupPayload(userId);
         }
 
         @Override
         public void applyRestoredPayload(byte[] payload, int userId) {
-            Log.i(TAG, "applyRestoredPayload() userId=" + userId);
             mBackupRestoreProcessor.applyRestoredPayload(payload, userId);
         }
 
@@ -1286,7 +1121,7 @@
             return new CompanionDeviceShellCommand(CompanionDeviceManagerService.this,
                     mAssociationStore, mDevicePresenceMonitor, mTransportManager,
                     mSystemDataTransferProcessor, mAssociationRequestsProcessor,
-                    mBackupRestoreProcessor, mAssociationRevokeProcessor)
+                    mBackupRestoreProcessor, mDisassociationProcessor)
                     .exec(this, in.getFileDescriptor(), out.getFileDescriptor(),
                             err.getFileDescriptor(), args);
         }
@@ -1314,88 +1149,6 @@
                 /* callback */ null, /* resultReceiver */ null);
     }
 
-    @NonNull
-    private Map<String, Set<Integer>> getPreviouslyUsedIdsForUser(@UserIdInt int userId) {
-        synchronized (mPreviouslyUsedIds) {
-            return getPreviouslyUsedIdsForUserLocked(userId);
-        }
-    }
-
-    @GuardedBy("mPreviouslyUsedIds")
-    @NonNull
-    private Map<String, Set<Integer>> getPreviouslyUsedIdsForUserLocked(@UserIdInt int userId) {
-        final Map<String, Set<Integer>> usedIdsForUser = mPreviouslyUsedIds.get(userId);
-        if (usedIdsForUser == null) {
-            return Collections.emptyMap();
-        }
-        return deepUnmodifiableCopy(usedIdsForUser);
-    }
-
-    @GuardedBy("mPreviouslyUsedIds")
-    @NonNull
-    private Set<Integer> getPreviouslyUsedIdsForPackageLocked(
-            @UserIdInt int userId, @NonNull String packageName) {
-        // "Deeply unmodifiable" map: the map itself and the Set<Integer> values it contains are all
-        // unmodifiable.
-        final Map<String, Set<Integer>> usedIdsForUser = getPreviouslyUsedIdsForUserLocked(userId);
-        final Set<Integer> usedIdsForPackage = usedIdsForUser.get(packageName);
-
-        if (usedIdsForPackage == null) {
-            return Collections.emptySet();
-        }
-
-        //The set is already unmodifiable.
-        return usedIdsForPackage;
-    }
-
-    /**
-     * Get a new association id for the package.
-     */
-    public int getNewAssociationIdForPackage(@UserIdInt int userId, @NonNull String packageName) {
-        synchronized (mPreviouslyUsedIds) {
-            // First: collect all IDs currently in use for this user's Associations.
-            final SparseBooleanArray usedIds = new SparseBooleanArray();
-
-            // We should really only be checking associations for the given user (i.e.:
-            // mAssociationStore.getAssociationsForUser(userId)), BUT in the past we've got in a
-            // state where association IDs were not assigned correctly in regard to
-            // user-to-association-ids-range (e.g. associations with IDs from 1 to 100,000 should
-            // always belong to u0), so let's check all the associations.
-            for (AssociationInfo it : mAssociationStore.getAssociations()) {
-                usedIds.put(it.getId(), true);
-            }
-
-            // Some IDs may be reserved by associations that aren't stored yet due to missing
-            // package after a backup restoration. We don't want the ID to have been taken by
-            // another association by the time when it is activated from the package installation.
-            final Set<AssociationInfo> pendingAssociations = mBackupRestoreProcessor
-                    .getAssociationsPendingAppInstallForUser(userId);
-            for (AssociationInfo it : pendingAssociations) {
-                usedIds.put(it.getId(), true);
-            }
-
-            // Second: collect all IDs that have been previously used for this package (and user).
-            final Set<Integer> previouslyUsedIds =
-                    getPreviouslyUsedIdsForPackageLocked(userId, packageName);
-
-            int id = getFirstAssociationIdForUser(userId);
-            final int lastAvailableIdForUser = getLastAssociationIdForUser(userId);
-
-            // Find first ID that isn't used now AND has never been used for the given package.
-            while (usedIds.get(id) || previouslyUsedIds.contains(id)) {
-                // Increment and try again
-                id++;
-                // ... but first check if the ID is valid (within the range allocated to the user).
-                if (id > lastAvailableIdForUser) {
-                    throw new RuntimeException("Cannot create a new Association ID for "
-                            + packageName + " for user " + userId);
-                }
-            }
-
-            return id;
-        }
-    }
-
     /**
      * Update special access for the association's package
      */
@@ -1403,20 +1156,27 @@
         final PackageInfo packageInfo =
                 getPackageInfo(getContext(), userId, packageName);
 
-        Binder.withCleanCallingIdentity(() -> updateSpecialAccessPermissionAsSystem(packageInfo));
+        Binder.withCleanCallingIdentity(() -> updateSpecialAccessPermissionAsSystem(packageInfo,
+                userId, packageName));
     }
 
-    private void updateSpecialAccessPermissionAsSystem(PackageInfo packageInfo) {
+    private void updateSpecialAccessPermissionAsSystem(PackageInfo packageInfo, int userId,
+            String packageName) {
         if (packageInfo == null) {
             return;
         }
+
+        List<AssociationInfo> associations = mAssociationStore.getActiveAssociationsByPackage(
+                userId, packageName);
+
         if (containsEither(packageInfo.requestedPermissions,
                 android.Manifest.permission.RUN_IN_BACKGROUND,
-                android.Manifest.permission.REQUEST_COMPANION_RUN_IN_BACKGROUND)) {
-            mPowerWhitelistManager.addToWhitelist(packageInfo.packageName);
+                android.Manifest.permission.REQUEST_COMPANION_RUN_IN_BACKGROUND)
+                && !associations.isEmpty()) {
+            mPowerExemptionManager.addToPermanentAllowList(packageInfo.packageName);
         } else {
             try {
-                mPowerWhitelistManager.removeFromWhitelist(packageInfo.packageName);
+                mPowerExemptionManager.removeFromPermanentAllowList(packageInfo.packageName);
             } catch (UnsupportedOperationException e) {
                 Slog.w(TAG, packageInfo.packageName + " can't be removed from power save"
                         + " whitelist. It might due to the package is whitelisted by the system.");
@@ -1427,7 +1187,8 @@
         try {
             if (containsEither(packageInfo.requestedPermissions,
                     android.Manifest.permission.USE_DATA_IN_BACKGROUND,
-                    android.Manifest.permission.REQUEST_COMPANION_USE_DATA_IN_BACKGROUND)) {
+                    android.Manifest.permission.REQUEST_COMPANION_USE_DATA_IN_BACKGROUND)
+                    && !associations.isEmpty()) {
                 networkPolicyManager.addUidPolicy(
                         packageInfo.applicationInfo.uid,
                         NetworkPolicyManager.POLICY_ALLOW_METERED_BACKGROUND);
@@ -1487,7 +1248,7 @@
 
             try {
                 final List<AssociationInfo> associations =
-                        mAssociationStore.getAssociationsForUser(userId);
+                        mAssociationStore.getActiveAssociationsByUser(userId);
                 for (AssociationInfo a : associations) {
                     try {
                         int uid = pm.getPackageUidAsUser(a.getPackageName(), userId);
@@ -1506,7 +1267,16 @@
             new AssociationStore.OnChangeListener() {
                 @Override
                 public void onAssociationChanged(int changeType, AssociationInfo association) {
-                    onAssociationChangedInternal(changeType, association);
+                    Slog.d(TAG, "onAssociationChanged changeType=[" + changeType
+                            + "], association=[" + association);
+
+                    final int userId = association.getUserId();
+                    final List<AssociationInfo> updatedAssociations =
+                            mAssociationStore.getActiveAssociationsByUser(userId);
+
+                    updateAtm(userId, updatedAssociations);
+                    updateSpecialAccessPermissionForAssociatedPackage(association.getUserId(),
+                            association.getPackageName());
                 }
             };
 
@@ -1634,64 +1404,4 @@
             }
         }
     }
-
-    /**
-     * This method must only be called from {@link CompanionDeviceShellCommand} for testing
-     * purposes only!
-     */
-    void persistState() {
-        mUserPersistenceHandler.clearMessages();
-        for (UserInfo user : mUserManager.getAliveUsers()) {
-            persistStateForUser(user.id);
-        }
-    }
-
-    /**
-     * This class is dedicated to handling requests to persist user state.
-     */
-    @SuppressLint("HandlerLeak")
-    private class PersistUserStateHandler extends Handler {
-        PersistUserStateHandler() {
-            super(BackgroundThread.get().getLooper());
-        }
-
-        /**
-         * Persists user state unless there is already an outstanding request for the given user.
-         */
-        synchronized void postPersistUserState(@UserIdInt int userId) {
-            if (!hasMessages(userId)) {
-                sendMessage(obtainMessage(userId));
-            }
-        }
-
-        /**
-         * Clears *ALL* outstanding persist requests for *ALL* users.
-         */
-        synchronized void clearMessages() {
-            removeCallbacksAndMessages(null);
-        }
-
-        @Override
-        public void handleMessage(@NonNull Message msg) {
-            final int userId = msg.what;
-            persistStateForUser(userId);
-        }
-    }
-
-    /**
-     * Persist associations
-     */
-    public void postPersistUserState(@UserIdInt int userId) {
-        mUserPersistenceHandler.postPersistUserState(userId);
-    }
-
-    /**
-     * Set to store associations
-     */
-    public static class PerUserAssociationSet extends PerUser<Set<AssociationInfo>> {
-        @Override
-        protected @NonNull Set<AssociationInfo> create(int userId) {
-            return new ArraySet<>();
-        }
-    }
 }
diff --git a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
index 16877dc..a7a73cb 100644
--- a/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
+++ b/services/companion/java/com/android/server/companion/CompanionDeviceShellCommand.java
@@ -33,8 +33,8 @@
 import android.util.proto.ProtoOutputStream;
 
 import com.android.server.companion.association.AssociationRequestsProcessor;
-import com.android.server.companion.association.AssociationRevokeProcessor;
 import com.android.server.companion.association.AssociationStore;
+import com.android.server.companion.association.DisassociationProcessor;
 import com.android.server.companion.datatransfer.SystemDataTransferProcessor;
 import com.android.server.companion.datatransfer.contextsync.BitmapUtils;
 import com.android.server.companion.datatransfer.contextsync.CrossDeviceSyncController;
@@ -49,7 +49,7 @@
     private static final String TAG = "CDM_CompanionDeviceShellCommand";
 
     private final CompanionDeviceManagerService mService;
-    private final AssociationRevokeProcessor mRevokeProcessor;
+    private final DisassociationProcessor mDisassociationProcessor;
     private final AssociationStore mAssociationStore;
     private final CompanionDevicePresenceMonitor mDevicePresenceMonitor;
     private final CompanionTransportManager mTransportManager;
@@ -65,7 +65,7 @@
             SystemDataTransferProcessor systemDataTransferProcessor,
             AssociationRequestsProcessor associationRequestsProcessor,
             BackupRestoreProcessor backupRestoreProcessor,
-            AssociationRevokeProcessor revokeProcessor) {
+            DisassociationProcessor disassociationProcessor) {
         mService = service;
         mAssociationStore = associationStore;
         mDevicePresenceMonitor = devicePresenceMonitor;
@@ -73,7 +73,7 @@
         mSystemDataTransferProcessor = systemDataTransferProcessor;
         mAssociationRequestsProcessor = associationRequestsProcessor;
         mBackupRestoreProcessor = backupRestoreProcessor;
-        mRevokeProcessor = revokeProcessor;
+        mDisassociationProcessor = disassociationProcessor;
     }
 
     @Override
@@ -105,12 +105,15 @@
                 case "list": {
                     final int userId = getNextIntArgRequired();
                     final List<AssociationInfo> associationsForUser =
-                            mAssociationStore.getAssociationsForUser(userId);
+                            mAssociationStore.getActiveAssociationsByUser(userId);
+                    final int maxId = mAssociationStore.getMaxId(userId);
+                    out.println("Max ID: " + maxId);
+                    out.println("Association ID | Package Name | Mac Address");
                     for (AssociationInfo association : associationsForUser) {
                         // TODO(b/212535524): use AssociationInfo.toShortString(), once it's not
                         //  longer referenced in tests.
-                        out.println(association.getPackageName() + " "
-                                + association.getDeviceMacAddress() + " " + association.getId());
+                        out.println(association.getId() + " | " + association.getPackageName()
+                                + " | " + association.getDeviceMacAddress());
                     }
                 }
                 break;
@@ -132,28 +135,24 @@
                     final String address = getNextArgRequired();
                     final AssociationInfo association =
                             mService.getAssociationWithCallerChecks(userId, packageName, address);
-                    if (association != null) {
-                        mRevokeProcessor.disassociateInternal(association.getId());
-                    }
+                    mDisassociationProcessor.disassociate(association.getId());
                 }
                 break;
 
                 case "disassociate-all": {
                     final int userId = getNextIntArgRequired();
-                    final String packageName = getNextArgRequired();
                     final List<AssociationInfo> userAssociations =
-                            mAssociationStore.getAssociationsForPackage(userId, packageName);
+                            mAssociationStore.getAssociationsByUser(userId);
                     for (AssociationInfo association : userAssociations) {
                         if (sanitizeWithCallerChecks(mService.getContext(), association) != null) {
-                            mRevokeProcessor.disassociateInternal(association.getId());
+                            mDisassociationProcessor.disassociate(association.getId());
                         }
                     }
                 }
                 break;
 
-                case "clear-association-memory-cache":
-                    mService.persistState();
-                    mService.loadAssociationsFromDisk();
+                case "refresh-cache":
+                    mAssociationStore.refreshCache();
                     break;
 
                 case "simulate-device-appeared":
diff --git a/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java b/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
index 75cb120..46d60f9 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationDiskStore.java
@@ -16,7 +16,6 @@
 
 package com.android.server.companion.association;
 
-import static com.android.internal.util.CollectionUtils.forEach;
 import static com.android.internal.util.XmlUtils.readBooleanAttribute;
 import static com.android.internal.util.XmlUtils.readIntAttribute;
 import static com.android.internal.util.XmlUtils.readLongAttribute;
@@ -26,7 +25,6 @@
 import static com.android.internal.util.XmlUtils.writeLongAttribute;
 import static com.android.internal.util.XmlUtils.writeStringAttribute;
 import static com.android.server.companion.utils.AssociationUtils.getFirstAssociationIdForUser;
-import static com.android.server.companion.utils.AssociationUtils.getLastAssociationIdForUser;
 import static com.android.server.companion.utils.DataStoreUtils.createStorageFileForUser;
 import static com.android.server.companion.utils.DataStoreUtils.fileToByteArray;
 import static com.android.server.companion.utils.DataStoreUtils.isEndOfTag;
@@ -40,10 +38,8 @@
 import android.companion.AssociationInfo;
 import android.net.MacAddress;
 import android.os.Environment;
-import android.util.ArrayMap;
 import android.util.AtomicFile;
 import android.util.Slog;
-import android.util.SparseArray;
 import android.util.Xml;
 
 import com.android.internal.util.XmlUtils;
@@ -59,11 +55,9 @@
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.util.Collection;
-import java.util.HashSet;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
@@ -82,8 +76,8 @@
  * <p>
  * Before Android T the data was stored using the v0 schema. See:
  * <ul>
- * <li>{@link #readAssociationsV0(TypedXmlPullParser, int, Collection) readAssociationsV0()}.
- * <li>{@link #readAssociationV0(TypedXmlPullParser, int, int, Collection) readAssociationV0()}.
+ * <li>{@link #readAssociationsV0(TypedXmlPullParser, int) readAssociationsV0()}.
+ * <li>{@link #readAssociationV0(TypedXmlPullParser, int, int) readAssociationV0()}.
  * </ul>
  *
  * The following snippet is a sample of a file that is using v0 schema.
@@ -116,15 +110,14 @@
  * optional.
  * <ul>
  * <li> {@link #CURRENT_PERSISTENCE_VERSION}
- * <li> {@link #readAssociationsV1(TypedXmlPullParser, int, Collection) readAssociationsV1()}
- * <li> {@link #readAssociationV1(TypedXmlPullParser, int, Collection) readAssociationV1()}
- * <li> {@link #readPreviouslyUsedIdsV1(TypedXmlPullParser, Map) readPreviouslyUsedIdsV1()}
+ * <li> {@link #readAssociationsV1(TypedXmlPullParser, int) readAssociationsV1()}
+ * <li> {@link #readAssociationV1(TypedXmlPullParser, int) readAssociationV1()}
  * </ul>
  *
  * The following snippet is a sample of a file that is using v1 schema.
  * <pre>{@code
  * <state persistence-version="1">
- *     <associations>
+ *     <associations max-id="3">
  *         <association
  *             id="1"
  *             package="com.sample.companion.app"
@@ -148,18 +141,12 @@
  *             time_approved="1634641160229"
  *             system_data_sync_flags="1"/>
  *     </associations>
- *
- *     <previously-used-ids>
- *         <package package_name="com.sample.companion.app">
- *             <id>2</id>
- *         </package>
- *     </previously-used-ids>
  * </state>
  * }</pre>
  */
 @SuppressLint("LongLogTag")
 public final class AssociationDiskStore {
-    private static final String TAG = "CompanionDevice_AssociationDiskStore";
+    private static final String TAG = "CDM_AssociationDiskStore";
 
     private static final int CURRENT_PERSISTENCE_VERSION = 1;
 
@@ -169,16 +156,11 @@
     private static final String XML_TAG_STATE = "state";
     private static final String XML_TAG_ASSOCIATIONS = "associations";
     private static final String XML_TAG_ASSOCIATION = "association";
-    private static final String XML_TAG_PREVIOUSLY_USED_IDS = "previously-used-ids";
-    private static final String XML_TAG_PACKAGE = "package";
     private static final String XML_TAG_TAG = "tag";
-    private static final String XML_TAG_ID = "id";
 
     private static final String XML_ATTR_PERSISTENCE_VERSION = "persistence-version";
+    private static final String XML_ATTR_MAX_ID = "max-id";
     private static final String XML_ATTR_ID = "id";
-    // Used in <package> elements, nested within <previously-used-ids> elements.
-    private static final String XML_ATTR_PACKAGE_NAME = "package_name";
-    // Used in <association> elements, nested within <associations> elements.
     private static final String XML_ATTR_PACKAGE = "package";
     private static final String XML_ATTR_MAC_ADDRESS = "mac_address";
     private static final String XML_ATTR_DISPLAY_NAME = "display_name";
@@ -199,38 +181,12 @@
     /**
      * Read all associations for given users
      */
-    public void readStateForUsers(@NonNull List<Integer> userIds,
-            @NonNull Set<AssociationInfo> allAssociationsOut,
-            @NonNull SparseArray<Map<String, Set<Integer>>> previouslyUsedIdsPerUserOut) {
+    public Map<Integer, Associations> readAssociationsByUsers(@NonNull List<Integer> userIds) {
+        Map<Integer, Associations> userToAssociationsMap = new HashMap<>();
         for (int userId : userIds) {
-            // Previously used IDs are stored in the "out" collection per-user.
-            final Map<String, Set<Integer>> previouslyUsedIds = new ArrayMap<>();
-
-            // Associations for all users are stored in a single "flat" set: so we read directly
-            // into it.
-            final Set<AssociationInfo> associationsForUser = new HashSet<>();
-            readStateForUser(userId, associationsForUser, previouslyUsedIds);
-
-            // Go through all the associations for the user and check if their IDs are within
-            // the allowed range (for the user).
-            final int firstAllowedId = getFirstAssociationIdForUser(userId);
-            final int lastAllowedId = getLastAssociationIdForUser(userId);
-            for (AssociationInfo association : associationsForUser) {
-                final int id = association.getId();
-                if (id < firstAllowedId || id > lastAllowedId) {
-                    Slog.e(TAG, "Wrong association ID assignment: " + id + ". "
-                            + "Association belongs to u" + userId + " and thus its ID should be "
-                            + "within [" + firstAllowedId + ", " + lastAllowedId + "] range.");
-                    // TODO(b/224736262): try fixing (re-assigning) the ID?
-                }
-            }
-
-            // Add user's association to the "output" set.
-            allAssociationsOut.addAll(associationsForUser);
-
-            // Save previously used IDs for this user into the "out" structure.
-            previouslyUsedIdsPerUserOut.append(userId, previouslyUsedIds);
+            userToAssociationsMap.put(userId, readAssociationsByUser(userId));
         }
+        return userToAssociationsMap;
     }
 
     /**
@@ -240,16 +196,12 @@
      * retrieval from this datastore because it is not persisted (by design). This means that
      * persisted data is not guaranteed to be identical to the initial data that was stored at the
      * time of association.
-     *
-     * @param userId Android UserID
-     * @param associationsOut a container to read the {@link AssociationInfo}s "into".
-     * @param previouslyUsedIdsPerPackageOut a container to read the used IDs "into".
      */
-    private void readStateForUser(@UserIdInt int userId,
-            @NonNull Collection<AssociationInfo> associationsOut,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
-        Slog.i(TAG, "Reading associations for user " + userId + " from disk");
+    @NonNull
+    private Associations readAssociationsByUser(@UserIdInt int userId) {
+        Slog.i(TAG, "Reading associations for user " + userId + " from disk.");
         final AtomicFile file = getStorageFileForUser(userId);
+        Associations associations;
 
         // getStorageFileForUser() ALWAYS returns the SAME OBJECT, which allows us to synchronize
         // accesses to the file on the file system using this AtomicFile object.
@@ -260,7 +212,7 @@
             if (!file.getBaseFile().exists()) {
                 legacyBaseFile = getBaseLegacyStorageFileForUser(userId);
                 if (!legacyBaseFile.exists()) {
-                    return;
+                    return new Associations();
                 }
 
                 readFrom = new AtomicFile(legacyBaseFile);
@@ -270,13 +222,12 @@
                 rootTag = XML_TAG_STATE;
             }
 
-            final int version = readStateFromFileLocked(userId, readFrom, rootTag,
-                    associationsOut, previouslyUsedIdsPerPackageOut);
+            associations = readAssociationsFromFile(userId, readFrom, rootTag);
 
-            if (legacyBaseFile != null || version < CURRENT_PERSISTENCE_VERSION) {
+            if (legacyBaseFile != null || associations.getVersion() < CURRENT_PERSISTENCE_VERSION) {
                 // The data is either in the legacy file or in the legacy format, or both.
                 // Save the data to right file in using the current format.
-                persistStateToFileLocked(file, associationsOut, previouslyUsedIdsPerPackageOut);
+                writeAssociationsToFile(file, associations);
 
                 if (legacyBaseFile != null) {
                     // We saved the data to the right file, can delete the old file now.
@@ -284,89 +235,75 @@
                 }
             }
         }
+        return associations;
     }
 
     /**
-     * Persisted data to the disk.
-     *
-     * Note that associatedDevice field in {@link AssociationInfo} is not persisted by this
-     * datastore implementation.
-     *
-     * @param userId Android UserID
-     * @param associations a set of user's associations.
-     * @param previouslyUsedIdsPerPackage a set previously used Association IDs for the user.
+     * Write associations to disk for the user.
      */
-    public void persistStateForUser(@UserIdInt int userId,
-            @NonNull Collection<AssociationInfo> associations,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) {
+    public void writeAssociationsForUser(@UserIdInt int userId,
+            @NonNull Associations associations) {
         Slog.i(TAG, "Writing associations for user " + userId + " to disk");
 
         final AtomicFile file = getStorageFileForUser(userId);
         // getStorageFileForUser() ALWAYS returns the SAME OBJECT, which allows us to synchronize
         // accesses to the file on the file system using this AtomicFile object.
         synchronized (file) {
-            persistStateToFileLocked(file, associations, previouslyUsedIdsPerPackage);
+            writeAssociationsToFile(file, associations);
         }
     }
 
-    private int readStateFromFileLocked(@UserIdInt int userId, @NonNull AtomicFile file,
-            @NonNull String rootTag, @Nullable Collection<AssociationInfo> associationsOut,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
+    @NonNull
+    private static Associations readAssociationsFromFile(@UserIdInt int userId,
+            @NonNull AtomicFile file, @NonNull String rootTag) {
         try (FileInputStream in = file.openRead()) {
-            return readStateFromInputStream(userId, in, rootTag, associationsOut,
-                    previouslyUsedIdsPerPackageOut);
+            return readAssociationsFromInputStream(userId, in, rootTag);
         } catch (XmlPullParserException | IOException e) {
             Slog.e(TAG, "Error while reading associations file", e);
-            return -1;
+            return new Associations();
         }
     }
 
-    private int readStateFromInputStream(@UserIdInt int userId, @NonNull InputStream in,
-            @NonNull String rootTag, @Nullable Collection<AssociationInfo> associationsOut,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut)
+    @NonNull
+    private static Associations readAssociationsFromInputStream(@UserIdInt int userId,
+            @NonNull InputStream in, @NonNull String rootTag)
             throws XmlPullParserException, IOException {
         final TypedXmlPullParser parser = Xml.resolvePullParser(in);
-
         XmlUtils.beginDocument(parser, rootTag);
+
         final int version = readIntAttribute(parser, XML_ATTR_PERSISTENCE_VERSION, 0);
+        Associations associations = new Associations();
+
         switch (version) {
             case 0:
-                readAssociationsV0(parser, userId, associationsOut);
+                associations = readAssociationsV0(parser, userId);
                 break;
             case 1:
                 while (true) {
                     parser.nextTag();
                     if (isStartOfTag(parser, XML_TAG_ASSOCIATIONS)) {
-                        readAssociationsV1(parser, userId, associationsOut);
-                    } else if (isStartOfTag(parser, XML_TAG_PREVIOUSLY_USED_IDS)) {
-                        readPreviouslyUsedIdsV1(parser, previouslyUsedIdsPerPackageOut);
+                        associations = readAssociationsV1(parser, userId);
                     } else if (isEndOfTag(parser, rootTag)) {
                         break;
                     }
                 }
                 break;
         }
-        return version;
+        return associations;
     }
 
-    private void persistStateToFileLocked(@NonNull AtomicFile file,
-            @Nullable Collection<AssociationInfo> associations,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) {
+    private void writeAssociationsToFile(@NonNull AtomicFile file,
+            @NonNull Associations associations) {
         // Writing to file could fail, for example, if the user has been recently removed and so was
         // their DE (/data/system_de/<user-id>/) directory.
         writeToFileSafely(file, out -> {
             final TypedXmlSerializer serializer = Xml.resolveSerializer(out);
-            serializer.setFeature(
-                    "http://xmlpull.org/v1/doc/features.html#indent-output", true);
-
+            serializer.setFeature("http://xmlpull.org/v1/doc/features.html#indent-output", true);
             serializer.startDocument(null, true);
             serializer.startTag(null, XML_TAG_STATE);
             writeIntAttribute(serializer,
                     XML_ATTR_PERSISTENCE_VERSION, CURRENT_PERSISTENCE_VERSION);
-
             writeAssociations(serializer, associations);
-            writePreviouslyUsedIds(serializer, previouslyUsedIdsPerPackage);
-
             serializer.endTag(null, XML_TAG_STATE);
             serializer.endDocument();
         });
@@ -379,7 +316,8 @@
      * IMPORTANT: the method will ALWAYS return the same {@link AtomicFile} object, which makes it
      * possible to synchronize reads and writes to the file using the returned object.
      */
-    private @NonNull AtomicFile getStorageFileForUser(@UserIdInt int userId) {
+    @NonNull
+    private AtomicFile getStorageFileForUser(@UserIdInt int userId) {
         return mUserIdToStorageFile.computeIfAbsent(userId,
                 u -> createStorageFileForUser(userId, FILE_NAME));
     }
@@ -399,14 +337,12 @@
     /**
      * Convert payload to a set of associations
      */
-    public void readStateFromPayload(byte[] payload, @UserIdInt int userId,
-                              @NonNull Set<AssociationInfo> associationsOut,
-                              @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackageOut) {
+    public static Associations readAssociationsFromPayload(byte[] payload, @UserIdInt int userId) {
         try (ByteArrayInputStream in = new ByteArrayInputStream(payload)) {
-            readStateFromInputStream(userId, in, XML_TAG_STATE, associationsOut,
-                    previouslyUsedIdsPerPackageOut);
+            return readAssociationsFromInputStream(userId, in, XML_TAG_STATE);
         } catch (XmlPullParserException | IOException e) {
             Slog.e(TAG, "Error while reading associations file", e);
+            return new Associations();
         }
     }
 
@@ -414,8 +350,8 @@
         return new File(Environment.getUserSystemDirectory(userId), FILE_NAME_LEGACY);
     }
 
-    private static void readAssociationsV0(@NonNull TypedXmlPullParser parser,
-            @UserIdInt int userId, @NonNull Collection<AssociationInfo> out)
+    private static Associations readAssociationsV0(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId)
             throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATIONS);
 
@@ -426,52 +362,70 @@
         // means that CDM hasn't assigned any IDs yet, so we can just start from the first available
         // id for each user (eg. 1 for user 0; 100 001 - for user 1; 200 001 - for user 2; etc).
         int associationId = getFirstAssociationIdForUser(userId);
+        Associations associations = new Associations();
+        associations.setVersion(0);
+
         while (true) {
             parser.nextTag();
             if (isEndOfTag(parser, XML_TAG_ASSOCIATIONS)) break;
             if (!isStartOfTag(parser, XML_TAG_ASSOCIATION)) continue;
 
-            readAssociationV0(parser, userId, associationId++, out);
+            associations.addAssociation(readAssociationV0(parser, userId, associationId++));
         }
+
+        associations.setMaxId(associationId - 1);
+
+        return associations;
     }
 
-    private static void readAssociationV0(@NonNull TypedXmlPullParser parser, @UserIdInt int userId,
-            int associationId, @NonNull Collection<AssociationInfo> out)
+    private static AssociationInfo readAssociationV0(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId, int associationId)
             throws XmlPullParserException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATION);
 
         final String appPackage = readStringAttribute(parser, XML_ATTR_PACKAGE);
         final String tag = readStringAttribute(parser, XML_TAG_TAG);
         final String deviceAddress = readStringAttribute(parser, LEGACY_XML_ATTR_DEVICE);
-
-        if (appPackage == null || deviceAddress == null) return;
-
         final String profile = readStringAttribute(parser, XML_ATTR_PROFILE);
         final boolean notify = readBooleanAttribute(parser, XML_ATTR_NOTIFY_DEVICE_NEARBY);
         final long timeApproved = readLongAttribute(parser, XML_ATTR_TIME_APPROVED, 0L);
 
-        out.add(new AssociationInfo(associationId, userId, appPackage, tag,
+        return new AssociationInfo(associationId, userId, appPackage, tag,
                 MacAddress.fromString(deviceAddress), null, profile, null,
                 /* managedByCompanionApp */ false, notify, /* revoked */ false, /* pending */ false,
-                timeApproved, Long.MAX_VALUE, /* systemDataSyncFlags */ 0));
+                timeApproved, Long.MAX_VALUE, /* systemDataSyncFlags */ 0);
     }
 
-    private static void readAssociationsV1(@NonNull TypedXmlPullParser parser,
-            @UserIdInt int userId, @NonNull Collection<AssociationInfo> out)
+    private static Associations readAssociationsV1(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId)
             throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATIONS);
 
+        // For old builds that don't have max-id attr,
+        // default maxId to 0 and get the maxId out of all association ids.
+        int maxId = readIntAttribute(parser, XML_ATTR_MAX_ID, 0);
+        Associations associations = new Associations();
+        associations.setVersion(1);
+
         while (true) {
             parser.nextTag();
             if (isEndOfTag(parser, XML_TAG_ASSOCIATIONS)) break;
             if (!isStartOfTag(parser, XML_TAG_ASSOCIATION)) continue;
 
-            readAssociationV1(parser, userId, out);
+            AssociationInfo association = readAssociationV1(parser, userId);
+            associations.addAssociation(association);
+
+            maxId = Math.max(maxId, association.getId());
         }
+
+        associations.setMaxId(maxId);
+
+        return associations;
     }
 
-    private static void readAssociationV1(@NonNull TypedXmlPullParser parser, @UserIdInt int userId,
-            @NonNull Collection<AssociationInfo> out) throws XmlPullParserException, IOException {
+    private static AssociationInfo readAssociationV1(@NonNull TypedXmlPullParser parser,
+            @UserIdInt int userId)
+            throws XmlPullParserException, IOException {
         requireStartOfTag(parser, XML_TAG_ASSOCIATION);
 
         final int associationId = readIntAttribute(parser, XML_ATTR_ID);
@@ -491,46 +445,19 @@
         final int systemDataSyncFlags = readIntAttribute(parser,
                 XML_ATTR_SYSTEM_DATA_SYNC_FLAGS, 0);
 
-        final AssociationInfo associationInfo = createAssociationInfoNoThrow(associationId, userId,
-                appPackage, tag, macAddress, displayName, profile, selfManaged, notify, revoked,
-                pending, timeApproved, lastTimeConnected, systemDataSyncFlags);
-        if (associationInfo != null) {
-            out.add(associationInfo);
-        }
-    }
-
-    private static void readPreviouslyUsedIdsV1(@NonNull TypedXmlPullParser parser,
-            @NonNull Map<String, Set<Integer>> out) throws XmlPullParserException, IOException {
-        requireStartOfTag(parser, XML_TAG_PREVIOUSLY_USED_IDS);
-
-        while (true) {
-            parser.nextTag();
-            if (isEndOfTag(parser, XML_TAG_PREVIOUSLY_USED_IDS)) break;
-            if (!isStartOfTag(parser, XML_TAG_PACKAGE)) continue;
-
-            final String packageName = readStringAttribute(parser, XML_ATTR_PACKAGE_NAME);
-            final Set<Integer> usedIds = new HashSet<>();
-
-            while (true) {
-                parser.nextTag();
-                if (isEndOfTag(parser, XML_TAG_PACKAGE)) break;
-                if (!isStartOfTag(parser, XML_TAG_ID)) continue;
-
-                parser.nextToken();
-                final int id = Integer.parseInt(parser.getText());
-                usedIds.add(id);
-            }
-
-            out.put(packageName, usedIds);
-        }
+        return new AssociationInfo(associationId, userId, appPackage, tag, macAddress, displayName,
+                profile, null, selfManaged, notify, revoked, pending, timeApproved,
+                lastTimeConnected, systemDataSyncFlags);
     }
 
     private static void writeAssociations(@NonNull XmlSerializer parent,
-            @Nullable Collection<AssociationInfo> associations) throws IOException {
+            @NonNull Associations associations)
+            throws IOException {
         final XmlSerializer serializer = parent.startTag(null, XML_TAG_ASSOCIATIONS);
-        for (AssociationInfo association : associations) {
+        for (AssociationInfo association : associations.getAssociations()) {
             writeAssociation(serializer, association);
         }
+        writeIntAttribute(serializer, XML_ATTR_MAX_ID, associations.getMaxId());
         serializer.endTag(null, XML_TAG_ASSOCIATIONS);
     }
 
@@ -557,26 +484,6 @@
         serializer.endTag(null, XML_TAG_ASSOCIATION);
     }
 
-    private static void writePreviouslyUsedIds(@NonNull XmlSerializer parent,
-            @NonNull Map<String, Set<Integer>> previouslyUsedIdsPerPackage) throws IOException {
-        final XmlSerializer serializer = parent.startTag(null, XML_TAG_PREVIOUSLY_USED_IDS);
-        for (Map.Entry<String, Set<Integer>> entry : previouslyUsedIdsPerPackage.entrySet()) {
-            writePreviouslyUsedIdsForPackage(serializer, entry.getKey(), entry.getValue());
-        }
-        serializer.endTag(null, XML_TAG_PREVIOUSLY_USED_IDS);
-    }
-
-    private static void writePreviouslyUsedIdsForPackage(@NonNull XmlSerializer parent,
-            @NonNull String packageName, @NonNull Set<Integer> previouslyUsedIds)
-            throws IOException {
-        final XmlSerializer serializer = parent.startTag(null, XML_TAG_PACKAGE);
-        writeStringAttribute(serializer, XML_ATTR_PACKAGE_NAME, packageName);
-        forEach(previouslyUsedIds, id -> serializer.startTag(null, XML_TAG_ID)
-                .text(Integer.toString(id))
-                .endTag(null, XML_TAG_ID));
-        serializer.endTag(null, XML_TAG_PACKAGE);
-    }
-
     private static void requireStartOfTag(@NonNull XmlPullParser parser, @NonNull String tag)
             throws XmlPullParserException {
         if (isStartOfTag(parser, tag)) return;
@@ -587,22 +494,4 @@
     private static @Nullable MacAddress stringToMacAddress(@Nullable String address) {
         return address != null ? MacAddress.fromString(address) : null;
     }
-
-    private static AssociationInfo createAssociationInfoNoThrow(int associationId,
-            @UserIdInt int userId, @NonNull String appPackage, @Nullable String tag,
-            @Nullable MacAddress macAddress, @Nullable CharSequence displayName,
-            @Nullable String profile, boolean selfManaged, boolean notify, boolean revoked,
-            boolean pending, long timeApproved, long lastTimeConnected, int systemDataSyncFlags) {
-        AssociationInfo associationInfo = null;
-        try {
-            // We do not persist AssociatedDevice, which means that AssociationInfo retrieved from
-            // datastore is not guaranteed to be identical to the one from initial association.
-            associationInfo = new AssociationInfo(associationId, userId, appPackage, tag,
-                    macAddress, displayName, profile, null, selfManaged, notify,
-                    revoked, pending, timeApproved, lastTimeConnected, systemDataSyncFlags);
-        } catch (Exception e) {
-            Slog.e(TAG, "Could not create AssociationInfo", e);
-        }
-        return associationInfo;
-    }
 }
diff --git a/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java b/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
index 29ec7c2..a02d9f9 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationRequestsProcessor.java
@@ -24,7 +24,6 @@
 import static android.content.ComponentName.createRelative;
 import static android.content.pm.PackageManager.FEATURE_WATCH;
 
-import static com.android.server.companion.utils.MetricUtils.logCreateAssociation;
 import static com.android.server.companion.utils.PackageUtils.enforceUsesCompanionDeviceFeature;
 import static com.android.server.companion.utils.PermissionsUtils.enforcePermissionForCreatingAssociation;
 import static com.android.server.companion.utils.RolesUtils.addRoleHolderForAssociation;
@@ -128,17 +127,16 @@
     private static final long ASSOCIATE_WITHOUT_PROMPT_WINDOW_MS = 60 * 60 * 1000; // 60 min;
 
     private final @NonNull Context mContext;
-    private final @NonNull CompanionDeviceManagerService mService;
-    private final @NonNull PackageManagerInternal mPackageManager;
+    private final @NonNull PackageManagerInternal mPackageManagerInternal;
     private final @NonNull AssociationStore mAssociationStore;
     @NonNull
     private final ComponentName mCompanionDeviceActivity;
 
-    public AssociationRequestsProcessor(@NonNull CompanionDeviceManagerService service,
+    public AssociationRequestsProcessor(@NonNull Context context,
+            @NonNull PackageManagerInternal packageManagerInternal,
             @NonNull AssociationStore associationStore) {
-        mContext = service.getContext();
-        mService = service;
-        mPackageManager = service.mPackageManagerInternal;
+        mContext = context;
+        mPackageManagerInternal = packageManagerInternal;
         mAssociationStore = associationStore;
         mCompanionDeviceActivity = createRelative(
                 mContext.getString(R.string.config_companionDeviceManagerPackage),
@@ -160,7 +158,7 @@
         requireNonNull(packageName, "Package name MUST NOT be null");
         requireNonNull(callback, "Callback MUST NOT be null");
 
-        final int packageUid = mPackageManager.getPackageUid(packageName, 0, userId);
+        final int packageUid = mPackageManagerInternal.getPackageUid(packageName, 0, userId);
         Slog.d(TAG, "processNewAssociationRequest() " + "request=" + request + ", " + "package=u"
                 + userId + "/" + packageName + " (uid=" + packageUid + ")");
 
@@ -226,7 +224,7 @@
 
         enforceUsesCompanionDeviceFeature(mContext, userId, packageName);
 
-        final int packageUid = mPackageManager.getPackageUid(packageName, 0, userId);
+        final int packageUid = mPackageManagerInternal.getPackageUid(packageName, 0, userId);
 
         final Bundle extras = new Bundle();
         extras.putBoolean(EXTRA_FORCE_CANCEL_CONFIRMATION, true);
@@ -243,7 +241,7 @@
             @NonNull ResultReceiver resultReceiver, @Nullable MacAddress macAddress) {
         final String packageName = request.getPackageName();
         final int userId = request.getUserId();
-        final int packageUid = mPackageManager.getPackageUid(packageName, 0, userId);
+        final int packageUid = mPackageManagerInternal.getPackageUid(packageName, 0, userId);
 
         // 1. Need to check permissions again in case something changed, since we first received
         // this request.
@@ -267,15 +265,12 @@
             @NonNull AssociationRequest request, @NonNull String packageName, @UserIdInt int userId,
             @Nullable MacAddress macAddress, @NonNull IAssociationRequestCallback callback,
             @NonNull ResultReceiver resultReceiver) {
-        final long callingIdentity = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             createAssociation(userId, packageName, macAddress, request.getDisplayName(),
                     request.getDeviceProfile(), request.getAssociatedDevice(),
                     request.isSelfManaged(),
                     callback, resultReceiver);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentity);
-        }
+        });
     }
 
     /**
@@ -286,7 +281,7 @@
             @Nullable String deviceProfile, @Nullable AssociatedDevice associatedDevice,
             boolean selfManaged, @Nullable IAssociationRequestCallback callback,
             @Nullable ResultReceiver resultReceiver) {
-        final int id = mService.getNewAssociationIdForPackage(userId, packageName);
+        final int id = mAssociationStore.getNextId(userId);
         final long timestamp = System.currentTimeMillis();
 
         final AssociationInfo association = new AssociationInfo(id, userId, packageName,
@@ -296,10 +291,6 @@
 
         // Add role holder for association (if specified) and add new association to store.
         maybeGrantRoleAndStoreAssociation(association, callback, resultReceiver);
-
-        // Don't need to update the mRevokedAssociationsPendingRoleHolderRemoval since
-        // maybeRemoveRoleHolderForAssociation in PackageInactivityListener will handle the case
-        // that there are other devices with the same profile, so the role holder won't be removed.
     }
 
     /**
@@ -311,12 +302,12 @@
         // If the "Device Profile" is specified, make the companion application a holder of the
         // corresponding role.
         // If it is null, then the operation will succeed without granting any role.
-        addRoleHolderForAssociation(mService.getContext(), association, success -> {
+        addRoleHolderForAssociation(mContext, association, success -> {
             if (success) {
                 Slog.i(TAG, "Added " + association.getDeviceProfile() + " role to userId="
                         + association.getUserId() + ", packageName="
                         + association.getPackageName());
-                addAssociationToStore(association);
+                mAssociationStore.addAssociation(association);
                 sendCallbackAndFinish(association, callback, resultReceiver);
             } else {
                 Slog.e(TAG, "Failed to add u" + association.getUserId()
@@ -347,17 +338,6 @@
         mAssociationStore.updateAssociation(updated);
     }
 
-    private void addAssociationToStore(@NonNull AssociationInfo association) {
-        Slog.i(TAG, "New CDM association created=" + association);
-
-        mAssociationStore.addAssociation(association);
-
-        mService.updateSpecialAccessPermissionForAssociatedPackage(association.getUserId(),
-                association.getPackageName());
-
-        logCreateAssociation(association.getDeviceProfile());
-    }
-
     private void sendCallbackAndFinish(@Nullable AssociationInfo association,
             @Nullable IAssociationRequestCallback callback,
             @Nullable ResultReceiver resultReceiver) {
@@ -409,27 +389,22 @@
 
     private PendingIntent createPendingIntent(int packageUid, Intent intent) {
         final PendingIntent pendingIntent;
-        final long token = Binder.clearCallingIdentity();
 
         // Using uid of the application that will own the association (usually the same
         // application that sent the request) allows us to have multiple "pending" association
         // requests at the same time.
         // If the application already has a pending association request, that PendingIntent
         // will be cancelled except application wants to cancel the request by the system.
-        try {
-            pendingIntent = PendingIntent.getActivityAsUser(
+        return Binder.withCleanCallingIdentity(() ->
+            PendingIntent.getActivityAsUser(
                     mContext, /*requestCode */ packageUid, intent,
                     FLAG_ONE_SHOT | FLAG_CANCEL_CURRENT | FLAG_IMMUTABLE,
                     ActivityOptions.makeBasic()
                             .setPendingIntentCreatorBackgroundActivityStartMode(
                                     ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED)
                             .toBundle(),
-                    UserHandle.CURRENT);
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
-
-        return pendingIntent;
+                    UserHandle.CURRENT)
+        );
     }
 
     private final ResultReceiver mOnRequestConfirmationReceiver =
@@ -470,7 +445,7 @@
         // Throttle frequent associations
         final long now = System.currentTimeMillis();
         final List<AssociationInfo> associationForPackage =
-                mAssociationStore.getAssociationsForPackage(userId, packageName);
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName);
         // Number of "recent" associations.
         int recent = 0;
         for (AssociationInfo association : associationForPackage) {
@@ -486,6 +461,6 @@
             }
         }
 
-        return PackageUtils.isPackageAllowlisted(mContext, mPackageManager, packageName);
+        return PackageUtils.isPackageAllowlisted(mContext, mPackageManagerInternal, packageName);
     }
 }
diff --git a/services/companion/java/com/android/server/companion/association/AssociationRevokeProcessor.java b/services/companion/java/com/android/server/companion/association/AssociationRevokeProcessor.java
deleted file mode 100644
index d1efbbc..0000000
--- a/services/companion/java/com/android/server/companion/association/AssociationRevokeProcessor.java
+++ /dev/null
@@ -1,383 +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.server.companion.association;
-
-import static android.app.ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE;
-import static android.companion.AssociationRequest.DEVICE_PROFILE_AUTOMOTIVE_PROJECTION;
-
-import static com.android.internal.util.CollectionUtils.any;
-import static com.android.server.companion.utils.MetricUtils.logRemoveAssociation;
-import static com.android.server.companion.utils.RolesUtils.removeRoleHolderForAssociation;
-import static com.android.server.companion.CompanionDeviceManagerService.PerUserAssociationSet;
-
-import android.annotation.NonNull;
-import android.annotation.SuppressLint;
-import android.annotation.UserIdInt;
-import android.app.ActivityManager;
-import android.companion.AssociationInfo;
-import android.content.Context;
-import android.content.pm.PackageManagerInternal;
-import android.os.Binder;
-import android.os.UserHandle;
-import android.util.ArraySet;
-import android.util.Log;
-import android.util.Slog;
-
-import com.android.internal.annotations.GuardedBy;
-import com.android.server.companion.CompanionApplicationController;
-import com.android.server.companion.CompanionDeviceManagerService;
-import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
-import com.android.server.companion.presence.CompanionDevicePresenceMonitor;
-import com.android.server.companion.transport.CompanionTransportManager;
-
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Set;
-
-/**
- * A class response for Association removal.
- */
-@SuppressLint("LongLogTag")
-public class AssociationRevokeProcessor {
-
-    private static final String TAG = "CDM_AssociationRevokeProcessor";
-    private static final boolean DEBUG = false;
-    private final @NonNull Context mContext;
-    private final @NonNull CompanionDeviceManagerService mService;
-    private final @NonNull AssociationStore mAssociationStore;
-    private final @NonNull PackageManagerInternal mPackageManagerInternal;
-    private final @NonNull CompanionDevicePresenceMonitor mDevicePresenceMonitor;
-    private final @NonNull SystemDataTransferRequestStore mSystemDataTransferRequestStore;
-    private final @NonNull CompanionApplicationController mCompanionAppController;
-    private final @NonNull CompanionTransportManager mTransportManager;
-    private final OnPackageVisibilityChangeListener mOnPackageVisibilityChangeListener;
-    private final ActivityManager mActivityManager;
-
-    /**
-     * A structure that consists of a set of revoked associations that pending for role holder
-     * removal per each user.
-     *
-     * @see #maybeRemoveRoleHolderForAssociation(AssociationInfo)
-     * @see #addToPendingRoleHolderRemoval(AssociationInfo)
-     * @see #removeFromPendingRoleHolderRemoval(AssociationInfo)
-     * @see #getPendingRoleHolderRemovalAssociationsForUser(int)
-     */
-    @GuardedBy("mRevokedAssociationsPendingRoleHolderRemoval")
-    private final PerUserAssociationSet mRevokedAssociationsPendingRoleHolderRemoval =
-            new PerUserAssociationSet();
-    /**
-     * Contains uid-s of packages pending to be removed from the role holder list (after
-     * revocation of an association), which will happen one the package is no longer visible to the
-     * user.
-     * For quicker uid -> (userId, packageName) look-up this is not a {@code Set<Integer>} but
-     * a {@code Map<Integer, String>} which maps uid-s to packageName-s (userId-s can be derived
-     * from uid-s using {@link UserHandle#getUserId(int)}).
-     *
-     * @see #maybeRemoveRoleHolderForAssociation(AssociationInfo)
-     * @see #addToPendingRoleHolderRemoval(AssociationInfo)
-     * @see #removeFromPendingRoleHolderRemoval(AssociationInfo)
-     */
-    @GuardedBy("mRevokedAssociationsPendingRoleHolderRemoval")
-    private final Map<Integer, String> mUidsPendingRoleHolderRemoval = new HashMap<>();
-
-    public AssociationRevokeProcessor(@NonNull CompanionDeviceManagerService service,
-            @NonNull AssociationStore associationStore,
-            @NonNull PackageManagerInternal packageManager,
-            @NonNull CompanionDevicePresenceMonitor devicePresenceMonitor,
-            @NonNull CompanionApplicationController applicationController,
-            @NonNull SystemDataTransferRequestStore systemDataTransferRequestStore,
-            @NonNull CompanionTransportManager companionTransportManager) {
-        mService = service;
-        mContext = service.getContext();
-        mActivityManager = mContext.getSystemService(ActivityManager.class);
-        mAssociationStore = associationStore;
-        mPackageManagerInternal = packageManager;
-        mOnPackageVisibilityChangeListener =
-                new OnPackageVisibilityChangeListener(mActivityManager);
-        mDevicePresenceMonitor = devicePresenceMonitor;
-        mCompanionAppController = applicationController;
-        mSystemDataTransferRequestStore = systemDataTransferRequestStore;
-        mTransportManager = companionTransportManager;
-    }
-
-    /**
-     * Disassociate an association
-     */
-    // TODO: also revoke notification access
-    public void disassociateInternal(int associationId) {
-        final AssociationInfo association = mAssociationStore.getAssociationById(associationId);
-        final int userId = association.getUserId();
-        final String packageName = association.getPackageName();
-        final String deviceProfile = association.getDeviceProfile();
-
-        // Detach transport if exists
-        mTransportManager.detachSystemDataTransport(packageName, userId, associationId);
-
-        if (!maybeRemoveRoleHolderForAssociation(association)) {
-            // Need to remove the app from list of the role holders, but will have to do it later
-            // (the app is in foreground at the moment).
-            addToPendingRoleHolderRemoval(association);
-        }
-
-        // Need to check if device still present now because CompanionDevicePresenceMonitor will
-        // remove current connected device after mAssociationStore.removeAssociation
-        final boolean wasPresent = mDevicePresenceMonitor.isDevicePresent(associationId);
-
-        // Removing the association.
-        mAssociationStore.removeAssociation(associationId);
-        // Do not need to persistUserState since CompanionDeviceManagerService will get callback
-        // from #onAssociationChanged, and it will handle the persistUserState which including
-        // active and revoked association.
-        logRemoveAssociation(deviceProfile);
-
-        // Remove all the system data transfer requests for the association.
-        mSystemDataTransferRequestStore.removeRequestsByAssociationId(userId, associationId);
-
-        if (!wasPresent || !association.isNotifyOnDeviceNearby()) return;
-        // The device was connected and the app was notified: check if we need to unbind the app
-        // now.
-        final boolean shouldStayBound = any(
-                mAssociationStore.getAssociationsForPackage(userId, packageName),
-                it -> it.isNotifyOnDeviceNearby()
-                        && mDevicePresenceMonitor.isDevicePresent(it.getId()));
-        if (shouldStayBound) return;
-        mCompanionAppController.unbindCompanionApplication(userId, packageName);
-    }
-
-    /**
-     * First, checks if the companion application should be removed from the list role holders when
-     * upon association's removal, i.e.: association's profile (matches the role) is not null,
-     * the application does not have other associations with the same profile, etc.
-     *
-     * <p>
-     * Then, if establishes that the application indeed has to be removed from the list of the role
-     * holders, checks if it could be done right now -
-     * {@link android.app.role.RoleManager#removeRoleHolderAsUser(String, String, int, UserHandle, java.util.concurrent.Executor, java.util.function.Consumer) RoleManager#removeRoleHolderAsUser()}
-     * will kill the application's process, which leads poor user experience if the application was
-     * in foreground when this happened, to avoid this CDMS delays invoking
-     * {@code RoleManager.removeRoleHolderAsUser()} until the app is no longer in foreground.
-     *
-     * @return {@code true} if the application does NOT need be removed from the list of the role
-     *         holders OR if the application was successfully removed from the list of role holders.
-     *         I.e.: from the role-management perspective the association is done with.
-     *         {@code false} if the application needs to be removed from the list of role the role
-     *         holders, BUT it CDMS would prefer to do it later.
-     *         I.e.: application is in the foreground at the moment, but invoking
-     *         {@code RoleManager.removeRoleHolderAsUser()} will kill the application's process,
-     *         which would lead to the poor UX, hence need to try later.
-     */
-    public boolean maybeRemoveRoleHolderForAssociation(@NonNull AssociationInfo association) {
-        if (DEBUG) Log.d(TAG, "maybeRemoveRoleHolderForAssociation() association=" + association);
-        final String deviceProfile = association.getDeviceProfile();
-
-        if (deviceProfile == null) {
-            // No role was granted to for this association, there is nothing else we need to here.
-            return true;
-        }
-        // Do not need to remove the system role since it was pre-granted by the system.
-        if (deviceProfile.equals(DEVICE_PROFILE_AUTOMOTIVE_PROJECTION)) {
-            return true;
-        }
-
-        // Check if the applications is associated with another devices with the profile. If so,
-        // it should remain the role holder.
-        final int id = association.getId();
-        final int userId = association.getUserId();
-        final String packageName = association.getPackageName();
-        final boolean roleStillInUse = any(
-                mAssociationStore.getAssociationsForPackage(userId, packageName),
-                it -> deviceProfile.equals(it.getDeviceProfile()) && id != it.getId());
-        if (roleStillInUse) {
-            // Application should remain a role holder, there is nothing else we need to here.
-            return true;
-        }
-
-        final int packageProcessImportance = getPackageProcessImportance(userId, packageName);
-        if (packageProcessImportance <= IMPORTANCE_VISIBLE) {
-            // Need to remove the app from the list of role holders, but the process is visible to
-            // the user at the moment, so we'll need to it later: log and return false.
-            Slog.i(TAG, "Cannot remove role holder for the removed association id=" + id
-                    + " now - process is visible.");
-            return false;
-        }
-
-        removeRoleHolderForAssociation(mContext, association.getUserId(),
-                association.getPackageName(), association.getDeviceProfile());
-        return true;
-    }
-
-    /**
-     * Set revoked flag for active association and add the revoked association and the uid into
-     * the caches.
-     *
-     * @see #mRevokedAssociationsPendingRoleHolderRemoval
-     * @see #mUidsPendingRoleHolderRemoval
-     * @see OnPackageVisibilityChangeListener
-     */
-    public void addToPendingRoleHolderRemoval(@NonNull AssociationInfo association) {
-        // First: set revoked flag
-        association = (new AssociationInfo.Builder(association)).setRevoked(true).build();
-        final String packageName = association.getPackageName();
-        final int userId = association.getUserId();
-        final int uid = mPackageManagerInternal.getPackageUid(packageName, /* flags */0, userId);
-        // Second: add to the set.
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            mRevokedAssociationsPendingRoleHolderRemoval.forUser(association.getUserId())
-                    .add(association);
-            if (!mUidsPendingRoleHolderRemoval.containsKey(uid)) {
-                mUidsPendingRoleHolderRemoval.put(uid, packageName);
-
-                if (mUidsPendingRoleHolderRemoval.size() == 1) {
-                    // Just added first uid: start the listener
-                    mOnPackageVisibilityChangeListener.startListening();
-                }
-            }
-        }
-    }
-
-    /**
-     * @return a copy of the revoked associations set (safeguarding against
-     *         {@code ConcurrentModificationException}-s).
-     */
-    @NonNull
-    public Set<AssociationInfo> getPendingRoleHolderRemovalAssociationsForUser(
-            @UserIdInt int userId) {
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            // Return a copy.
-            return new ArraySet<>(mRevokedAssociationsPendingRoleHolderRemoval.forUser(userId));
-        }
-    }
-
-    @SuppressLint("MissingPermission")
-    private int  getPackageProcessImportance(@UserIdInt int userId, @NonNull String packageName) {
-        return Binder.withCleanCallingIdentity(() -> {
-            final int uid =
-                    mPackageManagerInternal.getPackageUid(packageName, /* flags */0, userId);
-            return mActivityManager.getUidImportance(uid);
-        });
-    }
-
-    /**
-     * Remove the revoked association from the cache and also remove the uid from the map if
-     * there are other associations with the same package still pending for role holder removal.
-     *
-     * @see #mRevokedAssociationsPendingRoleHolderRemoval
-     * @see #mUidsPendingRoleHolderRemoval
-     * @see OnPackageVisibilityChangeListener
-     */
-    private void removeFromPendingRoleHolderRemoval(@NonNull AssociationInfo association) {
-        final String packageName = association.getPackageName();
-        final int userId = association.getUserId();
-        final int uid = mPackageManagerInternal.getPackageUid(packageName, /* flags */  0, userId);
-
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            mRevokedAssociationsPendingRoleHolderRemoval.forUser(userId)
-                    .remove(association);
-
-            final boolean shouldKeepUidForRemoval = any(
-                    getPendingRoleHolderRemovalAssociationsForUser(userId),
-                    ai -> packageName.equals(ai.getPackageName()));
-            // Do not remove the uid from the map since other associations with
-            // the same packageName still pending for role holder removal.
-            if (!shouldKeepUidForRemoval) {
-                mUidsPendingRoleHolderRemoval.remove(uid);
-            }
-
-            if (mUidsPendingRoleHolderRemoval.isEmpty()) {
-                // The set is empty now - can "turn off" the listener.
-                mOnPackageVisibilityChangeListener.stopListening();
-            }
-        }
-    }
-
-    private String getPackageNameByUid(int uid) {
-        synchronized (mRevokedAssociationsPendingRoleHolderRemoval) {
-            return mUidsPendingRoleHolderRemoval.get(uid);
-        }
-    }
-
-    /**
-     * An OnUidImportanceListener class which watches the importance of the packages.
-     * In this class, we ONLY interested in the importance of the running process is greater than
-     * {@link ActivityManager.RunningAppProcessInfo#IMPORTANCE_VISIBLE} for the uids have been added
-     * into the {@link #mUidsPendingRoleHolderRemoval}. Lastly remove the role holder for the
-     * revoked associations for the same packages.
-     *
-     * @see #maybeRemoveRoleHolderForAssociation(AssociationInfo)
-     * @see #removeFromPendingRoleHolderRemoval(AssociationInfo)
-     * @see #getPendingRoleHolderRemovalAssociationsForUser(int)
-     */
-    private class OnPackageVisibilityChangeListener implements
-            ActivityManager.OnUidImportanceListener {
-        final @NonNull ActivityManager mAm;
-
-        OnPackageVisibilityChangeListener(@NonNull ActivityManager am) {
-            this.mAm = am;
-        }
-
-        @SuppressLint("MissingPermission")
-        void startListening() {
-            Binder.withCleanCallingIdentity(
-                    () -> mAm.addOnUidImportanceListener(
-                            /* listener */ OnPackageVisibilityChangeListener.this,
-                            ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE));
-        }
-
-        @SuppressLint("MissingPermission")
-        void stopListening() {
-            Binder.withCleanCallingIdentity(
-                    () -> mAm.removeOnUidImportanceListener(
-                            /* listener */ OnPackageVisibilityChangeListener.this));
-        }
-
-        @Override
-        public void onUidImportance(int uid, int importance) {
-            if (importance <= ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE) {
-                // The lower the importance value the more "important" the process is.
-                // We are only interested when the process ceases to be visible.
-                return;
-            }
-
-            final String packageName = getPackageNameByUid(uid);
-            if (packageName == null) {
-                // Not interested in this uid.
-                return;
-            }
-
-            final int userId = UserHandle.getUserId(uid);
-
-            boolean needToPersistStateForUser = false;
-
-            for (AssociationInfo association :
-                    getPendingRoleHolderRemovalAssociationsForUser(userId)) {
-                if (!packageName.equals(association.getPackageName())) continue;
-
-                if (!maybeRemoveRoleHolderForAssociation(association)) {
-                    // Did not remove the role holder, will have to try again later.
-                    continue;
-                }
-
-                removeFromPendingRoleHolderRemoval(association);
-                needToPersistStateForUser = true;
-            }
-
-            if (needToPersistStateForUser) {
-                mService.postPersistUserState(userId);
-            }
-        }
-    }
-}
diff --git a/services/companion/java/com/android/server/companion/association/AssociationStore.java b/services/companion/java/com/android/server/companion/association/AssociationStore.java
index 2f94bde..29de764 100644
--- a/services/companion/java/com/android/server/companion/association/AssociationStore.java
+++ b/services/companion/java/com/android/server/companion/association/AssociationStore.java
@@ -16,15 +16,24 @@
 
 package com.android.server.companion.association;
 
+import static com.android.server.companion.utils.MetricUtils.logCreateAssociation;
+import static com.android.server.companion.utils.MetricUtils.logRemoveAssociation;
+
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.annotation.UserIdInt;
 import android.companion.AssociationInfo;
+import android.companion.IOnAssociationsChangedListener;
+import android.content.pm.UserInfo;
 import android.net.MacAddress;
+import android.os.Binder;
+import android.os.RemoteCallbackList;
+import android.os.RemoteException;
+import android.os.UserHandle;
+import android.os.UserManager;
 import android.util.Slog;
-import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.CollectionUtils;
@@ -33,15 +42,14 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 /**
  * Association store for CRUD.
@@ -109,44 +117,105 @@
 
     private final Object mLock = new Object();
 
-    @GuardedBy("mLock")
-    private final Map<Integer, AssociationInfo> mIdMap = new HashMap<>();
-    @GuardedBy("mLock")
-    private final Map<MacAddress, Set<Integer>> mAddressMap = new HashMap<>();
-    @GuardedBy("mLock")
-    private final SparseArray<List<AssociationInfo>> mCachedPerUser = new SparseArray<>();
+    private final ExecutorService mExecutor;
 
-    @GuardedBy("mListeners")
-    private final Set<OnChangeListener> mListeners = new LinkedHashSet<>();
+    @GuardedBy("mLock")
+    private boolean mPersisted = false;
+    @GuardedBy("mLock")
+    private final Map<Integer, AssociationInfo> mIdToAssociationMap = new HashMap<>();
+    @GuardedBy("mLock")
+    private final Map<Integer, Integer> mUserToMaxId = new HashMap<>();
+
+    @GuardedBy("mLocalListeners")
+    private final Set<OnChangeListener> mLocalListeners = new LinkedHashSet<>();
+    @GuardedBy("mRemoteListeners")
+    private final RemoteCallbackList<IOnAssociationsChangedListener> mRemoteListeners =
+            new RemoteCallbackList<>();
+
+    private final UserManager mUserManager;
+    private final AssociationDiskStore mDiskStore;
+
+    public AssociationStore(UserManager userManager, AssociationDiskStore diskStore) {
+        mUserManager = userManager;
+        mDiskStore = diskStore;
+        mExecutor = Executors.newSingleThreadExecutor();
+    }
+
+    /**
+     * Load all alive users' associations from disk to cache.
+     */
+    public void refreshCache() {
+        Binder.withCleanCallingIdentity(() -> {
+            List<Integer> userIds = new ArrayList<>();
+            for (UserInfo user : mUserManager.getAliveUsers()) {
+                userIds.add(user.id);
+            }
+
+            synchronized (mLock) {
+                mPersisted = false;
+
+                mIdToAssociationMap.clear();
+                mUserToMaxId.clear();
+
+                // The data is stored in DE directories, so we can read the data for all users now
+                // (which would not be possible if the data was stored to CE directories).
+                Map<Integer, Associations> userToAssociationsMap =
+                        mDiskStore.readAssociationsByUsers(userIds);
+                for (Map.Entry<Integer, Associations> entry : userToAssociationsMap.entrySet()) {
+                    for (AssociationInfo association : entry.getValue().getAssociations()) {
+                        mIdToAssociationMap.put(association.getId(), association);
+                    }
+                    mUserToMaxId.put(entry.getKey(), entry.getValue().getMaxId());
+                }
+
+                mPersisted = true;
+            }
+        });
+    }
+
+    /**
+     * Get the current max association id.
+     */
+    public int getMaxId(int userId) {
+        synchronized (mLock) {
+            return mUserToMaxId.getOrDefault(userId, 0);
+        }
+    }
+
+    /**
+     * Get the next available association id.
+     */
+    public int getNextId(int userId) {
+        synchronized (mLock) {
+            return getMaxId(userId) + 1;
+        }
+    }
 
     /**
      * Add an association.
      */
     public void addAssociation(@NonNull AssociationInfo association) {
-        Slog.i(TAG, "Adding new association=" + association);
-
-        // Validity check first.
-        checkNotRevoked(association);
+        Slog.i(TAG, "Adding new association=[" + association + "]...");
 
         final int id = association.getId();
+        final int userId = association.getUserId();
 
         synchronized (mLock) {
-            if (mIdMap.containsKey(id)) {
-                Slog.e(TAG, "Association with id " + id + " already exists.");
+            if (mIdToAssociationMap.containsKey(id)) {
+                Slog.e(TAG, "Association with id=[" + id + "] already exists.");
                 return;
             }
-            mIdMap.put(id, association);
 
-            final MacAddress address = association.getDeviceMacAddress();
-            if (address != null) {
-                mAddressMap.computeIfAbsent(address, it -> new HashSet<>()).add(id);
-            }
+            mIdToAssociationMap.put(id, association);
+            mUserToMaxId.put(userId, Math.max(mUserToMaxId.getOrDefault(userId, 0), id));
 
-            invalidateCacheForUserLocked(association.getUserId());
+            writeCacheToDisk(userId);
 
             Slog.i(TAG, "Done adding new association.");
         }
 
+        logCreateAssociation(association.getDeviceProfile());
+
         broadcastChange(CHANGE_TYPE_ADDED, association);
     }
 
@@ -154,18 +223,16 @@
      * Update an association.
      */
     public void updateAssociation(@NonNull AssociationInfo updated) {
-        Slog.i(TAG, "Updating new association=" + updated);
-        // Validity check first.
-        checkNotRevoked(updated);
+        Slog.i(TAG, "Updating new association=[" + updated + "]...");
 
         final int id = updated.getId();
-
         final AssociationInfo current;
         final boolean macAddressChanged;
+
         synchronized (mLock) {
-            current = mIdMap.get(id);
+            current = mIdToAssociationMap.get(id);
             if (current == null) {
-                Slog.w(TAG, "Can't update association. It does not exist.");
+                Slog.w(TAG, "Can't update association id=[" + id + "]. It does not exist.");
                 return;
             }
 
@@ -174,174 +241,238 @@
                 return;
             }
 
-            // Update the ID-to-Association map.
-            mIdMap.put(id, updated);
-            // Invalidate the corresponding user cache entry.
-            invalidateCacheForUserLocked(current.getUserId());
+            mIdToAssociationMap.put(id, updated);
 
-            // Update the MacAddress-to-List<Association> map if needed.
-            final MacAddress updatedAddress = updated.getDeviceMacAddress();
-            final MacAddress currentAddress = current.getDeviceMacAddress();
-            macAddressChanged = !Objects.equals(currentAddress, updatedAddress);
-            if (macAddressChanged) {
-                if (currentAddress != null) {
-                    mAddressMap.get(currentAddress).remove(id);
-                }
-                if (updatedAddress != null) {
-                    mAddressMap.computeIfAbsent(updatedAddress, it -> new HashSet<>()).add(id);
-                }
-            }
-            Slog.i(TAG, "Done updating association.");
+            writeCacheToDisk(updated.getUserId());
         }
 
+        Slog.i(TAG, "Done updating association.");
+
+        // Check if the MacAddress has changed.
+        final MacAddress updatedAddress = updated.getDeviceMacAddress();
+        final MacAddress currentAddress = current.getDeviceMacAddress();
+        macAddressChanged = !Objects.equals(currentAddress, updatedAddress);
+
         final int changeType = macAddressChanged ? CHANGE_TYPE_UPDATED_ADDRESS_CHANGED
                 : CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED;
+
         broadcastChange(changeType, updated);
     }
 
     /**
-     * Remove an association
+     * Remove an association.
      */
     public void removeAssociation(int id) {
-        Slog.i(TAG, "Removing association id=" + id);
+        Slog.i(TAG, "Removing association id=[" + id + "]...");
 
         final AssociationInfo association;
+
         synchronized (mLock) {
-            association = mIdMap.remove(id);
+            association = mIdToAssociationMap.remove(id);
 
             if (association == null) {
-                Slog.w(TAG, "Can't remove association. It does not exist.");
+                Slog.w(TAG, "Can't remove association id=[" + id + "]. It does not exist.");
                 return;
             }
 
-            final MacAddress macAddress = association.getDeviceMacAddress();
-            if (macAddress != null) {
-                mAddressMap.get(macAddress).remove(id);
-            }
-
-            invalidateCacheForUserLocked(association.getUserId());
+            writeCacheToDisk(association.getUserId());
 
             Slog.i(TAG, "Done removing association.");
         }
 
+        logRemoveAssociation(association.getDeviceProfile());
+
         broadcastChange(CHANGE_TYPE_REMOVED, association);
     }
 
+    private void writeCacheToDisk(@UserIdInt int userId) {
+        mExecutor.execute(() -> {
+            Associations associations = new Associations();
+            synchronized (mLock) {
+                associations.setMaxId(mUserToMaxId.getOrDefault(userId, 0));
+                associations.setAssociations(
+                        CollectionUtils.filter(mIdToAssociationMap.values().stream().toList(),
+                                a -> a.getUserId() == userId));
+            }
+            mDiskStore.writeAssociationsForUser(userId, associations);
+        });
+    }
+
     /**
-     * @return a "snapshot" of the current state of the existing associations.
+     * Get a copy of all associations including pending and revoked ones.
+     * Modifying the copy won't modify the actual associations.
+     *
+     * If a cache miss happens, read from disk.
      */
-    public @NonNull Collection<AssociationInfo> getAssociations() {
+    @NonNull
+    public List<AssociationInfo> getAssociations() {
         synchronized (mLock) {
-            // IMPORTANT: make and return a COPY of the mIdMap.values(), NOT a "direct" reference.
-            // The HashMap.values() returns a collection which is backed by the HashMap, so changes
-            // to the HashMap are reflected in this collection.
-            // For us this means that if mIdMap is modified while the iteration over mIdMap.values()
-            // is in progress it may lead to "undefined results" (according to the HashMap's
-            // documentation) or cause ConcurrentModificationExceptions in the iterator (according
-            // to the bugreports...).
-            return List.copyOf(mIdMap.values());
+            if (!mPersisted) {
+                refreshCache();
+            }
+            return List.copyOf(mIdToAssociationMap.values());
         }
     }
 
     /**
-     * Get associations for the user.
+     * Get a copy of active associations.
      */
-    public @NonNull List<AssociationInfo> getAssociationsForUser(@UserIdInt int userId) {
+    @NonNull
+    public List<AssociationInfo> getActiveAssociations() {
         synchronized (mLock) {
-            return getAssociationsForUserLocked(userId);
+            return CollectionUtils.filter(getAssociations(), AssociationInfo::isActive);
         }
     }
 
     /**
-     * Get associations for the package
+     * Get a copy of all associations by user.
      */
-    public @NonNull List<AssociationInfo> getAssociationsForPackage(
-            @UserIdInt int userId, @NonNull String packageName) {
-        final List<AssociationInfo> associationsForUser = getAssociationsForUser(userId);
-        final List<AssociationInfo> associationsForPackage =
-                CollectionUtils.filter(associationsForUser,
-                        it -> it.getPackageName().equals(packageName));
-        return Collections.unmodifiableList(associationsForPackage);
+    @NonNull
+    public List<AssociationInfo> getAssociationsByUser(@UserIdInt int userId) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(), a -> a.getUserId() == userId);
+        }
     }
 
     /**
-     * Get associations by mac address for the package.
+     * Get a copy of active associations by user.
      */
-    public @Nullable AssociationInfo getAssociationsForPackageWithAddress(
+    @NonNull
+    public List<AssociationInfo> getActiveAssociationsByUser(@UserIdInt int userId) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getActiveAssociations(), a -> a.getUserId() == userId);
+        }
+    }
+
+    /**
+     * Get a copy of all associations by package.
+     */
+    @NonNull
+    public List<AssociationInfo> getAssociationsByPackage(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociationsByUser(userId),
+                    a -> a.getPackageName().equals(packageName));
+        }
+    }
+
+    /**
+     * Get a copy of active associations by package.
+     */
+    @NonNull
+    public List<AssociationInfo> getActiveAssociationsByPackage(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getActiveAssociationsByUser(userId),
+                    a -> a.getPackageName().equals(packageName));
+        }
+    }
+
+    /**
+     * Get the first active association with the mac address.
+     */
+    @Nullable
+    public AssociationInfo getFirstAssociationByAddress(
             @UserIdInt int userId, @NonNull String packageName, @NonNull String macAddress) {
-        final List<AssociationInfo> associations = getAssociationsByAddress(macAddress);
-        return CollectionUtils.find(associations,
-                it -> it.belongsToPackage(userId, packageName));
-    }
-
-    /**
-     * Get association by id.
-     */
-    public @Nullable AssociationInfo getAssociationById(int id) {
         synchronized (mLock) {
-            return mIdMap.get(id);
+            return CollectionUtils.find(getActiveAssociationsByPackage(userId, packageName),
+                    a -> a.getDeviceMacAddress() != null && a.getDeviceMacAddress()
+                            .equals(MacAddress.fromString(macAddress)));
         }
     }
 
     /**
-     * Get associations by mac address.
+     * Get the association by id.
+     */
+    @Nullable
+    public AssociationInfo getAssociationById(int id) {
+        synchronized (mLock) {
+            return mIdToAssociationMap.get(id);
+        }
+    }
+
+    /**
+     * Get a copy of active associations by mac address.
      */
     @NonNull
-    public List<AssociationInfo> getAssociationsByAddress(@NonNull String macAddress) {
-        final MacAddress address = MacAddress.fromString(macAddress);
-
+    public List<AssociationInfo> getActiveAssociationsByAddress(@NonNull String macAddress) {
         synchronized (mLock) {
-            final Set<Integer> ids = mAddressMap.get(address);
-            if (ids == null) return Collections.emptyList();
-
-            final List<AssociationInfo> associations = new ArrayList<>(ids.size());
-            for (Integer id : ids) {
-                associations.add(mIdMap.get(id));
-            }
-
-            return Collections.unmodifiableList(associations);
+            return CollectionUtils.filter(getActiveAssociations(),
+                    a -> a.getDeviceMacAddress() != null && a.getDeviceMacAddress()
+                            .equals(MacAddress.fromString(macAddress)));
         }
     }
 
-    @GuardedBy("mLock")
+    /**
+     * Get a copy of revoked associations.
+     */
     @NonNull
-    private List<AssociationInfo> getAssociationsForUserLocked(@UserIdInt int userId) {
-        final List<AssociationInfo> cached = mCachedPerUser.get(userId);
-        if (cached != null) {
-            return cached;
-        }
-
-        final List<AssociationInfo> associationsForUser = new ArrayList<>();
-        for (AssociationInfo association : mIdMap.values()) {
-            if (association.getUserId() == userId) {
-                associationsForUser.add(association);
-            }
-        }
-        final List<AssociationInfo> set = Collections.unmodifiableList(associationsForUser);
-        mCachedPerUser.set(userId, set);
-        return set;
-    }
-
-    @GuardedBy("mLock")
-    private void invalidateCacheForUserLocked(@UserIdInt int userId) {
-        mCachedPerUser.delete(userId);
-    }
-
-    /**
-     * Register a listener for association changes.
-     */
-    public void registerListener(@NonNull OnChangeListener listener) {
-        synchronized (mListeners) {
-            mListeners.add(listener);
+    public List<AssociationInfo> getRevokedAssociations() {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(), AssociationInfo::isRevoked);
         }
     }
 
     /**
-     * Unregister a listener previously registered for association changes.
+     * Get a copy of revoked associations for the package.
      */
-    public void unregisterListener(@NonNull OnChangeListener listener) {
-        synchronized (mListeners) {
-            mListeners.remove(listener);
+    @NonNull
+    public List<AssociationInfo> getRevokedAssociations(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(),
+                    a -> packageName.equals(a.getPackageName()) && a.getUserId() == userId
+                            && a.isRevoked());
+        }
+    }
+
+    /**
+     * Get a copy of active associations.
+     */
+    @NonNull
+    public List<AssociationInfo> getPendingAssociations(@UserIdInt int userId,
+            @NonNull String packageName) {
+        synchronized (mLock) {
+            return CollectionUtils.filter(getAssociations(),
+                    a -> packageName.equals(a.getPackageName()) && a.getUserId() == userId
+                            && a.isPending());
+        }
+    }
+
+    /**
+     * Register a local listener for association changes.
+     */
+    public void registerLocalListener(@NonNull OnChangeListener listener) {
+        synchronized (mLocalListeners) {
+            mLocalListeners.add(listener);
+        }
+    }
+
+    /**
+     * Unregister a local listener previously registered for association changes.
+     */
+    public void unregisterLocalListener(@NonNull OnChangeListener listener) {
+        synchronized (mLocalListeners) {
+            mLocalListeners.remove(listener);
+        }
+    }
+
+    /**
+     * Register a remote listener for association changes.
+     */
+    public void registerRemoteListener(@NonNull IOnAssociationsChangedListener listener,
+            int userId) {
+        synchronized (mRemoteListeners) {
+            mRemoteListeners.register(listener, userId);
+        }
+    }
+
+    /**
+     * Unregister a remote listener previously registered for association changes.
+     */
+    public void unregisterRemoteListener(@NonNull IOnAssociationsChangedListener listener) {
+        synchronized (mRemoteListeners) {
+            mRemoteListeners.unregister(listener);
         }
     }
 
@@ -350,52 +481,39 @@
      */
     public void dump(@NonNull PrintWriter out) {
         out.append("Companion Device Associations: ");
-        if (getAssociations().isEmpty()) {
+        if (getActiveAssociations().isEmpty()) {
             out.append("<empty>\n");
         } else {
             out.append("\n");
-            for (AssociationInfo a : getAssociations()) {
+            for (AssociationInfo a : getActiveAssociations()) {
                 out.append("  ").append(a.toString()).append('\n');
             }
         }
     }
 
     private void broadcastChange(@ChangeType int changeType, AssociationInfo association) {
-        synchronized (mListeners) {
-            for (OnChangeListener listener : mListeners) {
+        synchronized (mLocalListeners) {
+            for (OnChangeListener listener : mLocalListeners) {
                 listener.onAssociationChanged(changeType, association);
             }
         }
-    }
-
-    /**
-     * Set associations to cache. It will clear the existing cache.
-     */
-    public void setAssociationsToCache(Collection<AssociationInfo> associations) {
-        // Validity check first.
-        associations.forEach(AssociationStore::checkNotRevoked);
-
-        synchronized (mLock) {
-            mIdMap.clear();
-            mAddressMap.clear();
-            mCachedPerUser.clear();
-
-            for (AssociationInfo association : associations) {
-                final int id = association.getId();
-                mIdMap.put(id, association);
-
-                final MacAddress address = association.getDeviceMacAddress();
-                if (address != null) {
-                    mAddressMap.computeIfAbsent(address, it -> new HashSet<>()).add(id);
-                }
+        synchronized (mRemoteListeners) {
+            final int userId = association.getUserId();
+            final List<AssociationInfo> updatedAssociations = getActiveAssociationsByUser(userId);
+            // Notify listeners if ADDED, REMOVED or UPDATED_ADDRESS_CHANGED.
+            // Do NOT notify when UPDATED_ADDRESS_UNCHANGED, which means a minor tweak in
+            // association's configs, which "listeners" won't (and shouldn't) be able to see.
+            if (changeType != CHANGE_TYPE_UPDATED_ADDRESS_UNCHANGED) {
+                mRemoteListeners.broadcast((listener, callbackUserId) -> {
+                    int listenerUserId = (int) callbackUserId;
+                    if (listenerUserId == userId || listenerUserId == UserHandle.USER_ALL) {
+                        try {
+                            listener.onAssociationsChanged(updatedAssociations);
+                        } catch (RemoteException ignored) {
+                        }
+                    }
+                });
             }
         }
     }
-
-    private static void checkNotRevoked(@NonNull AssociationInfo association) {
-        if (association.isRevoked()) {
-            throw new IllegalArgumentException(
-                    "Revoked (removed) associations MUST NOT appear in the AssociationStore");
-        }
-    }
 }
diff --git a/services/companion/java/com/android/server/companion/association/Associations.java b/services/companion/java/com/android/server/companion/association/Associations.java
new file mode 100644
index 0000000..7da3699
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/association/Associations.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.companion.association;
+
+import android.companion.AssociationInfo;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Represents associations per user. Should be only used by Association stores.
+ */
+public class Associations {
+
+    private int mVersion = 0;
+
+    private List<AssociationInfo> mAssociations = new ArrayList<>();
+
+    private int mMaxId = 0;
+
+    public Associations() {
+    }
+
+    public void setVersion(int version) {
+        mVersion = version;
+    }
+
+    /**
+     * Add an association.
+     */
+    public void addAssociation(AssociationInfo association) {
+        mAssociations.add(association);
+    }
+
+    public void setMaxId(int maxId) {
+        mMaxId = maxId;
+    }
+
+    public void setAssociations(List<AssociationInfo> associations) {
+        mAssociations = List.copyOf(associations);
+    }
+
+    public int getVersion() {
+        return mVersion;
+    }
+
+    public int getMaxId() {
+        return mMaxId;
+    }
+
+    public List<AssociationInfo> getAssociations() {
+        return mAssociations;
+    }
+}
diff --git a/services/companion/java/com/android/server/companion/association/DisassociationProcessor.java b/services/companion/java/com/android/server/companion/association/DisassociationProcessor.java
new file mode 100644
index 0000000..ec897791
--- /dev/null
+++ b/services/companion/java/com/android/server/companion/association/DisassociationProcessor.java
@@ -0,0 +1,218 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.companion.association;
+
+import static android.app.ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE;
+import static android.companion.AssociationRequest.DEVICE_PROFILE_AUTOMOTIVE_PROJECTION;
+
+import static com.android.internal.util.CollectionUtils.any;
+import static com.android.server.companion.utils.RolesUtils.removeRoleHolderForAssociation;
+
+import android.annotation.NonNull;
+import android.annotation.SuppressLint;
+import android.annotation.UserIdInt;
+import android.app.ActivityManager;
+import android.companion.AssociationInfo;
+import android.content.Context;
+import android.content.pm.PackageManagerInternal;
+import android.os.Binder;
+import android.os.UserHandle;
+import android.util.Slog;
+
+import com.android.server.companion.CompanionApplicationController;
+import com.android.server.companion.datatransfer.SystemDataTransferRequestStore;
+import com.android.server.companion.presence.CompanionDevicePresenceMonitor;
+import com.android.server.companion.transport.CompanionTransportManager;
+
+/**
+ * A class response for Association removal.
+ */
+@SuppressLint("LongLogTag")
+public class DisassociationProcessor {
+
+    private static final String TAG = "CDM_DisassociationProcessor";
+    @NonNull
+    private final Context mContext;
+    @NonNull
+    private final AssociationStore mAssociationStore;
+    @NonNull
+    private final PackageManagerInternal mPackageManagerInternal;
+    @NonNull
+    private final CompanionDevicePresenceMonitor mDevicePresenceMonitor;
+    @NonNull
+    private final SystemDataTransferRequestStore mSystemDataTransferRequestStore;
+    @NonNull
+    private final CompanionApplicationController mCompanionAppController;
+    @NonNull
+    private final CompanionTransportManager mTransportManager;
+    private final OnPackageVisibilityChangeListener mOnPackageVisibilityChangeListener;
+    private final ActivityManager mActivityManager;
+
+    public DisassociationProcessor(@NonNull Context context,
+            @NonNull ActivityManager activityManager,
+            @NonNull AssociationStore associationStore,
+            @NonNull PackageManagerInternal packageManager,
+            @NonNull CompanionDevicePresenceMonitor devicePresenceMonitor,
+            @NonNull CompanionApplicationController applicationController,
+            @NonNull SystemDataTransferRequestStore systemDataTransferRequestStore,
+            @NonNull CompanionTransportManager companionTransportManager) {
+        mContext = context;
+        mActivityManager = activityManager;
+        mAssociationStore = associationStore;
+        mPackageManagerInternal = packageManager;
+        mOnPackageVisibilityChangeListener =
+                new OnPackageVisibilityChangeListener();
+        mDevicePresenceMonitor = devicePresenceMonitor;
+        mCompanionAppController = applicationController;
+        mSystemDataTransferRequestStore = systemDataTransferRequestStore;
+        mTransportManager = companionTransportManager;
+    }
+
+    /**
+     * Disassociate an association by id.
+     */
+    // TODO: also revoke notification access
+    public void disassociate(int id) {
+        Slog.i(TAG, "Disassociating id=[" + id + "]...");
+
+        final AssociationInfo association = mAssociationStore.getAssociationById(id);
+        if (association == null) {
+            Slog.e(TAG, "Can't disassociate id=[" + id + "]. It doesn't exist.");
+            return;
+        }
+
+        final int userId = association.getUserId();
+        final String packageName = association.getPackageName();
+        final String deviceProfile = association.getDeviceProfile();
+
+        final boolean isRoleInUseByOtherAssociations = deviceProfile != null
+                && any(mAssociationStore.getActiveAssociationsByPackage(userId, packageName),
+                    it -> deviceProfile.equals(it.getDeviceProfile()) && id != it.getId());
+
+        final int packageProcessImportance = getPackageProcessImportance(userId, packageName);
+        if (packageProcessImportance <= IMPORTANCE_VISIBLE && deviceProfile != null
+                && !isRoleInUseByOtherAssociations) {
+            // Need to remove the app from the list of role holders, but the process is visible
+            // to the user at the moment, so we'll need to do it later.
+            Slog.i(TAG, "Cannot disassociate id=[" + id + "] now - process is visible. "
+                    + "Start listening to package importance...");
+
+            AssociationInfo revokedAssociation = (new AssociationInfo.Builder(
+                    association)).setRevoked(true).build();
+            mAssociationStore.updateAssociation(revokedAssociation);
+            startListening();
+            return;
+        }
+
+        // Association cleanup.
+        mAssociationStore.removeAssociation(association.getId());
+        mSystemDataTransferRequestStore.removeRequestsByAssociationId(userId, id);
+
+        // Detach transport if exists
+        mTransportManager.detachSystemDataTransport(packageName, userId, id);
+
+        // If role is not in use by other associations, revoke the role.
+        // Do not need to remove the system role since it was pre-granted by the system.
+        if (!isRoleInUseByOtherAssociations && deviceProfile != null && !deviceProfile.equals(
+                DEVICE_PROFILE_AUTOMOTIVE_PROJECTION)) {
+            removeRoleHolderForAssociation(mContext, association.getUserId(),
+                    association.getPackageName(), association.getDeviceProfile());
+        }
+
+        // Unbind the app if needed.
+        final boolean wasPresent = mDevicePresenceMonitor.isDevicePresent(id);
+        if (!wasPresent || !association.isNotifyOnDeviceNearby()) {
+            return;
+        }
+        final boolean shouldStayBound = any(
+                mAssociationStore.getActiveAssociationsByPackage(userId, packageName),
+                it -> it.isNotifyOnDeviceNearby()
+                        && mDevicePresenceMonitor.isDevicePresent(it.getId()));
+        if (!shouldStayBound) {
+            mCompanionAppController.unbindCompanionApplication(userId, packageName);
+        }
+    }
+
+    @SuppressLint("MissingPermission")
+    private int getPackageProcessImportance(@UserIdInt int userId, @NonNull String packageName) {
+        return Binder.withCleanCallingIdentity(() -> {
+            final int uid =
+                    mPackageManagerInternal.getPackageUid(packageName, /* flags */0, userId);
+            return mActivityManager.getUidImportance(uid);
+        });
+    }
+
+    private void startListening() {
+        Slog.i(TAG, "Start listening to uid importance changes...");
+        try {
+            Binder.withCleanCallingIdentity(
+                    () -> mActivityManager.addOnUidImportanceListener(
+                            mOnPackageVisibilityChangeListener,
+                            ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE));
+        }  catch (IllegalArgumentException e) {
+            Slog.e(TAG, "Failed to start listening to uid importance changes.");
+        }
+    }
+
+    private void stopListening() {
+        Slog.i(TAG, "Stop listening to uid importance changes.");
+        try {
+            Binder.withCleanCallingIdentity(() -> mActivityManager.removeOnUidImportanceListener(
+                    mOnPackageVisibilityChangeListener));
+        } catch (IllegalArgumentException e) {
+            Slog.e(TAG, "Failed to stop listening to uid importance changes.");
+        }
+    }
+
+    /**
+     * An OnUidImportanceListener class which watches the importance of the packages.
+     * In this class, we ONLY interested in the importance of the running process is greater than
+     * {@link ActivityManager.RunningAppProcessInfo#IMPORTANCE_VISIBLE}.
+     *
+     * Lastly remove the role holder for the revoked associations for the same packages.
+     *
+     * @see #disassociate(int)
+     */
+    private class OnPackageVisibilityChangeListener implements
+            ActivityManager.OnUidImportanceListener {
+
+        @Override
+        public void onUidImportance(int uid, int importance) {
+            if (importance <= ActivityManager.RunningAppProcessInfo.IMPORTANCE_VISIBLE) {
+                // The lower the importance value the more "important" the process is.
+                // We are only interested when the process ceases to be visible.
+                return;
+            }
+
+            final String packageName = mPackageManagerInternal.getNameForUid(uid);
+            if (packageName == null) {
+                // Not interested in this uid.
+                return;
+            }
+
+            int userId = UserHandle.getUserId(uid);
+            for (AssociationInfo association : mAssociationStore.getRevokedAssociations(userId,
+                    packageName)) {
+                disassociate(association.getId());
+            }
+
+            if (mAssociationStore.getRevokedAssociations().isEmpty()) {
+                stopListening();
+            }
+        }
+    }
+}
diff --git a/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java b/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java
index 894c49a..f287315 100644
--- a/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java
+++ b/services/companion/java/com/android/server/companion/association/InactiveAssociationsRemovalService.java
@@ -33,7 +33,7 @@
  * A Job Service responsible for clean up idle self-managed associations.
  *
  * The job will be executed only if the device is charging and in idle mode due to the application
- * will be killed if association/role are revoked. See {@link AssociationRevokeProcessor}
+ * will be killed if association/role are revoked. See {@link DisassociationProcessor}
  */
 public class InactiveAssociationsRemovalService extends JobService {
 
diff --git a/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java b/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
index a08e0da..c5ca0bf 100644
--- a/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
+++ b/services/companion/java/com/android/server/companion/datatransfer/SystemDataTransferProcessor.java
@@ -186,18 +186,14 @@
         intent.putExtras(extras);
 
         // Create a PendingIntent
-        final long token = Binder.clearCallingIdentity();
-        try {
-            return PendingIntent.getActivityAsUser(mContext, /*requestCode */ associationId, intent,
-                    FLAG_ONE_SHOT | FLAG_CANCEL_CURRENT | FLAG_IMMUTABLE,
-                    ActivityOptions.makeBasic()
-                            .setPendingIntentCreatorBackgroundActivityStartMode(
-                                    ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED)
-                            .toBundle(),
-                    UserHandle.CURRENT);
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
+        return Binder.withCleanCallingIdentity(() ->
+                PendingIntent.getActivityAsUser(mContext, /*requestCode */ associationId,
+                        intent, FLAG_ONE_SHOT | FLAG_CANCEL_CURRENT | FLAG_IMMUTABLE,
+                        ActivityOptions.makeBasic()
+                                .setPendingIntentCreatorBackgroundActivityStartMode(
+                                        ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED)
+                                .toBundle(),
+                        UserHandle.CURRENT));
     }
 
     /**
@@ -228,8 +224,7 @@
         }
 
         // Start permission sync
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             // TODO: refactor to work with streams of data
             mPermissionControllerManager.getRuntimePermissionBackup(UserHandle.of(userId),
                     mExecutor, backup -> {
@@ -237,39 +232,31 @@
                                 .requestPermissionRestore(associationId, backup);
                         translateFutureToCallback(future, callback);
                     });
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
      * Enable perm sync for the association
      */
     public void enablePermissionsSync(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             PermissionSyncRequest request = new PermissionSyncRequest(associationId);
             request.setUserConsented(true);
             mSystemDataTransferRequestStore.writeRequest(userId, request);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
      * Disable perm sync for the association
      */
     public void disablePermissionsSync(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             PermissionSyncRequest request = new PermissionSyncRequest(associationId);
             request.setUserConsented(false);
             mSystemDataTransferRequestStore.writeRequest(userId, request);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
@@ -277,8 +264,7 @@
      */
     @Nullable
     public PermissionSyncRequest getPermissionSyncRequest(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        return Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             List<SystemDataTransferRequest> requests =
                     mSystemDataTransferRequestStore.readRequestsByAssociationId(userId,
@@ -289,22 +275,17 @@
                 }
             }
             return null;
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     /**
      * Remove perm sync request for the association.
      */
     public void removePermissionSyncRequest(int associationId) {
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+        Binder.withCleanCallingIdentity(() -> {
             int userId = mAssociationStore.getAssociationById(associationId).getUserId();
             mSystemDataTransferRequestStore.removeRequestsByAssociationId(userId, associationId);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     private void onReceivePermissionRestore(byte[] message) {
@@ -318,14 +299,12 @@
         Slog.i(LOG_TAG, "Applying permissions.");
         // Start applying permissions
         UserHandle user = mContext.getUser();
-        final long callingIdentityToken = Binder.clearCallingIdentity();
-        try {
+
+        Binder.withCleanCallingIdentity(() -> {
             // TODO: refactor to work with streams of data
             mPermissionControllerManager.stageAndApplyRuntimePermissionsBackup(
                     message, user);
-        } finally {
-            Binder.restoreCallingIdentity(callingIdentityToken);
-        }
+        });
     }
 
     private static void translateFutureToCallback(@NonNull Future<?> future,
diff --git a/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java b/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java
index 99466a9..c89ce11 100644
--- a/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java
+++ b/services/companion/java/com/android/server/companion/presence/BleCompanionDeviceScanner.java
@@ -106,7 +106,7 @@
         checkBleState();
         registerBluetoothStateBroadcastReceiver(context);
 
-        mAssociationStore.registerListener(this);
+        mAssociationStore.registerLocalListener(this);
     }
 
     @MainThread
@@ -183,7 +183,7 @@
 
         // Collect MAC addresses from all associations.
         final Set<String> macAddresses = new HashSet<>();
-        for (AssociationInfo association : mAssociationStore.getAssociations()) {
+        for (AssociationInfo association : mAssociationStore.getActiveAssociations()) {
             if (!association.isNotifyOnDeviceNearby()) continue;
 
             // Beware that BT stack does not consider low-case MAC addresses valid, while
@@ -255,7 +255,7 @@
         if (DEBUG) Log.i(TAG, "notifyDevice_Found()" + btDeviceToString(device));
 
         final List<AssociationInfo> associations =
-                mAssociationStore.getAssociationsByAddress(device.getAddress());
+                mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
         if (DEBUG) Log.d(TAG, "  > associations=" + Arrays.toString(associations.toArray()));
 
         for (AssociationInfo association : associations) {
@@ -268,7 +268,7 @@
         if (DEBUG) Log.i(TAG, "notifyDevice_Lost()" + btDeviceToString(device));
 
         final List<AssociationInfo> associations =
-                mAssociationStore.getAssociationsByAddress(device.getAddress());
+                mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
         if (DEBUG) Log.d(TAG, "  > associations=" + Arrays.toString(associations.toArray()));
 
         for (AssociationInfo association : associations) {
@@ -319,7 +319,7 @@
                 Log.v(TAG, "  > scanResult=" + result);
 
                 final List<AssociationInfo> associations =
-                        mAssociationStore.getAssociationsByAddress(device.getAddress());
+                        mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
                 Log.v(TAG, "  > associations=" + Arrays.toString(associations.toArray()));
             }
 
diff --git a/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java b/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java
index 4da3f9b..cb363a7 100644
--- a/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java
+++ b/services/companion/java/com/android/server/companion/presence/BluetoothCompanionDeviceConnectionListener.java
@@ -93,7 +93,7 @@
 
         btAdapter.registerBluetoothConnectionCallback(
                 new HandlerExecutor(Handler.getMain()), /* callback */this);
-        mAssociationStore.registerListener(this);
+        mAssociationStore.registerLocalListener(this);
     }
 
     /**
@@ -168,7 +168,7 @@
     private void onDeviceConnectivityChanged(@NonNull BluetoothDevice device, boolean connected) {
         int userId = UserHandle.myUserId();
         final List<AssociationInfo> associations =
-                mAssociationStore.getAssociationsByAddress(device.getAddress());
+                mAssociationStore.getActiveAssociationsByAddress(device.getAddress());
         final List<ObservableUuid> observableUuids =
                 mObservableUuidStore.getObservableUuidsForUser(userId);
         final ParcelUuid[] bluetoothDeviceUuids = device.getUuids();
diff --git a/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java b/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java
index 37bbb93..7a1a83f 100644
--- a/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java
+++ b/services/companion/java/com/android/server/companion/presence/CompanionDevicePresenceMonitor.java
@@ -145,7 +145,7 @@
             Log.w(TAG, "BluetoothAdapter is NOT available.");
         }
 
-        mAssociationStore.registerListener(this);
+        mAssociationStore.registerLocalListener(this);
     }
 
     /**
@@ -481,7 +481,7 @@
      * BT connected and BLE presence and are not pending to report BLE lost.
      */
     private boolean canStopBleScan() {
-        for (AssociationInfo ai : mAssociationStore.getAssociations()) {
+        for (AssociationInfo ai : mAssociationStore.getActiveAssociations()) {
             int id = ai.getId();
             synchronized (mBtDisconnectedDevices) {
                 if (ai.isNotifyOnDeviceNearby() && !(isBtConnected(id)
diff --git a/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java b/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java
index ee8b106..db15da29 100644
--- a/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java
+++ b/services/companion/java/com/android/server/companion/presence/ObservableUuidStore.java
@@ -90,6 +90,8 @@
      * Remove the observable uuid.
      */
     public void removeObservableUuid(@UserIdInt int userId, ParcelUuid uuid, String packageName) {
+        Slog.i(TAG, "Removing uuid=[" + uuid.getUuid() + "] from store...");
+
         List<ObservableUuid> cachedObservableUuids;
 
         synchronized (mLock) {
@@ -108,7 +110,7 @@
      * Write the observable uuid.
      */
     public void writeObservableUuid(@UserIdInt int userId, ObservableUuid uuid) {
-        Slog.i(TAG, "Writing uuid=" + uuid.getUuid() + " to store.");
+        Slog.i(TAG, "Writing uuid=[" + uuid.getUuid() + "] to store...");
 
         List<ObservableUuid> cachedObservableUuids;
         synchronized (mLock) {
diff --git a/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java b/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java
index c75b1a5..369a925 100644
--- a/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java
+++ b/services/companion/java/com/android/server/companion/utils/DataStoreUtils.java
@@ -64,8 +64,8 @@
      * IMPORTANT: the method will ALWAYS return the same {@link AtomicFile} object, which makes it
      * possible to synchronize reads and writes to the file using the returned object.
      *
-     * @param userId              the userId to retrieve the storage file
-     * @param fileName         the storage file name
+     * @param userId the userId to retrieve the storage file
+     * @param fileName the storage file name
      * @return an AtomicFile for the user
      */
     @NonNull
diff --git a/services/companion/java/com/android/server/companion/utils/RolesUtils.java b/services/companion/java/com/android/server/companion/utils/RolesUtils.java
index f798e21..dd12e04 100644
--- a/services/companion/java/com/android/server/companion/utils/RolesUtils.java
+++ b/services/companion/java/com/android/server/companion/utils/RolesUtils.java
@@ -93,8 +93,8 @@
 
         Slog.i(TAG, "Removing CDM role=" + deviceProfile
                 + " for userId=" + userId + ", packageName=" + packageName);
-        final long identity = Binder.clearCallingIdentity();
-        try {
+
+        Binder.withCleanCallingIdentity(() ->
             roleManager.removeRoleHolderAsUser(deviceProfile, packageName,
                     MANAGE_HOLDERS_FLAG_DONT_KILL_APP, userHandle, context.getMainExecutor(),
                     success -> {
@@ -103,11 +103,9 @@
                                     + packageName + " from the list of " + deviceProfile
                                     + " holders.");
                         }
-                    });
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
+                    })
+        );
     }
 
-    private RolesUtils() {};
+    private RolesUtils() {}
 }
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 29e0586..84eebe8 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -22,6 +22,7 @@
 import static android.Manifest.permission.CAPTURE_VIDEO_OUTPUT;
 import static android.Manifest.permission.INTERNAL_SYSTEM_WINDOW;
 import static android.Manifest.permission.MANAGE_DISPLAYS;
+import static android.Manifest.permission.RESTRICT_DISPLAY_MODES;
 import static android.app.ActivityManager.RunningAppProcessInfo.IMPORTANCE_CACHED;
 import static android.app.ActivityManager.RunningAppProcessInfo.IMPORTANCE_GONE;
 import static android.hardware.display.DisplayManager.EventsMask;
@@ -4531,6 +4532,14 @@
             disableConnectedDisplay_enforcePermission();
             DisplayManagerService.this.enableConnectedDisplay(displayId, false);
         }
+
+        @EnforcePermission(RESTRICT_DISPLAY_MODES)
+        @Override // Binder call
+        public void requestDisplayModes(IBinder token, int displayId, @Nullable int[] modeIds) {
+            requestDisplayModes_enforcePermission();
+            DisplayManagerService.this.mDisplayModeDirector.requestDisplayModes(
+                    token, displayId, modeIds);
+        }
     }
 
     private static boolean isValidBrightness(float brightness) {
diff --git a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
index 31524dc..e1a166e 100644
--- a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
+++ b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
@@ -92,9 +92,9 @@
             Flags.FLAG_BRIGHTNESS_INT_RANGE_USER_PERCEPTION,
             Flags::brightnessIntRangeUserPerception);
 
-    private final FlagState mVsyncProximityVote = new FlagState(
-            Flags.FLAG_ENABLE_EXTERNAL_VSYNC_PROXIMITY_VOTE,
-            Flags::enableExternalVsyncProximityVote);
+    private final FlagState mRestrictDisplayModes = new FlagState(
+            Flags.FLAG_ENABLE_RESTRICT_DISPLAY_MODES,
+            Flags::enableRestrictDisplayModes);
 
     private final FlagState mVsyncLowPowerVote = new FlagState(
             Flags.FLAG_ENABLE_VSYNC_LOW_POWER_VOTE,
@@ -242,8 +242,8 @@
         return mBrightnessIntRangeUserPerceptionFlagState.isEnabled();
     }
 
-    public boolean isVsyncProximityVoteEnabled() {
-        return mVsyncProximityVote.isEnabled();
+    public boolean isRestrictDisplayModesEnabled() {
+        return mRestrictDisplayModes.isEnabled();
     }
 
     public boolean isVsyncLowPowerVoteEnabled() {
@@ -311,7 +311,7 @@
         pw.println(" " + mPowerThrottlingClamperFlagState);
         pw.println(" " + mSmallAreaDetectionFlagState);
         pw.println(" " + mBrightnessIntRangeUserPerceptionFlagState);
-        pw.println(" " + mVsyncProximityVote);
+        pw.println(" " + mRestrictDisplayModes);
         pw.println(" " + mBrightnessWearBedtimeModeClamperFlagState);
         pw.println(" " + mAutoBrightnessModesFlagState);
         pw.println(" " + mFastHdrTransitions);
diff --git a/services/core/java/com/android/server/display/feature/display_flags.aconfig b/services/core/java/com/android/server/display/feature/display_flags.aconfig
index ff0f597..a5f241f 100644
--- a/services/core/java/com/android/server/display/feature/display_flags.aconfig
+++ b/services/core/java/com/android/server/display/feature/display_flags.aconfig
@@ -130,9 +130,9 @@
 }
 
 flag {
-    name: "enable_external_vsync_proximity_vote"
+    name: "enable_restrict_display_modes"
     namespace: "display_manager"
-    description: "Feature flag for external vsync proximity vote"
+    description: "Feature flag for restriction display modes api"
     bug: "284866750"
     is_fixed_read_only: true
 }
diff --git a/services/core/java/com/android/server/display/mode/BaseModeRefreshRateVote.java b/services/core/java/com/android/server/display/mode/BaseModeRefreshRateVote.java
index c538231..6d750c0 100644
--- a/services/core/java/com/android/server/display/mode/BaseModeRefreshRateVote.java
+++ b/services/core/java/com/android/server/display/mode/BaseModeRefreshRateVote.java
@@ -16,6 +16,8 @@
 
 package com.android.server.display.mode;
 
+import android.annotation.NonNull;
+
 import java.util.Objects;
 
 class BaseModeRefreshRateVote implements Vote {
@@ -31,7 +33,7 @@
     }
 
     @Override
-    public void updateSummary(VoteSummary summary) {
+    public void updateSummary(@NonNull VoteSummary summary) {
         if (summary.appRequestBaseModeRefreshRate == 0f
                 && mAppRequestBaseModeRefreshRate > 0f) {
             summary.appRequestBaseModeRefreshRate = mAppRequestBaseModeRefreshRate;
diff --git a/services/core/java/com/android/server/display/mode/CombinedVote.java b/services/core/java/com/android/server/display/mode/CombinedVote.java
index 4b68791..3cd16bf 100644
--- a/services/core/java/com/android/server/display/mode/CombinedVote.java
+++ b/services/core/java/com/android/server/display/mode/CombinedVote.java
@@ -16,6 +16,8 @@
 
 package com.android.server.display.mode;
 
+import android.annotation.NonNull;
+
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -28,7 +30,7 @@
     }
 
     @Override
-    public void updateSummary(VoteSummary summary) {
+    public void updateSummary(@NonNull VoteSummary summary) {
         mVotes.forEach(vote -> vote.updateSummary(summary));
     }
 
diff --git a/services/core/java/com/android/server/display/mode/DisableRefreshRateSwitchingVote.java b/services/core/java/com/android/server/display/mode/DisableRefreshRateSwitchingVote.java
index 7f57406..7abb518 100644
--- a/services/core/java/com/android/server/display/mode/DisableRefreshRateSwitchingVote.java
+++ b/services/core/java/com/android/server/display/mode/DisableRefreshRateSwitchingVote.java
@@ -16,6 +16,8 @@
 
 package com.android.server.display.mode;
 
+import android.annotation.NonNull;
+
 import java.util.Objects;
 
 class DisableRefreshRateSwitchingVote implements Vote {
@@ -31,7 +33,7 @@
     }
 
     @Override
-    public void updateSummary(VoteSummary summary) {
+    public void updateSummary(@NonNull VoteSummary summary) {
         summary.disableRefreshRateSwitching =
                 summary.disableRefreshRateSwitching || mDisableRefreshRateSwitching;
     }
diff --git a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
index 64cbd54..495ae87 100644
--- a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
+++ b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
@@ -41,6 +41,7 @@
 import android.hardware.fingerprint.IUdfpsRefreshRateRequestCallback;
 import android.net.Uri;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.IThermalEventListener;
 import android.os.IThermalService;
 import android.os.Looper;
@@ -80,7 +81,6 @@
 import com.android.server.display.utils.DeviceConfigParsingUtils;
 import com.android.server.display.utils.SensorUtils;
 import com.android.server.sensors.SensorManagerInternal;
-import com.android.server.sensors.SensorManagerInternal.ProximityActiveListener;
 import com.android.server.statusbar.StatusBarManagerInternal;
 
 import java.io.PrintWriter;
@@ -128,9 +128,12 @@
     private final SettingsObserver mSettingsObserver;
     private final DisplayObserver mDisplayObserver;
     private final UdfpsObserver mUdfpsObserver;
-    private final SensorObserver mSensorObserver;
+    private final ProximitySensorObserver mSensorObserver;
     private final HbmObserver mHbmObserver;
     private final SkinThermalStatusObserver mSkinThermalStatusObserver;
+
+    @Nullable
+    private final SystemRequestObserver mSystemRequestObserver;
     private final DeviceConfigParameterProvider mConfigParameterProvider;
     private final DeviceConfigDisplaySettings mDeviceConfigDisplaySettings;
 
@@ -203,6 +206,7 @@
             .isDisplaysRefreshRatesSynchronizationEnabled();
         mIsBackUpSmoothDisplayAndForcePeakRefreshRateEnabled = displayManagerFlags
                 .isBackUpSmoothDisplayAndForcePeakRefreshRateEnabled();
+
         mContext = context;
         mHandler = new DisplayModeDirectorHandler(handler.getLooper());
         mInjector = injector;
@@ -222,10 +226,15 @@
         mVotesStorage = new VotesStorage(this::notifyDesiredDisplayModeSpecsChangedLocked,
                 mVotesStatsReporter);
         mDisplayObserver = new DisplayObserver(context, handler, mVotesStorage);
-        mSensorObserver = new SensorObserver(context, mVotesStorage, injector);
+        mSensorObserver = new ProximitySensorObserver(mVotesStorage, injector);
         mSkinThermalStatusObserver = new SkinThermalStatusObserver(injector, mVotesStorage);
         mHbmObserver = new HbmObserver(injector, mVotesStorage, BackgroundThread.getHandler(),
                 mDeviceConfigDisplaySettings);
+        if (mDvrrSupported && displayManagerFlags.isRestrictDisplayModesEnabled()) {
+            mSystemRequestObserver = new SystemRequestObserver(mVotesStorage);
+        } else {
+            mSystemRequestObserver = null;
+        }
         mAlwaysRespectAppRequest = false;
         mSupportsFrameRateOverride = injector.supportsFrameRateOverride();
     }
@@ -520,6 +529,15 @@
     }
 
     /**
+     * Delegates requestDisplayModes call to SystemRequestObserver
+     */
+    public void requestDisplayModes(IBinder token, int displayId, int[] modeIds) {
+        if (mSystemRequestObserver != null) {
+            mSystemRequestObserver.requestDisplayModes(token, displayId, modeIds);
+        }
+    }
+
+    /**
      * Print the object's state and debug information into the given stream.
      *
      * @param pw The stream to dump information to.
@@ -970,10 +988,10 @@
                     Settings.Global.LOW_POWER_MODE, 0 /*default*/) != 0;
             final Vote vote;
             if (inLowPowerMode && mVsynLowPowerVoteEnabled) {
-                vote = Vote.forSupportedModes(List.of(
-                        new SupportedModesVote.SupportedMode(/* peakRefreshRate= */ 60f,
+                vote = Vote.forSupportedRefreshRates(List.of(
+                        new SupportedRefreshRatesVote.RefreshRates(/* peakRefreshRate= */ 60f,
                                 /* vsyncRate= */ 240f),
-                        new SupportedModesVote.SupportedMode(/* peakRefreshRate= */ 60f,
+                        new SupportedRefreshRatesVote.RefreshRates(/* peakRefreshRate= */ 60f,
                                 /* vsyncRate= */ 60f)
                 ));
             } else if (inLowPowerMode) {
@@ -2158,11 +2176,11 @@
                 }
 
                 if (mVsyncLowLightBlockingVoteEnabled) {
-                    refreshRateSwitchingVote = Vote.forSupportedModesAndDisableRefreshRateSwitching(
+                    refreshRateSwitchingVote = Vote.forSupportedRefreshRatesAndDisableSwitching(
                             List.of(
-                                    new SupportedModesVote.SupportedMode(
+                                    new SupportedRefreshRatesVote.RefreshRates(
                                             /* peakRefreshRate= */ 60f, /* vsyncRate= */ 60f),
-                                    new SupportedModesVote.SupportedMode(
+                                    new SupportedRefreshRatesVote.RefreshRates(
                                             /* peakRefreshRate= */120f, /* vsyncRate= */ 120f)));
                 } else {
                     refreshRateSwitchingVote = Vote.forDisableRefreshRateSwitching();
@@ -2498,116 +2516,6 @@
         }
     }
 
-    protected static final class SensorObserver implements ProximityActiveListener,
-            DisplayManager.DisplayListener {
-        private final String mProximitySensorName = null;
-        private final String mProximitySensorType = Sensor.STRING_TYPE_PROXIMITY;
-
-        private final VotesStorage mVotesStorage;
-        private final Context mContext;
-        private final Injector mInjector;
-        @GuardedBy("mSensorObserverLock")
-        private final SparseBooleanArray mDozeStateByDisplay = new SparseBooleanArray();
-        private final Object mSensorObserverLock = new Object();
-
-        private DisplayManager mDisplayManager;
-        private DisplayManagerInternal mDisplayManagerInternal;
-        @GuardedBy("mSensorObserverLock")
-        private boolean mIsProxActive = false;
-
-        SensorObserver(Context context, VotesStorage votesStorage, Injector injector) {
-            mContext = context;
-            mVotesStorage = votesStorage;
-            mInjector = injector;
-        }
-
-        @Override
-        public void onProximityActive(boolean isActive) {
-            synchronized (mSensorObserverLock) {
-                if (mIsProxActive != isActive) {
-                    mIsProxActive = isActive;
-                    recalculateVotesLocked();
-                }
-            }
-        }
-
-        public void observe() {
-            mDisplayManager = mContext.getSystemService(DisplayManager.class);
-            mDisplayManagerInternal = mInjector.getDisplayManagerInternal();
-
-            final SensorManagerInternal sensorManager = mInjector.getSensorManagerInternal();
-            sensorManager.addProximityActiveListener(BackgroundThread.getExecutor(), this);
-
-            synchronized (mSensorObserverLock) {
-                for (Display d : mInjector.getDisplays()) {
-                    mDozeStateByDisplay.put(d.getDisplayId(), mInjector.isDozeState(d));
-                }
-            }
-            mInjector.registerDisplayListener(this, BackgroundThread.getHandler(),
-                    DisplayManager.EVENT_FLAG_DISPLAY_ADDED
-                            | DisplayManager.EVENT_FLAG_DISPLAY_CHANGED
-                            | DisplayManager.EVENT_FLAG_DISPLAY_REMOVED);
-        }
-
-        private void recalculateVotesLocked() {
-            final Display[] displays = mInjector.getDisplays();
-            for (Display d : displays) {
-                int displayId = d.getDisplayId();
-                Vote vote = null;
-                if (mIsProxActive && !mDozeStateByDisplay.get(displayId)) {
-                    final RefreshRateRange rate =
-                            mDisplayManagerInternal.getRefreshRateForDisplayAndSensor(
-                                    displayId, mProximitySensorName, mProximitySensorType);
-                    if (rate != null) {
-                        vote = Vote.forPhysicalRefreshRates(rate.min, rate.max);
-                    }
-                }
-                mVotesStorage.updateVote(displayId, Vote.PRIORITY_PROXIMITY, vote);
-            }
-        }
-
-        void dump(PrintWriter pw) {
-            pw.println("  SensorObserver");
-            synchronized (mSensorObserverLock) {
-                pw.println("    mIsProxActive=" + mIsProxActive);
-                pw.println("    mDozeStateByDisplay:");
-                for (int i = 0; i < mDozeStateByDisplay.size(); i++) {
-                    final int id = mDozeStateByDisplay.keyAt(i);
-                    final boolean dozed = mDozeStateByDisplay.valueAt(i);
-                    pw.println("      " + id + " -> " + dozed);
-                }
-            }
-        }
-
-        @Override
-        public void onDisplayAdded(int displayId) {
-            boolean isDozeState = mInjector.isDozeState(mInjector.getDisplay(displayId));
-            synchronized (mSensorObserverLock) {
-                mDozeStateByDisplay.put(displayId, isDozeState);
-                recalculateVotesLocked();
-            }
-        }
-
-        @Override
-        public void onDisplayChanged(int displayId) {
-            boolean wasDozeState = mDozeStateByDisplay.get(displayId);
-            synchronized (mSensorObserverLock) {
-                mDozeStateByDisplay.put(displayId,
-                        mInjector.isDozeState(mInjector.getDisplay(displayId)));
-                if (wasDozeState != mDozeStateByDisplay.get(displayId)) {
-                    recalculateVotesLocked();
-                }
-            }
-        }
-
-        @Override
-        public void onDisplayRemoved(int displayId) {
-            synchronized (mSensorObserverLock) {
-                mDozeStateByDisplay.delete(displayId);
-                recalculateVotesLocked();
-            }
-        }
-    }
 
     /**
      * Listens to DisplayManager for HBM status and applies any refresh-rate restrictions for
diff --git a/services/core/java/com/android/server/display/mode/ProximitySensorObserver.java b/services/core/java/com/android/server/display/mode/ProximitySensorObserver.java
new file mode 100644
index 0000000..11418c1
--- /dev/null
+++ b/services/core/java/com/android/server/display/mode/ProximitySensorObserver.java
@@ -0,0 +1,138 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.mode;
+
+import android.hardware.Sensor;
+import android.hardware.display.DisplayManager;
+import android.hardware.display.DisplayManagerInternal;
+import android.util.SparseBooleanArray;
+import android.view.Display;
+import android.view.SurfaceControl;
+
+import com.android.internal.annotations.GuardedBy;
+import com.android.internal.os.BackgroundThread;
+import com.android.server.sensors.SensorManagerInternal;
+
+import java.io.PrintWriter;
+
+class ProximitySensorObserver implements
+        SensorManagerInternal.ProximityActiveListener,
+        DisplayManager.DisplayListener {
+    private final String mProximitySensorName = null;
+    private final String mProximitySensorType = Sensor.STRING_TYPE_PROXIMITY;
+
+    private final VotesStorage mVotesStorage;
+    private final DisplayModeDirector.Injector mInjector;
+    @GuardedBy("mSensorObserverLock")
+    private final SparseBooleanArray mDozeStateByDisplay = new SparseBooleanArray();
+    private final Object mSensorObserverLock = new Object();
+    private DisplayManagerInternal mDisplayManagerInternal;
+    @GuardedBy("mSensorObserverLock")
+    private boolean mIsProxActive = false;
+
+    ProximitySensorObserver(VotesStorage votesStorage, DisplayModeDirector.Injector injector) {
+        mVotesStorage = votesStorage;
+        mInjector = injector;
+    }
+
+    @Override
+    public void onProximityActive(boolean isActive) {
+        synchronized (mSensorObserverLock) {
+            if (mIsProxActive != isActive) {
+                mIsProxActive = isActive;
+                recalculateVotesLocked();
+            }
+        }
+    }
+
+    void observe() {
+        mDisplayManagerInternal = mInjector.getDisplayManagerInternal();
+
+        final SensorManagerInternal sensorManager = mInjector.getSensorManagerInternal();
+        sensorManager.addProximityActiveListener(BackgroundThread.getExecutor(), this);
+
+        synchronized (mSensorObserverLock) {
+            for (Display d : mInjector.getDisplays()) {
+                mDozeStateByDisplay.put(d.getDisplayId(), mInjector.isDozeState(d));
+            }
+        }
+        mInjector.registerDisplayListener(this, BackgroundThread.getHandler(),
+                DisplayManager.EVENT_FLAG_DISPLAY_ADDED
+                        | DisplayManager.EVENT_FLAG_DISPLAY_CHANGED
+                        | DisplayManager.EVENT_FLAG_DISPLAY_REMOVED);
+    }
+
+    @GuardedBy("mSensorObserverLock")
+    private void recalculateVotesLocked() {
+        final Display[] displays = mInjector.getDisplays();
+        for (Display d : displays) {
+            int displayId = d.getDisplayId();
+            Vote vote = null;
+            if (mIsProxActive && !mDozeStateByDisplay.get(displayId)) {
+                final SurfaceControl.RefreshRateRange rate =
+                        mDisplayManagerInternal.getRefreshRateForDisplayAndSensor(
+                                displayId, mProximitySensorName, mProximitySensorType);
+                if (rate != null) {
+                    vote = Vote.forPhysicalRefreshRates(rate.min, rate.max);
+                }
+            }
+            mVotesStorage.updateVote(displayId, Vote.PRIORITY_PROXIMITY, vote);
+        }
+    }
+
+    void dump(PrintWriter pw) {
+        pw.println("  SensorObserver");
+        synchronized (mSensorObserverLock) {
+            pw.println("    mIsProxActive=" + mIsProxActive);
+            pw.println("    mDozeStateByDisplay:");
+            for (int i = 0; i < mDozeStateByDisplay.size(); i++) {
+                final int id = mDozeStateByDisplay.keyAt(i);
+                final boolean dozed = mDozeStateByDisplay.valueAt(i);
+                pw.println("      " + id + " -> " + dozed);
+            }
+        }
+    }
+
+    @Override
+    public void onDisplayAdded(int displayId) {
+        boolean isDozeState = mInjector.isDozeState(mInjector.getDisplay(displayId));
+        synchronized (mSensorObserverLock) {
+            mDozeStateByDisplay.put(displayId, isDozeState);
+            recalculateVotesLocked();
+        }
+    }
+
+    @Override
+    public void onDisplayChanged(int displayId) {
+        synchronized (mSensorObserverLock) {
+            boolean wasDozeState = mDozeStateByDisplay.get(displayId);
+            mDozeStateByDisplay.put(displayId,
+                    mInjector.isDozeState(mInjector.getDisplay(displayId)));
+            if (wasDozeState != mDozeStateByDisplay.get(displayId)) {
+                recalculateVotesLocked();
+            }
+        }
+    }
+
+    @Override
+    public void onDisplayRemoved(int displayId) {
+        synchronized (mSensorObserverLock) {
+            mDozeStateByDisplay.delete(displayId);
+            recalculateVotesLocked();
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/display/mode/RefreshRateVote.java b/services/core/java/com/android/server/display/mode/RefreshRateVote.java
index 670b8a1..b96ab3b 100644
--- a/services/core/java/com/android/server/display/mode/RefreshRateVote.java
+++ b/services/core/java/com/android/server/display/mode/RefreshRateVote.java
@@ -16,6 +16,8 @@
 
 package com.android.server.display.mode;
 
+import android.annotation.NonNull;
+
 import java.util.Objects;
 
 
@@ -64,7 +66,7 @@
          *  Vote: min(ignored)     min(applied)  min(applied+physical)  max(applied)  max(ignored)
          */
         @Override
-        public void updateSummary(VoteSummary summary) {
+        public void updateSummary(@NonNull VoteSummary summary) {
             summary.minRenderFrameRate = Math.max(summary.minRenderFrameRate, mMinRefreshRate);
             summary.maxRenderFrameRate = Math.min(summary.maxRenderFrameRate, mMaxRefreshRate);
             // Physical refresh rate cannot be lower than the minimal render frame rate.
@@ -97,7 +99,7 @@
          *  Vote: min(ignored) min(applied)  max(applied+render)     max(applied)  max(ignored)
          */
         @Override
-        public void updateSummary(VoteSummary summary) {
+        public void updateSummary(@NonNull VoteSummary summary) {
             summary.minPhysicalRefreshRate = Math.max(summary.minPhysicalRefreshRate,
                     mMinRefreshRate);
             summary.maxPhysicalRefreshRate = Math.min(summary.maxPhysicalRefreshRate,
diff --git a/services/core/java/com/android/server/display/mode/SizeVote.java b/services/core/java/com/android/server/display/mode/SizeVote.java
index f2f8dc4..f5a5abe 100644
--- a/services/core/java/com/android/server/display/mode/SizeVote.java
+++ b/services/core/java/com/android/server/display/mode/SizeVote.java
@@ -16,6 +16,8 @@
 
 package com.android.server.display.mode;
 
+import android.annotation.NonNull;
+
 import java.util.Objects;
 
 class SizeVote implements Vote {
@@ -48,7 +50,7 @@
     }
 
     @Override
-    public void updateSummary(VoteSummary summary) {
+    public void updateSummary(@NonNull VoteSummary summary) {
         if (mHeight > 0 && mWidth > 0) {
             // For display size, disable refresh rate switching and base mode refresh rate use
             // only the first vote we come across (i.e. the highest priority vote that includes
diff --git a/services/core/java/com/android/server/display/mode/SupportedModesVote.java b/services/core/java/com/android/server/display/mode/SupportedModesVote.java
index 7eebcc0..0cf8311 100644
--- a/services/core/java/com/android/server/display/mode/SupportedModesVote.java
+++ b/services/core/java/com/android/server/display/mode/SupportedModesVote.java
@@ -16,77 +16,42 @@
 
 package com.android.server.display.mode;
 
-import java.util.ArrayList;
+import android.annotation.NonNull;
+
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 
-class SupportedModesVote implements Vote {
+public class SupportedModesVote implements Vote {
 
-    final List<SupportedMode> mSupportedModes;
+    final List<Integer> mModeIds;
 
-    SupportedModesVote(List<SupportedMode> supportedModes) {
-        mSupportedModes = Collections.unmodifiableList(supportedModes);
+    SupportedModesVote(List<Integer> modeIds) {
+        mModeIds = Collections.unmodifiableList(modeIds);
+    }
+    @Override
+    public void updateSummary(@NonNull VoteSummary summary) {
+        if (summary.supportedModeIds == null) {
+            summary.supportedModeIds = mModeIds;
+        } else {
+            summary.supportedModeIds.retainAll(mModeIds);
+        }
     }
 
-    /**
-     * Summary should have subset of supported modes.
-     * If Vote1.supportedModes=(A,B), Vote2.supportedModes=(B,C) then summary.supportedModes=(B)
-     * If summary.supportedModes==null then there is no restriction on supportedModes
-     */
     @Override
-    public void updateSummary(VoteSummary summary) {
-        if (summary.supportedModes == null) {
-            summary.supportedModes = new ArrayList<>(mSupportedModes);
-        } else {
-            summary.supportedModes.retainAll(mSupportedModes);
-        }
+    public String toString() {
+        return "SupportedModesVote{ mModeIds=" + mModeIds + " }";
     }
 
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         if (!(o instanceof SupportedModesVote that)) return false;
-        return mSupportedModes.equals(that.mSupportedModes);
+        return mModeIds.equals(that.mModeIds);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(mSupportedModes);
-    }
-
-    @Override
-    public String toString() {
-        return "SupportedModesVote{ mSupportedModes=" + mSupportedModes + " }";
-    }
-
-    static class SupportedMode {
-        final float mPeakRefreshRate;
-        final float mVsyncRate;
-
-
-        SupportedMode(float peakRefreshRate, float vsyncRate) {
-            mPeakRefreshRate = peakRefreshRate;
-            mVsyncRate = vsyncRate;
-        }
-
-        @Override
-        public boolean equals(Object o) {
-            if (this == o) return true;
-            if (!(o instanceof SupportedMode that)) return false;
-            return Float.compare(that.mPeakRefreshRate, mPeakRefreshRate) == 0
-                    && Float.compare(that.mVsyncRate, mVsyncRate) == 0;
-        }
-
-        @Override
-        public int hashCode() {
-            return Objects.hash(mPeakRefreshRate, mVsyncRate);
-        }
-
-        @Override
-        public String toString() {
-            return "SupportedMode{ mPeakRefreshRate=" + mPeakRefreshRate
-                    + ", mVsyncRate=" + mVsyncRate + " }";
-        }
+        return Objects.hash(mModeIds);
     }
 }
diff --git a/services/core/java/com/android/server/display/mode/SupportedRefreshRatesVote.java b/services/core/java/com/android/server/display/mode/SupportedRefreshRatesVote.java
new file mode 100644
index 0000000..5305487
--- /dev/null
+++ b/services/core/java/com/android/server/display/mode/SupportedRefreshRatesVote.java
@@ -0,0 +1,94 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.mode;
+
+import android.annotation.NonNull;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+
+class SupportedRefreshRatesVote implements Vote {
+
+    final List<RefreshRates> mRefreshRates;
+
+    SupportedRefreshRatesVote(List<RefreshRates> refreshRates) {
+        mRefreshRates = Collections.unmodifiableList(refreshRates);
+    }
+
+    /**
+     * Summary should have subset of supported modes.
+     * If Vote1.refreshRates=(A,B), Vote2.refreshRates=(B,C)
+     *  then summary.supportedRefreshRates=(B)
+     * If summary.supportedRefreshRates==null then there is no restriction on supportedRefreshRates
+     */
+    @Override
+    public void updateSummary(@NonNull VoteSummary summary) {
+        if (summary.supportedRefreshRates == null) {
+            summary.supportedRefreshRates = new ArrayList<>(mRefreshRates);
+        } else {
+            summary.supportedRefreshRates.retainAll(mRefreshRates);
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof SupportedRefreshRatesVote that)) return false;
+        return mRefreshRates.equals(that.mRefreshRates);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(mRefreshRates);
+    }
+
+    @Override
+    public String toString() {
+        return "SupportedRefreshRatesVote{ mSupportedModes=" + mRefreshRates + " }";
+    }
+
+    static class RefreshRates {
+        final float mPeakRefreshRate;
+        final float mVsyncRate;
+
+        RefreshRates(float peakRefreshRate, float vsyncRate) {
+            mPeakRefreshRate = peakRefreshRate;
+            mVsyncRate = vsyncRate;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof RefreshRates that)) return false;
+            return Float.compare(that.mPeakRefreshRate, mPeakRefreshRate) == 0
+                    && Float.compare(that.mVsyncRate, mVsyncRate) == 0;
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(mPeakRefreshRate, mVsyncRate);
+        }
+
+        @Override
+        public String toString() {
+            return "RefreshRates{ mPeakRefreshRate=" + mPeakRefreshRate
+                    + ", mVsyncRate=" + mVsyncRate + " }";
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/display/mode/SystemRequestObserver.java b/services/core/java/com/android/server/display/mode/SystemRequestObserver.java
new file mode 100644
index 0000000..15f19cc
--- /dev/null
+++ b/services/core/java/com/android/server/display/mode/SystemRequestObserver.java
@@ -0,0 +1,139 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.mode;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.IBinder;
+import android.os.RemoteException;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.GuardedBy;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * SystemRequestObserver responsible for handling system requests to filter allowable display
+ * modes
+ */
+class SystemRequestObserver {
+    private final VotesStorage mVotesStorage;
+
+    private final IBinder.DeathRecipient mDeathRecipient = new IBinder.DeathRecipient() {
+        @Override
+        public void binderDied() {
+            // noop, binderDied(@NonNull IBinder who) is overridden
+        }
+        @Override
+        public void binderDied(@NonNull IBinder who) {
+            removeSystemRequestedVotes(who);
+            who.unlinkToDeath(mDeathRecipient, 0);
+        }
+    };
+
+    private final Object mLock = new Object();
+    @GuardedBy("mLock")
+    private final Map<IBinder, SparseArray<List<Integer>>> mDisplaysRestrictions = new HashMap<>();
+
+    SystemRequestObserver(VotesStorage storage) {
+        mVotesStorage = storage;
+    }
+
+    void requestDisplayModes(IBinder token, int displayId, @Nullable int[] modeIds) {
+        if (modeIds == null) {
+            removeSystemRequestedVote(token, displayId);
+        } else {
+            addSystemRequestedVote(token, displayId, modeIds);
+        }
+    }
+
+    private void addSystemRequestedVote(IBinder token, int displayId, @NonNull int[] modeIds) {
+        try {
+            boolean needLinkToDeath = false;
+            List<Integer> modeIdsList = new ArrayList<>();
+            for (int mode: modeIds) {
+                modeIdsList.add(mode);
+            }
+            synchronized (mLock) {
+                SparseArray<List<Integer>> modesByDisplay = mDisplaysRestrictions.get(token);
+                if (modesByDisplay == null) {
+                    needLinkToDeath = true;
+                    modesByDisplay = new SparseArray<>();
+                    mDisplaysRestrictions.put(token, modesByDisplay);
+                }
+
+                modesByDisplay.put(displayId, modeIdsList);
+                updateStorageLocked(displayId);
+            }
+            if (needLinkToDeath) {
+                token.linkToDeath(mDeathRecipient, 0);
+            }
+        } catch (RemoteException re) {
+            removeSystemRequestedVotes(token);
+        }
+    }
+
+    private void removeSystemRequestedVote(IBinder token, int displayId) {
+        boolean needToUnlink = false;
+        synchronized (mLock) {
+            SparseArray<List<Integer>> modesByDisplay = mDisplaysRestrictions.get(token);
+            if (modesByDisplay != null) {
+                modesByDisplay.remove(displayId);
+                needToUnlink = modesByDisplay.size() == 0;
+                updateStorageLocked(displayId);
+            }
+        }
+        if (needToUnlink) {
+            token.unlinkToDeath(mDeathRecipient, 0);
+        }
+    }
+
+    private void removeSystemRequestedVotes(IBinder token) {
+        synchronized (mLock) {
+            SparseArray<List<Integer>> removed = mDisplaysRestrictions.remove(token);
+            if (removed != null) {
+                for (int i = 0; i < removed.size(); i++) {
+                    updateStorageLocked(removed.keyAt(i));
+                }
+            }
+        }
+    }
+
+    @GuardedBy("mLock")
+    private void updateStorageLocked(int displayId) {
+        List<Integer> modeIds = new ArrayList<>();
+        boolean[] modesFound = new boolean[1];
+
+        mDisplaysRestrictions.forEach((key, value) -> {
+            List<Integer> modesForDisplay = value.get(displayId);
+            if (modesForDisplay != null) {
+                if (!modesFound[0]) {
+                    modeIds.addAll(modesForDisplay);
+                    modesFound[0] = true;
+                } else {
+                    modeIds.retainAll(modesForDisplay);
+                }
+            }
+        });
+
+        mVotesStorage.updateVote(displayId, Vote.PRIORITY_SYSTEM_REQUESTED_MODES,
+                modesFound[0] ? Vote.forSupportedModes(modeIds) : null);
+    }
+}
diff --git a/services/core/java/com/android/server/display/mode/Vote.java b/services/core/java/com/android/server/display/mode/Vote.java
index e8d5a19..5b987f4 100644
--- a/services/core/java/com/android/server/display/mode/Vote.java
+++ b/services/core/java/com/android/server/display/mode/Vote.java
@@ -16,6 +16,8 @@
 
 package com.android.server.display.mode;
 
+import android.annotation.NonNull;
+
 import java.util.List;
 
 interface Vote {
@@ -91,26 +93,29 @@
     // For concurrent displays we want to limit refresh rate on all displays
     int PRIORITY_LAYOUT_LIMITED_FRAME_RATE = 12;
 
+    // For internal application to limit display modes to specific ids
+    int PRIORITY_SYSTEM_REQUESTED_MODES = 13;
+
     // LOW_POWER_MODE force the render frame rate to [0, 60HZ] if
     // Settings.Global.LOW_POWER_MODE is on.
-    int PRIORITY_LOW_POWER_MODE = 13;
+    int PRIORITY_LOW_POWER_MODE = 14;
 
     // PRIORITY_FLICKER_REFRESH_RATE_SWITCH votes for disabling refresh rate switching. If the
     // higher priority voters' result is a range, it will fix the rate to a single choice.
     // It's used to avoid refresh rate switches in certain conditions which may result in the
     // user seeing the display flickering when the switches occur.
-    int PRIORITY_FLICKER_REFRESH_RATE_SWITCH = 14;
+    int PRIORITY_FLICKER_REFRESH_RATE_SWITCH = 15;
 
     // Force display to [0, 60HZ] if skin temperature is at or above CRITICAL.
-    int PRIORITY_SKIN_TEMPERATURE = 15;
+    int PRIORITY_SKIN_TEMPERATURE = 16;
 
     // The proximity sensor needs the refresh rate to be locked in order to function, so this is
     // set to a high priority.
-    int PRIORITY_PROXIMITY = 16;
+    int PRIORITY_PROXIMITY = 17;
 
     // The Under-Display Fingerprint Sensor (UDFPS) needs the refresh rate to be locked in order
     // to function, so this needs to be the highest priority of all votes.
-    int PRIORITY_UDFPS = 17;
+    int PRIORITY_UDFPS = 18;
 
     // Whenever a new priority is added, remember to update MIN_PRIORITY, MAX_PRIORITY, and
     // APP_REQUEST_REFRESH_RATE_RANGE_PRIORITY_CUTOFF, as well as priorityToString.
@@ -128,7 +133,7 @@
      */
     int INVALID_SIZE = -1;
 
-    void updateSummary(VoteSummary summary);
+    void updateSummary(@NonNull VoteSummary summary);
 
     static Vote forPhysicalRefreshRates(float minRefreshRate, float maxRefreshRate) {
         return new CombinedVote(
@@ -166,15 +171,22 @@
         return new BaseModeRefreshRateVote(baseModeRefreshRate);
     }
 
-    static Vote forSupportedModes(List<SupportedModesVote.SupportedMode> supportedModes) {
-        return new SupportedModesVote(supportedModes);
+    static Vote forSupportedRefreshRates(
+            List<SupportedRefreshRatesVote.RefreshRates> refreshRates) {
+        return new SupportedRefreshRatesVote(refreshRates);
+    }
+
+    static Vote forSupportedModes(List<Integer> modeIds) {
+        return new SupportedModesVote(modeIds);
     }
 
 
-    static Vote forSupportedModesAndDisableRefreshRateSwitching(
-            List<SupportedModesVote.SupportedMode> supportedModes) {
+
+    static Vote forSupportedRefreshRatesAndDisableSwitching(
+            List<SupportedRefreshRatesVote.RefreshRates> supportedRefreshRates) {
         return new CombinedVote(
-                List.of(forDisableRefreshRateSwitching(), forSupportedModes(supportedModes)));
+                List.of(forDisableRefreshRateSwitching(),
+                        forSupportedRefreshRates(supportedRefreshRates)));
     }
 
     static String priorityToString(int priority) {
diff --git a/services/core/java/com/android/server/display/mode/VoteSummary.java b/services/core/java/com/android/server/display/mode/VoteSummary.java
index 5fc36b5..d4ce892 100644
--- a/services/core/java/com/android/server/display/mode/VoteSummary.java
+++ b/services/core/java/com/android/server/display/mode/VoteSummary.java
@@ -16,6 +16,7 @@
 
 package com.android.server.display.mode;
 
+import android.annotation.Nullable;
 import android.util.Slog;
 import android.util.SparseArray;
 import android.view.Display;
@@ -39,7 +40,11 @@
     public boolean disableRefreshRateSwitching;
     public float appRequestBaseModeRefreshRate;
 
-    public List<SupportedModesVote.SupportedMode> supportedModes;
+    @Nullable
+    public List<SupportedRefreshRatesVote.RefreshRates> supportedRefreshRates;
+
+    @Nullable
+    public List<Integer> supportedModeIds;
 
     final boolean mIsDisplayResolutionRangeVotingEnabled;
 
@@ -112,6 +117,9 @@
         boolean missingBaseModeRefreshRate = appRequestBaseModeRefreshRate > 0f;
 
         for (Display.Mode mode : modes) {
+            if (!validateRefreshRatesSupported(mode)) {
+                continue;
+            }
             if (!validateModeSupported(mode)) {
                 continue;
             }
@@ -253,21 +261,37 @@
     }
 
     private boolean validateModeSupported(Display.Mode mode) {
-        if (supportedModes == null || !mSupportedModesVoteEnabled) {
+        if (supportedModeIds == null || !mSupportedModesVoteEnabled) {
             return true;
         }
-        for (SupportedModesVote.SupportedMode supportedMode : supportedModes) {
-            if (equalsWithinFloatTolerance(mode.getRefreshRate(), supportedMode.mPeakRefreshRate)
-                    && equalsWithinFloatTolerance(mode.getVsyncRate(), supportedMode.mVsyncRate)) {
+        if (supportedModeIds.contains(mode.getModeId())) {
+            return true;
+        }
+        if (mLoggingEnabled) {
+            Slog.w(TAG, "Discarding mode " + mode.getModeId()
+                    + ", supportedMode not found"
+                    + ": mode.modeId=" + mode.getModeId()
+                    + ", supportedModeIds=" + supportedModeIds);
+        }
+        return false;
+    }
+
+    private boolean validateRefreshRatesSupported(Display.Mode mode) {
+        if (supportedRefreshRates == null || !mSupportedModesVoteEnabled) {
+            return true;
+        }
+        for (SupportedRefreshRatesVote.RefreshRates refreshRates : this.supportedRefreshRates) {
+            if (equalsWithinFloatTolerance(mode.getRefreshRate(), refreshRates.mPeakRefreshRate)
+                    && equalsWithinFloatTolerance(mode.getVsyncRate(), refreshRates.mVsyncRate)) {
                 return true;
             }
         }
         if (mLoggingEnabled) {
             Slog.w(TAG, "Discarding mode " + mode.getModeId()
-                    + ", supportedMode not found"
+                    + ", supportedRefreshRates not found"
                     + ": mode.refreshRate=" + mode.getRefreshRate()
                     + ", mode.vsyncRate=" + mode.getVsyncRate()
-                    + ", supportedModes=" + supportedModes);
+                    + ", supportedRefreshRates=" + supportedRefreshRates);
         }
         return false;
     }
@@ -298,7 +322,8 @@
             return false;
         }
 
-        if (supportedModes != null && mSupportedModesVoteEnabled && supportedModes.isEmpty()) {
+        if (supportedRefreshRates != null && mSupportedModesVoteEnabled
+                && supportedRefreshRates.isEmpty()) {
             if (mLoggingEnabled) {
                 Slog.w(TAG, "Vote summary resulted in empty set (empty supportedModes)");
             }
@@ -345,7 +370,8 @@
         minHeight = 0;
         disableRefreshRateSwitching = false;
         appRequestBaseModeRefreshRate = 0f;
-        supportedModes = null;
+        supportedRefreshRates = null;
+        supportedModeIds = null;
         if (mLoggingEnabled) {
             Slog.i(TAG, "Summary reset: " + this);
         }
@@ -367,7 +393,8 @@
                 + ", minHeight=" + minHeight
                 + ", disableRefreshRateSwitching=" + disableRefreshRateSwitching
                 + ", appRequestBaseModeRefreshRate=" + appRequestBaseModeRefreshRate
-                + ", supportedModes=" + supportedModes
+                + ", supportedRefreshRates=" + supportedRefreshRates
+                + ", supportedModeIds=" + supportedModeIds
                 + ", mIsDisplayResolutionRangeVotingEnabled="
                 + mIsDisplayResolutionRangeVotingEnabled
                 + ", mSupportedModesVoteEnabled=" + mSupportedModesVoteEnabled
diff --git a/services/core/java/com/android/server/display/mode/VotesStatsReporter.java b/services/core/java/com/android/server/display/mode/VotesStatsReporter.java
index e80b9451..7562a52 100644
--- a/services/core/java/com/android/server/display/mode/VotesStatsReporter.java
+++ b/services/core/java/com/android/server/display/mode/VotesStatsReporter.java
@@ -117,11 +117,11 @@
             maxRefreshRate = (int) physicalVote.mMaxRefreshRate;
         } else if (!ignoreRenderRate && (vote instanceof RefreshRateVote.RenderVote renderVote)) {
             maxRefreshRate = (int)  renderVote.mMaxRefreshRate;
-        } else if (vote instanceof SupportedModesVote supportedModesVote) {
-            // SupportedModesVote limits mode by specific refreshRates, so highest rr is allowed
+        } else if (vote instanceof SupportedRefreshRatesVote refreshRatesVote) {
+            // SupportedRefreshRatesVote limits mode by refreshRates, so highest rr is allowed
             maxRefreshRate = 0;
-            for (SupportedModesVote.SupportedMode mode : supportedModesVote.mSupportedModes) {
-                maxRefreshRate = Math.max(maxRefreshRate, (int) mode.mPeakRefreshRate);
+            for (SupportedRefreshRatesVote.RefreshRates rr : refreshRatesVote.mRefreshRates) {
+                maxRefreshRate = Math.max(maxRefreshRate, (int) rr.mPeakRefreshRate);
             }
         } else if (vote instanceof CombinedVote combinedVote) {
             for (Vote subVote: combinedVote.mVotes) {
diff --git a/services/core/java/com/android/server/display/mode/VotesStorage.java b/services/core/java/com/android/server/display/mode/VotesStorage.java
index 56c7c18..6becf1c 100644
--- a/services/core/java/com/android/server/display/mode/VotesStorage.java
+++ b/services/core/java/com/android/server/display/mode/VotesStorage.java
@@ -18,6 +18,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.util.IntArray;
 import android.util.Slog;
 import android.util.SparseArray;
 
@@ -124,6 +125,44 @@
         }
     }
 
+    /** removes all votes with certain priority from vote storage */
+    void removeAllVotesForPriority(int priority) {
+        if (mLoggingEnabled) {
+            Slog.i(TAG, "removeAllVotesForPriority(priority="
+                    + Vote.priorityToString(priority) + ")");
+        }
+        if (priority < Vote.MIN_PRIORITY || priority > Vote.MAX_PRIORITY) {
+            Slog.w(TAG, "Received an invalid priority, ignoring:"
+                    + " priority=" + Vote.priorityToString(priority));
+            return;
+        }
+        IntArray removedVotesDisplayIds = new IntArray();
+        synchronized (mStorageLock) {
+            int size = mVotesByDisplay.size();
+            for (int i = 0; i < size; i++) {
+                SparseArray<Vote> votes = mVotesByDisplay.valueAt(i);
+                if (votes.get(priority) != null) {
+                    votes.remove(priority);
+                    removedVotesDisplayIds.add(mVotesByDisplay.keyAt(i));
+                }
+            }
+        }
+        if (mLoggingEnabled) {
+            Slog.i(TAG, "Removed votes with priority=" + priority
+                    + " for displays=" + removedVotesDisplayIds);
+        }
+        int removedVotesSize = removedVotesDisplayIds.size();
+        if (removedVotesSize > 0) {
+            if (mVotesStatsReporter != null) {
+                for (int i = 0; i < removedVotesSize; i++) {
+                    mVotesStatsReporter.reportVoteChanged(
+                            removedVotesDisplayIds.get(i), priority, null);
+                }
+            }
+            mListener.onChanged();
+        }
+    }
+
     /** dump class values, for debugging */
     void dump(@NonNull PrintWriter pw) {
         SparseArray<SparseArray<Vote>> votesByDisplayLocal = new SparseArray<>();
diff --git a/services/core/java/com/android/server/input/KeyboardMetricsCollector.java b/services/core/java/com/android/server/input/KeyboardMetricsCollector.java
index b8ae737..f21fd41 100644
--- a/services/core/java/com/android/server/input/KeyboardMetricsCollector.java
+++ b/services/core/java/com/android/server/input/KeyboardMetricsCollector.java
@@ -146,6 +146,10 @@
         SPLIT_SCREEN_NAVIGATION(
                 FrameworkStatsLog.KEYBOARD_SYSTEMS_EVENT_REPORTED__KEYBOARD_SYSTEM_EVENT__SPLIT_SCREEN_NAVIGATION,
                 "SPLIT_SCREEN_NAVIGATION"),
+
+        CHANGE_SPLITSCREEN_FOCUS(
+                FrameworkStatsLog.KEYBOARD_SYSTEMS_EVENT_REPORTED__KEYBOARD_SYSTEM_EVENT__CHANGE_SPLITSCREEN_FOCUS,
+                "CHANGE_SPLITSCREEN_FOCUS"),
         TRIGGER_BUG_REPORT(
                 FrameworkStatsLog.KEYBOARD_SYSTEMS_EVENT_REPORTED__KEYBOARD_SYSTEM_EVENT__TRIGGER_BUG_REPORT,
                 "TRIGGER_BUG_REPORT"),
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 266418f..ecd7035 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -3526,6 +3526,9 @@
                         moveFocusedTaskToStageSplit(getTargetDisplayIdForKeyEvent(event),
                                 true /* leftOrTop */);
                         logKeyboardSystemsEvent(event, KeyboardLogEvent.SPLIT_SCREEN_NAVIGATION);
+                    } else if (event.isAltPressed()) {
+                        setSplitscreenFocus(true /* leftOrTop */);
+                        logKeyboardSystemsEvent(event, KeyboardLogEvent.CHANGE_SPLITSCREEN_FOCUS);
                     } else {
                         logKeyboardSystemsEvent(event, KeyboardLogEvent.BACK);
                         injectBackGesture(event.getDownTime());
@@ -3534,11 +3537,17 @@
                 }
                 break;
             case KeyEvent.KEYCODE_DPAD_RIGHT:
-                if (firstDown && event.isMetaPressed() && event.isCtrlPressed()) {
-                    moveFocusedTaskToStageSplit(getTargetDisplayIdForKeyEvent(event),
-                            false /* leftOrTop */);
-                    logKeyboardSystemsEvent(event, KeyboardLogEvent.SPLIT_SCREEN_NAVIGATION);
-                    return true;
+                if (firstDown && event.isMetaPressed()) {
+                    if (event.isCtrlPressed()) {
+                        moveFocusedTaskToStageSplit(getTargetDisplayIdForKeyEvent(event),
+                                false /* leftOrTop */);
+                        logKeyboardSystemsEvent(event, KeyboardLogEvent.SPLIT_SCREEN_NAVIGATION);
+                        return true;
+                    } else if (event.isAltPressed()) {
+                        setSplitscreenFocus(false /* leftOrTop */);
+                        logKeyboardSystemsEvent(event, KeyboardLogEvent.CHANGE_SPLITSCREEN_FOCUS);
+                        return true;
+                    }
                 }
                 break;
             case KeyEvent.KEYCODE_SLASH:
@@ -4398,6 +4407,13 @@
         }
     }
 
+    private void setSplitscreenFocus(boolean leftOrTop) {
+        StatusBarManagerInternal statusbar = getStatusBarManagerInternal();
+        if (statusbar != null) {
+            statusbar.setSplitscreenFocus(leftOrTop);
+        }
+    }
+
     void launchHomeFromHotKey(int displayId) {
         launchHomeFromHotKey(displayId, true /* awakenFromDreams */, true /*respectKeyguard*/);
     }
diff --git a/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java b/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java
index c73f89c..f7c236a 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java
@@ -238,6 +238,14 @@
     void moveFocusedTaskToStageSplit(int displayId, boolean leftOrTop);
 
     /**
+     * Change the split screen focus to the left / top app or the right / bottom app based on
+     * {@param leftOrTop}.
+     *
+     * @see com.android.internal.statusbar.IStatusBar#setSplitscreenFocus
+     */
+    void setSplitscreenFocus(boolean leftOrTop);
+
+    /**
      * Shows the media output switcher dialog.
      *
      * @param packageName of the session for which the output switcher is shown.
diff --git a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
index 214dbe0..7b3e237 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
@@ -830,6 +830,15 @@
         }
 
         @Override
+        public void setSplitscreenFocus(boolean leftOrTop) {
+            IStatusBar bar = mBar;
+            if (bar != null) {
+                try {
+                    bar.setSplitscreenFocus(leftOrTop);
+                } catch (RemoteException ex) { }
+            }
+        }
+        @Override
         public void enterDesktop(int displayId) {
             IStatusBar bar = mBar;
             if (bar != null) {
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 46bac16..2b337ae 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -5700,7 +5700,8 @@
             // window becomes visible while the sync group is still active.
             return true;
         }
-        if (mSyncState == SYNC_STATE_WAITING_FOR_DRAW && mLastConfigReportedToClient && isDrawn()) {
+        if (mSyncState == SYNC_STATE_WAITING_FOR_DRAW && mLastConfigReportedToClient && isDrawn()
+                && mPrepareSyncSeqId <= 0) {
             // Complete the sync state immediately for a drawn window that doesn't need to redraw.
             onSyncFinishedDrawing();
         }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/BrightnessObserverTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/BrightnessObserverTest.kt
index 638924e..b182cce 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/BrightnessObserverTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/BrightnessObserverTest.kt
@@ -95,9 +95,9 @@
     ) {
         ALL_ENABLED(true, true, CombinedVote(
                 listOf(DisableRefreshRateSwitchingVote(true),
-                        SupportedModesVote(
-                                listOf(SupportedModesVote.SupportedMode(60f, 60f),
-                                        SupportedModesVote.SupportedMode(120f, 120f)))))),
+                        SupportedRefreshRatesVote(
+                                listOf(SupportedRefreshRatesVote.RefreshRates(60f, 60f),
+                                        SupportedRefreshRatesVote.RefreshRates(120f, 120f)))))),
         VRR_NOT_SUPPORTED(false, true, DisableRefreshRateSwitchingVote(true)),
         VSYNC_VOTE_DISABLED(true, false, DisableRefreshRateSwitchingVote(true))
     }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/ProximitySensorObserverTest.java b/services/tests/displayservicetests/src/com/android/server/display/mode/ProximitySensorObserverTest.java
new file mode 100644
index 0000000..e93e5bc
--- /dev/null
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/ProximitySensorObserverTest.java
@@ -0,0 +1,127 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.mode;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.when;
+
+import android.hardware.display.DisplayManagerInternal;
+import android.util.SparseArray;
+import android.view.Display;
+import android.view.SurfaceControl;
+
+import androidx.test.filters.SmallTest;
+
+import com.android.server.sensors.SensorManagerInternal;
+
+import junitparams.JUnitParamsRunner;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+@SmallTest
+@RunWith(JUnitParamsRunner.class)
+public class ProximitySensorObserverTest {
+
+    private static final float FLOAT_TOLERANCE = 0.01f;
+    private static final int DISPLAY_ID = 1;
+    private static final SurfaceControl.RefreshRateRange REFRESH_RATE_RANGE =
+            new SurfaceControl.RefreshRateRange(60, 90);
+
+    private final VotesStorage mStorage = new VotesStorage(() -> { }, null);
+    private final FakesInjector mInjector = new FakesInjector();
+    private ProximitySensorObserver mSensorObserver;
+
+    @Mock
+    DisplayManagerInternal mMockDisplayManagerInternal;
+    @Mock
+    SensorManagerInternal mMockSensorManagerInternal;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        when(mMockDisplayManagerInternal.getRefreshRateForDisplayAndSensor(eq(DISPLAY_ID),
+                any(), any())).thenReturn(REFRESH_RATE_RANGE);
+        mSensorObserver = new ProximitySensorObserver(mStorage, mInjector);
+        mSensorObserver.observe();
+    }
+
+    @Test
+    public void testAddsProximityVoteIfSensorManagerProximityActive() {
+        mSensorObserver.onProximityActive(true);
+
+        SparseArray<Vote> displayVotes = mStorage.getVotes(DISPLAY_ID);
+        assertThat(displayVotes.size()).isEqualTo(1);
+        Vote vote = displayVotes.get(Vote.PRIORITY_PROXIMITY);
+        assertThat(vote).isNotNull();
+        assertThat(vote).isInstanceOf(CombinedVote.class);
+        CombinedVote combinedVote = (CombinedVote) vote;
+        RefreshRateVote.PhysicalVote physicalVote =
+                (RefreshRateVote.PhysicalVote) combinedVote.mVotes.get(0);
+        assertThat(physicalVote.mMinRefreshRate).isWithin(FLOAT_TOLERANCE).of(60);
+        assertThat(physicalVote.mMaxRefreshRate).isWithin(FLOAT_TOLERANCE).of(90);
+    }
+
+    @Test
+    public void testDoesNotAddProximityVoteIfSensorManagerProximityNotActive() {
+        mSensorObserver.onProximityActive(false);
+
+        SparseArray<Vote> displayVotes = mStorage.getVotes(DISPLAY_ID);
+        assertThat(displayVotes.size()).isEqualTo(0);
+    }
+
+    @Test
+    public void testDoesNotAddProximityVoteIfDoze() {
+        mInjector.mDozeState = true;
+        mSensorObserver.onDisplayChanged(DISPLAY_ID);
+        mSensorObserver.onProximityActive(true);
+
+        SparseArray<Vote> displayVotes = mStorage.getVotes(DISPLAY_ID);
+        assertThat(displayVotes.size()).isEqualTo(0);
+    }
+
+    private class FakesInjector extends DisplayModeDirectorTest.FakesInjector {
+
+        private boolean mDozeState = false;
+
+        @Override
+        public Display[] getDisplays() {
+            return new Display[] { createDisplay(DISPLAY_ID) };
+        }
+
+        @Override
+        public DisplayManagerInternal getDisplayManagerInternal() {
+            return mMockDisplayManagerInternal;
+        }
+
+        @Override
+        public SensorManagerInternal getSensorManagerInternal() {
+            return mMockSensorManagerInternal;
+        }
+
+        @Override
+        public boolean isDozeState(Display d) {
+            return mDozeState;
+        }
+    }
+}
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt
index ebb4f18..230317b 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt
@@ -85,9 +85,9 @@
             internal val expectedVote: Vote?
     ) {
         ALL_ENABLED(true, true, true,
-                SupportedModesVote(listOf(
-                        SupportedModesVote.SupportedMode(60f, 240f),
-                        SupportedModesVote.SupportedMode(60f, 60f)
+                SupportedRefreshRatesVote(listOf(
+                        SupportedRefreshRatesVote.RefreshRates(60f, 240f),
+                        SupportedRefreshRatesVote.RefreshRates(60f, 60f)
                 ))),
         LOW_POWER_OFF(true, true, false, null),
         DVRR_NOT_SUPPORTED_LOW_POWER_ON(false, true, true,
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt
index 04e6265..6ce49b8 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2023 The Android Open Source Project
+ * 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.
@@ -27,12 +27,9 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class SupportedModesVoteTest {
-    private val supportedModes = listOf(
-            SupportedModesVote.SupportedMode(60f, 90f ),
-            SupportedModesVote.SupportedMode(120f, 240f )
-    )
+    private val supportedModes = listOf(1, 2, 4)
 
-    private val otherMode = SupportedModesVote.SupportedMode(120f, 120f )
+    private val otherMode = 5
 
     private lateinit var supportedModesVote: SupportedModesVote
 
@@ -42,31 +39,31 @@
     }
 
     @Test
-    fun `adds supported modes if supportedModes in summary is null`() {
+    fun `adds supported mode ids if supportedModeIds in summary is null`() {
         val summary = createVotesSummary()
 
         supportedModesVote.updateSummary(summary)
 
-        assertThat(summary.supportedModes).containsExactlyElementsIn(supportedModes)
+        assertThat(summary.supportedModeIds).containsExactlyElementsIn(supportedModes)
     }
 
     @Test
-    fun `does not add supported modes if summary has empty list of modes`() {
+    fun `does not add supported mode ids if summary has empty list of modeIds`() {
         val summary = createVotesSummary()
-        summary.supportedModes = ArrayList()
+        summary.supportedModeIds = ArrayList()
 
         supportedModesVote.updateSummary(summary)
 
-        assertThat(summary.supportedModes).isEmpty()
+        assertThat(summary.supportedModeIds).isEmpty()
     }
 
     @Test
     fun `filters out modes that does not match vote`() {
         val summary = createVotesSummary()
-        summary.supportedModes = ArrayList(listOf(otherMode, supportedModes[0]))
+        summary.supportedModeIds = ArrayList(listOf(otherMode, supportedModes[0]))
 
         supportedModesVote.updateSummary(summary)
 
-        assertThat(summary.supportedModes).containsExactly(supportedModes[0])
+        assertThat(summary.supportedModeIds).containsExactly(supportedModes[0])
     }
 }
\ No newline at end of file
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedRefreshRatesVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedRefreshRatesVoteTest.kt
new file mode 100644
index 0000000..d0c112b
--- /dev/null
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedRefreshRatesVoteTest.kt
@@ -0,0 +1,72 @@
+/*
+ * 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.mode
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class SupportedRefreshRatesVoteTest {
+    private val refreshRates = listOf(
+            SupportedRefreshRatesVote.RefreshRates(60f, 90f),
+            SupportedRefreshRatesVote.RefreshRates(120f, 240f)
+    )
+
+    private val otherMode = SupportedRefreshRatesVote.RefreshRates(120f, 120f)
+
+    private lateinit var supportedRefreshRatesVote: SupportedRefreshRatesVote
+
+    @Before
+    fun setUp() {
+        supportedRefreshRatesVote = SupportedRefreshRatesVote(refreshRates)
+    }
+
+    @Test
+    fun `adds supported refresh rates if supportedModes in summary is null`() {
+        val summary = createVotesSummary()
+
+        supportedRefreshRatesVote.updateSummary(summary)
+
+        assertThat(summary.supportedRefreshRates).containsExactlyElementsIn(refreshRates)
+    }
+
+    @Test
+    fun `does not add supported refresh rates if summary has empty list of refresh rates`() {
+        val summary = createVotesSummary()
+        summary.supportedRefreshRates = ArrayList()
+
+        supportedRefreshRatesVote.updateSummary(summary)
+
+        assertThat(summary.supportedRefreshRates).isEmpty()
+    }
+
+    @Test
+    fun `filters out supported refresh rates that does not match vote`() {
+        val summary = createVotesSummary()
+        summary.supportedRefreshRates = ArrayList(listOf(otherMode, refreshRates[0]))
+
+        supportedRefreshRatesVote.updateSummary(summary)
+
+        assertThat(summary.supportedRefreshRates).containsExactly(refreshRates[0])
+    }
+}
\ No newline at end of file
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SystemRequestObserverTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SystemRequestObserverTest.kt
new file mode 100644
index 0000000..c49205b
--- /dev/null
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SystemRequestObserverTest.kt
@@ -0,0 +1,212 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.mode
+
+import android.os.IBinder
+import android.os.RemoteException
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import com.google.testing.junit.testparameterinjector.TestParameterInjector
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.junit.MockitoJUnit
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doThrow
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+private const val DISPLAY_ID = 1
+private const val DISPLAY_ID_OTHER = 2
+
+@SmallTest
+@RunWith(TestParameterInjector::class)
+class SystemRequestObserverTest {
+
+
+    @get:Rule
+    val mockitoRule = MockitoJUnit.rule()
+
+    private val mockToken = mock<IBinder>()
+    private val mockOtherToken = mock<IBinder>()
+
+    private val storage = VotesStorage({}, null)
+
+    @Test
+    fun `requestDisplayModes adds vote to storage`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+
+        val votes = storage.getVotes(DISPLAY_ID)
+        assertThat(votes.size()).isEqualTo(1)
+        val vote = votes.get(Vote.PRIORITY_SYSTEM_REQUESTED_MODES)
+        assertThat(vote).isInstanceOf(SupportedModesVote::class.java)
+        val supportedModesVote = vote as SupportedModesVote
+        assertThat(supportedModesVote.mModeIds.size).isEqualTo(requestedModes.size)
+        for (mode in requestedModes) {
+            assertThat(supportedModesVote.mModeIds).contains(mode)
+        }
+    }
+
+    @Test
+    fun `requestDisplayModes overrides votes in storage`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, intArrayOf(1, 2, 3))
+
+        val overrideModes = intArrayOf(10, 20, 30)
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, overrideModes)
+
+        val votes = storage.getVotes(DISPLAY_ID)
+        assertThat(votes.size()).isEqualTo(1)
+        val vote = votes.get(Vote.PRIORITY_SYSTEM_REQUESTED_MODES)
+        assertThat(vote).isInstanceOf(SupportedModesVote::class.java)
+        val supportedModesVote = vote as SupportedModesVote
+        assertThat(supportedModesVote.mModeIds.size).isEqualTo(overrideModes.size)
+        for (mode in overrideModes) {
+            assertThat(supportedModesVote.mModeIds).contains(mode)
+        }
+    }
+
+    @Test
+    fun `requestDisplayModes removes vote to storage`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, null)
+
+        val votes = storage.getVotes(DISPLAY_ID)
+        assertThat(votes.size()).isEqualTo(0)
+    }
+
+    @Test
+    fun `requestDisplayModes calls linkToDeath to token`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+
+        verify(mockToken).linkToDeath(any(), eq(0))
+    }
+
+    @Test
+    fun `does not add votes to storage if binder died when requestDisplayModes called`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+
+        doThrow(RemoteException()).whenever(mockOtherToken).linkToDeath(any(), eq(0))
+        systemRequestObserver.requestDisplayModes(mockOtherToken, DISPLAY_ID, requestedModes)
+
+        val votes = storage.getVotes(DISPLAY_ID)
+        assertThat(votes.size()).isEqualTo(0)
+    }
+
+    @Test
+    fun `removes all votes from storage when binder dies`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+        val deathRecipientCaptor = argumentCaptor<IBinder.DeathRecipient>()
+        verify(mockToken).linkToDeath(deathRecipientCaptor.capture(), eq(0))
+
+        deathRecipientCaptor.lastValue.binderDied(mockToken)
+
+        val votes = storage.getVotes(DISPLAY_ID)
+        assertThat(votes.size()).isEqualTo(0)
+    }
+
+    @Test
+    fun `calls unlinkToDeath on token when no votes remaining`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, null)
+
+        verify(mockToken).unlinkToDeath(any(), eq(0))
+    }
+
+    @Test
+    fun `does not call unlinkToDeath on token when votes for other display in storage`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID_OTHER, requestedModes)
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, null)
+
+        verify(mockToken, never()).unlinkToDeath(any(), eq(0))
+    }
+
+    @Test
+    fun `requestDisplayModes subset modes from different tokens`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+
+        val requestedOtherModes = intArrayOf(2, 3, 4)
+        systemRequestObserver.requestDisplayModes(mockOtherToken, DISPLAY_ID, requestedOtherModes)
+
+        verify(mockToken).linkToDeath(any(), eq(0))
+        verify(mockOtherToken).linkToDeath(any(), eq(0))
+        verify(mockToken, never()).unlinkToDeath(any(), eq(0))
+        verify(mockOtherToken, never()).unlinkToDeath(any(), eq(0))
+
+        val expectedModes = intArrayOf(2, 3)
+        val votes = storage.getVotes(DISPLAY_ID)
+        assertThat(votes.size()).isEqualTo(1)
+        val vote = votes.get(Vote.PRIORITY_SYSTEM_REQUESTED_MODES)
+        assertThat(vote).isInstanceOf(SupportedModesVote::class.java)
+        val supportedModesVote = vote as SupportedModesVote
+        assertThat(supportedModesVote.mModeIds.size).isEqualTo(expectedModes.size)
+        for (mode in expectedModes) {
+            assertThat(supportedModesVote.mModeIds).contains(mode)
+        }
+    }
+
+    @Test
+    fun `recalculates vote if one binder dies`() {
+        val systemRequestObserver = SystemRequestObserver(storage)
+        val requestedModes = intArrayOf(1, 2, 3)
+        systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
+
+        val requestedOtherModes = intArrayOf(2, 3, 4)
+        systemRequestObserver.requestDisplayModes(mockOtherToken, DISPLAY_ID, requestedOtherModes)
+
+        val deathRecipientCaptor = argumentCaptor<IBinder.DeathRecipient>()
+        verify(mockOtherToken).linkToDeath(deathRecipientCaptor.capture(), eq(0))
+        deathRecipientCaptor.lastValue.binderDied(mockOtherToken)
+
+        val votes = storage.getVotes(DISPLAY_ID)
+        assertThat(votes.size()).isEqualTo(1)
+        val vote = votes.get(Vote.PRIORITY_SYSTEM_REQUESTED_MODES)
+        assertThat(vote).isInstanceOf(SupportedModesVote::class.java)
+        val supportedModesVote = vote as SupportedModesVote
+        assertThat(supportedModesVote.mModeIds.size).isEqualTo(requestedModes.size)
+        for (mode in requestedModes) {
+            assertThat(supportedModesVote.mModeIds).contains(mode)
+        }
+    }
+}
\ No newline at end of file
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/TestUtils.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/TestUtils.kt
index 910e03c..6b90bde 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/TestUtils.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/TestUtils.kt
@@ -18,10 +18,10 @@
 
 internal fun createVotesSummary(
         isDisplayResolutionRangeVotingEnabled: Boolean = true,
-        vsyncProximityVoteEnabled: Boolean = true,
+        supportedModesVoteEnabled: Boolean = true,
         loggingEnabled: Boolean = true,
         supportsFrameRateOverride: Boolean = true
 ): VoteSummary {
-    return VoteSummary(isDisplayResolutionRangeVotingEnabled, vsyncProximityVoteEnabled,
+    return VoteSummary(isDisplayResolutionRangeVotingEnabled, supportedModesVoteEnabled,
             loggingEnabled, supportsFrameRateOverride)
 }
\ No newline at end of file
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt
index d6c8469..04b35f1 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt
@@ -28,29 +28,29 @@
 @RunWith(TestParameterInjector::class)
 class VoteSummaryTest {
 
-    enum class SupportedModesVoteTestCase(
-            val vsyncProximityVoteEnabled: Boolean,
-            internal val summarySupportedModes: List<SupportedModesVote.SupportedMode>?,
+    enum class SupportedRefreshRatesTestCase(
+            val supportedModesVoteEnabled: Boolean,
+            internal val summaryRefreshRates: List<SupportedRefreshRatesVote.RefreshRates>?,
             val modesToFilter: Array<Display.Mode>,
             val expectedModeIds: List<Int>
     ) {
         HAS_NO_MATCHING_VOTE(true,
-                listOf(SupportedModesVote.SupportedMode(60f, 60f)),
+                listOf(SupportedRefreshRatesVote.RefreshRates(60f, 60f)),
                 arrayOf(createMode(1, 90f, 90f),
                         createMode(2, 90f, 60f),
                         createMode(3, 60f, 90f)),
                 listOf()
         ),
         HAS_SINGLE_MATCHING_VOTE(true,
-                listOf(SupportedModesVote.SupportedMode(60f, 90f)),
+                listOf(SupportedRefreshRatesVote.RefreshRates(60f, 90f)),
                 arrayOf(createMode(1, 90f, 90f),
                         createMode(2, 90f, 60f),
                         createMode(3, 60f, 90f)),
                 listOf(3)
         ),
         HAS_MULTIPLE_MATCHING_VOTES(true,
-                listOf(SupportedModesVote.SupportedMode(60f, 90f),
-                        SupportedModesVote.SupportedMode(90f, 90f)),
+                listOf(SupportedRefreshRatesVote.RefreshRates(60f, 90f),
+                        SupportedRefreshRatesVote.RefreshRates(90f, 90f)),
                 arrayOf(createMode(1, 90f, 90f),
                         createMode(2, 90f, 60f),
                         createMode(3, 60f, 90f)),
@@ -70,7 +70,69 @@
                         createMode(3, 60f, 90f)),
                 listOf(1, 2, 3)
         ),
-        HAS_VSYNC_PROXIMITY_DISABLED(false,
+        HAS_SUPPORTED_MODES_VOTE_DISABLED(false,
+                listOf(),
+                arrayOf(createMode(1, 90f, 90f),
+                        createMode(2, 90f, 60f),
+                        createMode(3, 60f, 90f)),
+                listOf(1, 2, 3)
+        ),
+    }
+
+    @Test
+    fun `filters modes for summary supportedRefreshRates`(
+            @TestParameter testCase: SupportedRefreshRatesTestCase
+    ) {
+        val summary = createSummary(testCase.supportedModesVoteEnabled)
+        summary.supportedRefreshRates = testCase.summaryRefreshRates
+
+        val result = summary.filterModes(testCase.modesToFilter)
+
+        assertThat(result.map { it.modeId }).containsExactlyElementsIn(testCase.expectedModeIds)
+    }
+
+    enum class SupportedModesTestCase(
+            val supportedModesVoteEnabled: Boolean,
+            internal val summarySupportedModes: List<Int>?,
+            val modesToFilter: Array<Display.Mode>,
+            val expectedModeIds: List<Int>
+    ) {
+        HAS_NO_MATCHING_VOTE(true,
+                listOf(4, 5),
+                arrayOf(createMode(1, 90f, 90f),
+                        createMode(2, 90f, 60f),
+                        createMode(3, 60f, 90f)),
+                listOf()
+        ),
+        HAS_SINGLE_MATCHING_VOTE(true,
+                listOf(3),
+                arrayOf(createMode(1, 90f, 90f),
+                        createMode(2, 90f, 60f),
+                        createMode(3, 60f, 90f)),
+                listOf(3)
+        ),
+        HAS_MULTIPLE_MATCHING_VOTES(true,
+                listOf(1, 3),
+                arrayOf(createMode(1, 90f, 90f),
+                        createMode(2, 90f, 60f),
+                        createMode(3, 60f, 90f)),
+                listOf(1, 3)
+        ),
+        HAS_NO_SUPPORTED_MODES(true,
+                listOf(),
+                arrayOf(createMode(1, 90f, 90f),
+                        createMode(2, 90f, 60f),
+                        createMode(3, 60f, 90f)),
+                listOf()
+        ),
+        HAS_NULL_SUPPORTED_MODES(true,
+                null,
+                arrayOf(createMode(1, 90f, 90f),
+                        createMode(2, 90f, 60f),
+                        createMode(3, 60f, 90f)),
+                listOf(1, 2, 3)
+        ),
+        HAS_SUPPORTED_MODES_VOTE_DISABLED(false,
                 listOf(),
                 arrayOf(createMode(1, 90f, 90f),
                         createMode(2, 90f, 60f),
@@ -81,10 +143,10 @@
 
     @Test
     fun `filters modes for summary supportedModes`(
-            @TestParameter testCase: SupportedModesVoteTestCase
+            @TestParameter testCase: SupportedModesTestCase
     ) {
-        val summary = createSummary(testCase.vsyncProximityVoteEnabled)
-        summary.supportedModes = testCase.summarySupportedModes
+        val summary = createSummary(testCase.supportedModesVoteEnabled)
+        summary.supportedModeIds = testCase.summarySupportedModes
 
         val result = summary.filterModes(testCase.modesToFilter)
 
@@ -96,8 +158,8 @@
             FloatArray(0), IntArray(0))
 }
 
-private fun createSummary(vsyncVoteEnabled: Boolean): VoteSummary {
-    val summary = createVotesSummary(vsyncProximityVoteEnabled = vsyncVoteEnabled)
+private fun createSummary(supportedModesVoteEnabled: Boolean): VoteSummary {
+    val summary = createVotesSummary(supportedModesVoteEnabled = supportedModesVoteEnabled)
     summary.width = 600
     summary.height = 800
     summary.maxPhysicalRefreshRate = Float.POSITIVE_INFINITY
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/VotesStorageTest.java b/services/tests/displayservicetests/src/com/android/server/display/mode/VotesStorageTest.java
index 1f6f1a4..a248d6de 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/VotesStorageTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/VotesStorageTest.java
@@ -18,6 +18,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -153,4 +154,51 @@
         assertThat(mVotesStorage.getVotes(DISPLAY_ID).size()).isEqualTo(0);
         verify(mVotesListener, never()).onChanged();
     }
+
+
+    @Test
+    public void removesAllVotesForPriority() {
+        // GIVEN vote storage with votes
+        mVotesStorage.updateVote(DISPLAY_ID, PRIORITY, VOTE);
+        mVotesStorage.updateVote(DISPLAY_ID, PRIORITY_OTHER, VOTE_OTHER);
+        mVotesStorage.updateVote(DISPLAY_ID_OTHER, PRIORITY, VOTE);
+        mVotesStorage.updateVote(DISPLAY_ID_OTHER, PRIORITY_OTHER, VOTE_OTHER);
+        // WHEN removeAllVotesForPriority is called
+        mVotesStorage.removeAllVotesForPriority(PRIORITY);
+        // THEN votes with priority are removed from the storage
+        SparseArray<Vote> votes = mVotesStorage.getVotes(DISPLAY_ID);
+        assertThat(votes.size()).isEqualTo(1);
+        assertThat(votes.get(PRIORITY)).isNull();
+        votes = mVotesStorage.getVotes(DISPLAY_ID_OTHER);
+        assertThat(votes.size()).isEqualTo(1);
+        assertThat(votes.get(PRIORITY)).isNull();
+    }
+
+    @Test
+    public void removesAllVotesForPriority_notifiesListenerOnce() {
+        // GIVEN vote storage with votes
+        mVotesStorage.updateVote(DISPLAY_ID, PRIORITY, VOTE);
+        mVotesStorage.updateVote(DISPLAY_ID, PRIORITY_OTHER, VOTE_OTHER);
+        mVotesStorage.updateVote(DISPLAY_ID_OTHER, PRIORITY, VOTE);
+        mVotesStorage.updateVote(DISPLAY_ID_OTHER, PRIORITY_OTHER, VOTE_OTHER);
+        clearInvocations(mVotesListener);
+        // WHEN removeAllVotesForPriority is called
+        mVotesStorage.removeAllVotesForPriority(PRIORITY);
+        // THEN listener notified once
+        verify(mVotesListener).onChanged();
+    }
+
+    @Test
+    public void removesAllVotesForPriority_noChangesIfNothingRemoved() {
+        // GIVEN vote storage with votes
+        mVotesStorage.updateVote(DISPLAY_ID, PRIORITY, VOTE);
+        clearInvocations(mVotesListener);
+        // WHEN removeAllVotesForPriority is called for missing priority
+        mVotesStorage.removeAllVotesForPriority(PRIORITY_OTHER);
+        // THEN no changes to votes storage
+        SparseArray<Vote> votes = mVotesStorage.getVotes(DISPLAY_ID);
+        assertThat(votes.size()).isEqualTo(1);
+        assertThat(votes.get(PRIORITY)).isEqualTo(VOTE);
+        verify(mVotesListener, never()).onChanged();
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index 03b695d..43b424f 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -1376,8 +1376,9 @@
         assertTrue(w1.syncNextBuffer());
         assertTrue(w2.syncNextBuffer());
 
-        // A drawn window can complete the sync state automatically.
+        // A drawn window in non-explicit sync can complete the sync state automatically.
         w1.mWinAnimator.mDrawState = WindowStateAnimator.HAS_DRAWN;
+        w1.mPrepareSyncSeqId = 0;
         makeLastConfigReportedToClient(w1, true /* visible */);
         mWm.mSyncEngine.onSurfacePlacement();
         verify(mockCallback).onTransactionReady(anyInt(), any());