Merge "Fix NPE when reading the allowlist" into main
diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 5e9af6a..a16aa2d 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -411,6 +411,13 @@
     host_supported: true,
 }
 
+cc_aconfig_library {
+    name: "android.companion.virtualdevice.flags-aconfig-cc-test",
+    aconfig_declarations: "android.companion.virtualdevice.flags-aconfig",
+    host_supported: true,
+    mode: "test",
+}
+
 java_aconfig_library {
     name: "android.companion.virtualdevice.flags-aconfig-java",
     aconfig_declarations: "android.companion.virtualdevice.flags-aconfig",
diff --git a/core/java/android/companion/virtual/flags/flags.aconfig b/core/java/android/companion/virtual/flags/flags.aconfig
index b63e2cf..d05d23c 100644
--- a/core/java/android/companion/virtual/flags/flags.aconfig
+++ b/core/java/android/companion/virtual/flags/flags.aconfig
@@ -103,3 +103,10 @@
     description: "API for on-demand rotation of virtual displays"
     bug: "291748430"
 }
+
+flag {
+    namespace: "virtual_devices"
+    name: "high_resolution_scroll"
+    description: "Enable high resolution scroll"
+    bug: "335160780"
+}
diff --git a/core/java/com/android/internal/os/BinderTransactionNameResolver.java b/core/java/com/android/internal/os/BinderTransactionNameResolver.java
index 5f6f427..f1dc1f2 100644
--- a/core/java/com/android/internal/os/BinderTransactionNameResolver.java
+++ b/core/java/com/android/internal/os/BinderTransactionNameResolver.java
@@ -17,6 +17,7 @@
 package com.android.internal.os;
 
 import android.os.Binder;
+import android.text.TextUtils;
 
 import com.android.internal.annotations.VisibleForTesting;
 
@@ -37,6 +38,8 @@
 @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
 public class BinderTransactionNameResolver {
     private static final Method NO_GET_DEFAULT_TRANSACTION_NAME_METHOD;
+    private static final boolean USE_TRANSACTION_CODES_FOR_UNKNOWN_METHODS =
+            Flags.useTransactionCodesForUnknownMethods();
 
     /**
      * Generates the default transaction method name, which is just the transaction code.
@@ -81,7 +84,14 @@
         }
 
         try {
-            return (String) method.invoke(null, transactionCode);
+            String methodName = (String) method.invoke(null, transactionCode);
+            if (USE_TRANSACTION_CODES_FOR_UNKNOWN_METHODS) {
+                return TextUtils.isEmpty(methodName)
+                        ? String.valueOf(transactionCode)
+                        : methodName;
+            } else {
+                return methodName;
+            }
         } catch (IllegalAccessException | InvocationTargetException e) {
             throw new RuntimeException(e);
         }
diff --git a/core/java/com/android/internal/os/flags.aconfig b/core/java/com/android/internal/os/flags.aconfig
index c8d6810..30fa4f1 100644
--- a/core/java/com/android/internal/os/flags.aconfig
+++ b/core/java/com/android/internal/os/flags.aconfig
@@ -9,3 +9,14 @@
     is_fixed_read_only: true
     bug: "241474956"
 }
+
+flag {
+    name: "use_transaction_codes_for_unknown_methods"
+    namespace: "dropbox"
+    description: "Use transaction codes when the method names is unknown"
+    bug: "350041302"
+    is_fixed_read_only: true
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
\ No newline at end of file
diff --git a/core/java/com/android/internal/widget/NotificationRowIconView.java b/core/java/com/android/internal/widget/NotificationRowIconView.java
index f5f04a7..98e6e85 100644
--- a/core/java/com/android/internal/widget/NotificationRowIconView.java
+++ b/core/java/com/android/internal/widget/NotificationRowIconView.java
@@ -35,8 +35,10 @@
 import android.view.RemotableViewMethod;
 import android.widget.RemoteViews;
 
+import com.android.internal.R;
+
 /**
- * An image view that holds the icon displayed on the left side of a notification row.
+ * An image view that holds the icon displayed at the start of a notification row.
  */
 @RemoteViews.RemoteView
 public class NotificationRowIconView extends CachingIconView {
@@ -98,9 +100,12 @@
                 setPadding(0, 0, 0, 0);
 
                 // Make the background white in case the icon itself doesn't have one.
-                int white = Color.rgb(255, 255, 255);
-                ColorFilter colorFilter = new PorterDuffColorFilter(white,
+                ColorFilter colorFilter = new PorterDuffColorFilter(Color.WHITE,
                         PorterDuff.Mode.SRC_ATOP);
+
+                if (mOriginalBackground == null) {
+                    setBackground(getContext().getDrawable(R.drawable.notification_icon_circle));
+                }
                 getBackground().mutate().setColorFilter(colorFilter);
             } else {
                 // Restore original padding and background if needed
diff --git a/core/res/res/layout/notification_template_conversation_icon_container.xml b/core/res/res/layout/notification_template_conversation_icon_container.xml
index 0438dc5..b45483b 100644
--- a/core/res/res/layout/notification_template_conversation_icon_container.xml
+++ b/core/res/res/layout/notification_template_conversation_icon_container.xml
@@ -77,7 +77,7 @@
                 android:scaleType="center"
                 />
 
-            <com.android.internal.widget.CachingIconView
+            <com.android.internal.widget.NotificationRowIconView
                 android:id="@+id/icon"
                 android:layout_width="match_parent"
                 android:layout_height="match_parent"
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index c9b5d41..419a615 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -4059,6 +4059,7 @@
   <java-symbol type="id" name="snooze_button" />
   <java-symbol type="dimen" name="text_size_body_2_material" />
   <java-symbol type="dimen" name="notification_icon_circle_size" />
+  <java-symbol type="drawable" name="notification_icon_circle" />
   <java-symbol type="dimen" name="messaging_avatar_size" />
   <java-symbol type="dimen" name="messaging_group_sending_progress_size" />
   <java-symbol type="dimen" name="messaging_image_rounding" />
diff --git a/libs/WindowManager/Shell/OWNERS b/libs/WindowManager/Shell/OWNERS
index ebebd8a..cb422ea 100644
--- a/libs/WindowManager/Shell/OWNERS
+++ b/libs/WindowManager/Shell/OWNERS
@@ -1,5 +1,5 @@
 xutan@google.com
 
 # Give submodule owners in shell resource approval
-per-file res*/*/*.xml = atsjenk@google.com, hwwang@google.com, jorgegil@google.com, lbill@google.com, madym@google.com, nmusgrave@google.com, pbdr@google.com, tkachenkoi@google.com, mpodolian@google.com, liranb@google.com
+per-file res*/*/*.xml = atsjenk@google.com, hwwang@google.com, jorgegil@google.com, lbill@google.com, madym@google.com, vaniadesmonda@google.com, pbdr@google.com, tkachenkoi@google.com, mpodolian@google.com, liranb@google.com
 per-file res*/*/tv_*.xml = bronger@google.com
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
index 5983168..8c2cfd1e 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
@@ -37,10 +37,6 @@
   DESKTOP_WINDOWING_MODE(Flags::enableDesktopWindowingMode, true),
   WALLPAPER_ACTIVITY(Flags::enableDesktopWindowingWallpaperActivity, true);
 
-  // Local cache for toggle override, which is initialized once on its first access. It needs to be
-  // refreshed only on reboots as overridden state takes effect on reboots.
-  private var cachedToggleOverride: ToggleOverride? = null
-
   /**
    * Determines state of flag based on the actual flag and desktop mode developer option overrides.
    *
@@ -138,13 +134,19 @@
         }
   }
 
-  private companion object {
-    const val TAG = "DesktopModeFlags"
+  companion object {
+    private const val TAG = "DesktopModeFlags"
 
     /**
      * Key for non-persistent System Property which is used to store desktop windowing developer
      * option overrides.
      */
-    const val SYSTEM_PROPERTY_OVERRIDE_KEY = "sys.wmshell.desktopmode.dev_toggle_override"
+    private const val SYSTEM_PROPERTY_OVERRIDE_KEY = "sys.wmshell.desktopmode.dev_toggle_override"
+
+    /**
+     * Local cache for toggle override, which is initialized once on its first access. It needs to
+     * be refreshed only on reboots as overridden state takes effect on reboots.
+     */
+    private var cachedToggleOverride: ToggleOverride? = null
   }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index ebdea1b..d41600b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -24,6 +24,9 @@
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 
+import static com.android.wm.shell.compatui.impl.CompatUIEventsKt.CAMERA_CONTROL_STATE_UPDATE;
+import static com.android.wm.shell.compatui.impl.CompatUIEventsKt.SIZE_COMPAT_RESTART_BUTTON_APPEARED;
+import static com.android.wm.shell.compatui.impl.CompatUIEventsKt.SIZE_COMPAT_RESTART_BUTTON_CLICKED;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_TASK_ORG;
 import static com.android.wm.shell.transition.Transitions.ENABLE_SHELL_TRANSITIONS;
 
@@ -31,7 +34,6 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.ActivityManager.RunningTaskInfo;
-import android.app.CameraCompatTaskInfo.CameraCompatControlState;
 import android.app.TaskInfo;
 import android.app.WindowConfiguration;
 import android.content.LocusId;
@@ -57,6 +59,11 @@
 import com.android.wm.shell.common.ScreenshotUtils;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.compatui.CompatUIController;
+import com.android.wm.shell.compatui.api.CompatUIHandler;
+import com.android.wm.shell.compatui.api.CompatUIInfo;
+import com.android.wm.shell.compatui.impl.CompatUIEvents.CameraControlStateUpdated;
+import com.android.wm.shell.compatui.impl.CompatUIEvents.SizeCompatRestartButtonAppeared;
+import com.android.wm.shell.compatui.impl.CompatUIEvents.SizeCompatRestartButtonClicked;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.startingsurface.StartingWindowController;
 import com.android.wm.shell.sysui.ShellCommandHandler;
@@ -75,8 +82,7 @@
  * Unified task organizer for all components in the shell.
  * TODO(b/167582004): may consider consolidating this class and TaskOrganizer
  */
-public class ShellTaskOrganizer extends TaskOrganizer implements
-        CompatUIController.CompatUICallback {
+public class ShellTaskOrganizer extends TaskOrganizer {
     private static final String TAG = "ShellTaskOrganizer";
 
     // Intentionally using negative numbers here so the positive numbers can be used
@@ -194,12 +200,11 @@
      * In charge of showing compat UI. Can be {@code null} if the device doesn't support size
      * compat or if this isn't the main {@link ShellTaskOrganizer}.
      *
-     * <p>NOTE: only the main {@link ShellTaskOrganizer} should have a {@link CompatUIController},
-     * and register itself as a {@link CompatUIController.CompatUICallback}. Subclasses should be
-     * initialized with a {@code null} {@link CompatUIController}.
+     * <p>NOTE: only the main {@link ShellTaskOrganizer} should have a {@link CompatUIHandler},
+     * Subclasses should be initialized with a {@code null} {@link CompatUIHandler}.
      */
     @Nullable
-    private final CompatUIController mCompatUI;
+    private final CompatUIHandler mCompatUI;
 
     @NonNull
     private final ShellCommandHandler mShellCommandHandler;
@@ -223,7 +228,7 @@
 
     public ShellTaskOrganizer(ShellInit shellInit,
             ShellCommandHandler shellCommandHandler,
-            @Nullable CompatUIController compatUI,
+            @Nullable CompatUIHandler compatUI,
             Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasks,
             ShellExecutor mainExecutor) {
@@ -235,7 +240,7 @@
     protected ShellTaskOrganizer(ShellInit shellInit,
             ShellCommandHandler shellCommandHandler,
             ITaskOrganizerController taskOrganizerController,
-            @Nullable CompatUIController compatUI,
+            @Nullable CompatUIHandler compatUI,
             Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasks,
             ShellExecutor mainExecutor) {
@@ -252,7 +257,21 @@
     private void onInit() {
         mShellCommandHandler.addDumpCallback(this::dump, this);
         if (mCompatUI != null) {
-            mCompatUI.setCompatUICallback(this);
+            mCompatUI.setCallback(compatUIEvent -> {
+                switch(compatUIEvent.getEventId()) {
+                    case SIZE_COMPAT_RESTART_BUTTON_APPEARED:
+                        onSizeCompatRestartButtonAppeared(compatUIEvent.asType());
+                        break;
+                    case SIZE_COMPAT_RESTART_BUTTON_CLICKED:
+                        onSizeCompatRestartButtonClicked(compatUIEvent.asType());
+                        break;
+                    case CAMERA_CONTROL_STATE_UPDATE:
+                        onCameraControlStateUpdated(compatUIEvent.asType());
+                        break;
+                    default:
+
+                }
+            });
         }
         registerOrganizer();
     }
@@ -727,45 +746,6 @@
         }
     }
 
-    @Override
-    public void onSizeCompatRestartButtonAppeared(int taskId) {
-        final TaskAppearedInfo info;
-        synchronized (mLock) {
-            info = mTasks.get(taskId);
-        }
-        if (info == null) {
-            return;
-        }
-        logSizeCompatRestartButtonEventReported(info,
-                FrameworkStatsLog.SIZE_COMPAT_RESTART_BUTTON_EVENT_REPORTED__EVENT__APPEARED);
-    }
-
-    @Override
-    public void onSizeCompatRestartButtonClicked(int taskId) {
-        final TaskAppearedInfo info;
-        synchronized (mLock) {
-            info = mTasks.get(taskId);
-        }
-        if (info == null) {
-            return;
-        }
-        logSizeCompatRestartButtonEventReported(info,
-                FrameworkStatsLog.SIZE_COMPAT_RESTART_BUTTON_EVENT_REPORTED__EVENT__CLICKED);
-        restartTaskTopActivityProcessIfVisible(info.getTaskInfo().token);
-    }
-
-    @Override
-    public void onCameraControlStateUpdated(int taskId, @CameraCompatControlState int state) {
-        final TaskAppearedInfo info;
-        synchronized (mLock) {
-            info = mTasks.get(taskId);
-        }
-        if (info == null) {
-            return;
-        }
-        updateCameraCompatControlState(info.getTaskInfo().token, state);
-    }
-
     /** Reparents a child window surface to the task surface. */
     public void reparentChildSurfaceToTask(int taskId, SurfaceControl sc,
             SurfaceControl.Transaction t) {
@@ -783,6 +763,50 @@
         taskListener.reparentChildSurfaceToTask(taskId, sc, t);
     }
 
+    @VisibleForTesting
+    void onSizeCompatRestartButtonAppeared(@NonNull SizeCompatRestartButtonAppeared compatUIEvent) {
+        final int taskId = compatUIEvent.getTaskId();
+        final TaskAppearedInfo info;
+        synchronized (mLock) {
+            info = mTasks.get(taskId);
+        }
+        if (info == null) {
+            return;
+        }
+        logSizeCompatRestartButtonEventReported(info,
+                FrameworkStatsLog.SIZE_COMPAT_RESTART_BUTTON_EVENT_REPORTED__EVENT__APPEARED);
+    }
+
+    @VisibleForTesting
+    void onSizeCompatRestartButtonClicked(@NonNull SizeCompatRestartButtonClicked compatUIEvent) {
+        final int taskId = compatUIEvent.getTaskId();
+        final TaskAppearedInfo info;
+        synchronized (mLock) {
+            info = mTasks.get(taskId);
+        }
+        if (info == null) {
+            return;
+        }
+        logSizeCompatRestartButtonEventReported(info,
+                FrameworkStatsLog.SIZE_COMPAT_RESTART_BUTTON_EVENT_REPORTED__EVENT__CLICKED);
+        restartTaskTopActivityProcessIfVisible(info.getTaskInfo().token);
+    }
+
+    @VisibleForTesting
+    void onCameraControlStateUpdated(@NonNull CameraControlStateUpdated compatUIEvent) {
+        final int taskId = compatUIEvent.getTaskId();
+        final int state = compatUIEvent.getState();
+        final TaskAppearedInfo info;
+        synchronized (mLock) {
+            info = mTasks.get(taskId);
+        }
+        if (info == null) {
+            return;
+        }
+        updateCameraCompatControlState(info.getTaskInfo().token, state);
+    }
+
+
     private void logSizeCompatRestartButtonEventReported(@NonNull TaskAppearedInfo info,
             int event) {
         ActivityInfo topActivityInfo = info.getTaskInfo().topActivityInfo;
@@ -810,10 +834,10 @@
         // on this Task if there is any.
         if (taskListener == null || !taskListener.supportCompatUI()
                 || !taskInfo.appCompatTaskInfo.hasCompatUI() || !taskInfo.isVisible) {
-            mCompatUI.onCompatInfoChanged(taskInfo, null /* taskListener */);
+            mCompatUI.onCompatInfoChanged(new CompatUIInfo(taskInfo, null /* taskListener */));
             return;
         }
-        mCompatUI.onCompatInfoChanged(taskInfo, taskListener);
+        mCompatUI.onCompatInfoChanged(new CompatUIInfo(taskInfo, taskListener));
     }
 
     private TaskListener getTaskListener(RunningTaskInfo runningTaskInfo) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java
index 2520c25..c02c9cf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java
@@ -20,7 +20,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.CameraCompatTaskInfo.CameraCompatControlState;
 import android.app.TaskInfo;
 import android.content.ComponentName;
 import android.content.Context;
@@ -50,6 +49,10 @@
 import com.android.wm.shell.common.DockStateReader;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
+import com.android.wm.shell.compatui.api.CompatUIEvent;
+import com.android.wm.shell.compatui.api.CompatUIHandler;
+import com.android.wm.shell.compatui.api.CompatUIInfo;
+import com.android.wm.shell.compatui.impl.CompatUIEvents.SizeCompatRestartButtonClicked;
 import com.android.wm.shell.sysui.KeyguardChangeListener;
 import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
@@ -71,17 +74,7 @@
  * activities are in compatibility mode.
  */
 public class CompatUIController implements OnDisplaysChangedListener,
-        DisplayImeController.ImePositionProcessor, KeyguardChangeListener {
-
-    /** Callback for compat UI interaction. */
-    public interface CompatUICallback {
-        /** Called when the size compat restart button appears. */
-        void onSizeCompatRestartButtonAppeared(int taskId);
-        /** Called when the size compat restart button is clicked. */
-        void onSizeCompatRestartButtonClicked(int taskId);
-        /** Called when the camera compat control state is updated. */
-        void onCameraControlStateUpdated(int taskId, @CameraCompatControlState int state);
-    }
+        DisplayImeController.ImePositionProcessor, KeyguardChangeListener, CompatUIHandler {
 
     private static final String TAG = "CompatUIController";
 
@@ -170,7 +163,7 @@
     private final Function<Integer, Integer> mDisappearTimeSupplier;
 
     @Nullable
-    private CompatUICallback mCompatUICallback;
+    private Consumer<CompatUIEvent> mCallback;
 
     // Indicates if the keyguard is currently showing, in which case compat UIs shouldn't
     // be shown.
@@ -230,20 +223,21 @@
         mCompatUIShellCommandHandler.onInit();
     }
 
-    /** Sets the callback for Compat UI interactions. */
-    public void setCompatUICallback(@NonNull CompatUICallback compatUiCallback) {
-        mCompatUICallback = compatUiCallback;
+    /** Sets the callback for UI interactions. */
+    @Override
+    public void setCallback(@Nullable Consumer<CompatUIEvent> callback) {
+        mCallback = callback;
     }
 
     /**
      * Called when the Task info changed. Creates and updates the compat UI if there is an
      * activity in size compat, or removes the UI if there is no size compat activity.
      *
-     * @param taskInfo {@link TaskInfo} task the activity is in.
-     * @param taskListener listener to handle the Task Surface placement.
+     * @param compatUIInfo {@link CompatUIInfo} encapsulates information about the task and listener
      */
-    public void onCompatInfoChanged(@NonNull TaskInfo taskInfo,
-            @Nullable ShellTaskOrganizer.TaskListener taskListener) {
+    public void onCompatInfoChanged(@NonNull CompatUIInfo compatUIInfo) {
+        final TaskInfo taskInfo = compatUIInfo.getTaskInfo();
+        final ShellTaskOrganizer.TaskListener taskListener = compatUIInfo.getListener();
         if (taskInfo != null && !taskInfo.appCompatTaskInfo.topActivityInSizeCompat) {
             mSetOfTaskIdsShowingRestartDialog.remove(taskInfo.taskId);
         }
@@ -466,7 +460,7 @@
     CompatUIWindowManager createCompatUiWindowManager(Context context, TaskInfo taskInfo,
             ShellTaskOrganizer.TaskListener taskListener) {
         return new CompatUIWindowManager(context,
-                taskInfo, mSyncQueue, mCompatUICallback, taskListener,
+                taskInfo, mSyncQueue, mCallback, taskListener,
                 mDisplayController.getDisplayLayout(taskInfo.displayId), mCompatUIHintsState,
                 mCompatUIConfiguration, this::onRestartButtonClicked);
     }
@@ -478,9 +472,9 @@
                 taskInfoState.first)) {
             // We need to show the dialog
             mSetOfTaskIdsShowingRestartDialog.add(taskInfoState.first.taskId);
-            onCompatInfoChanged(taskInfoState.first, taskInfoState.second);
+            onCompatInfoChanged(new CompatUIInfo(taskInfoState.first, taskInfoState.second));
         } else {
-            mCompatUICallback.onSizeCompatRestartButtonClicked(taskInfoState.first.taskId);
+            mCallback.accept(new SizeCompatRestartButtonClicked(taskInfoState.first.taskId));
         }
     }
 
@@ -575,13 +569,13 @@
     private void onRestartDialogCallback(
             Pair<TaskInfo, ShellTaskOrganizer.TaskListener> stateInfo) {
         mTaskIdToRestartDialogWindowManagerMap.remove(stateInfo.first.taskId);
-        mCompatUICallback.onSizeCompatRestartButtonClicked(stateInfo.first.taskId);
+        mCallback.accept(new SizeCompatRestartButtonClicked(stateInfo.first.taskId));
     }
 
     private void onRestartDialogDismissCallback(
             Pair<TaskInfo, ShellTaskOrganizer.TaskListener> stateInfo) {
         mSetOfTaskIdsShowingRestartDialog.remove(stateInfo.first.taskId);
-        onCompatInfoChanged(stateInfo.first, stateInfo.second);
+        onCompatInfoChanged(new CompatUIInfo(stateInfo.first, stateInfo.second));
     }
 
     private void createOrUpdateReachabilityEduLayout(@NonNull TaskInfo taskInfo,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
index 3ab1fad..1931212 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIWindowManager.java
@@ -40,8 +40,10 @@
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
-import com.android.wm.shell.compatui.CompatUIController.CompatUICallback;
 import com.android.wm.shell.compatui.CompatUIController.CompatUIHintsState;
+import com.android.wm.shell.compatui.api.CompatUIEvent;
+import com.android.wm.shell.compatui.impl.CompatUIEvents.CameraControlStateUpdated;
+import com.android.wm.shell.compatui.impl.CompatUIEvents.SizeCompatRestartButtonAppeared;
 
 import java.util.function.Consumer;
 
@@ -50,10 +52,13 @@
  */
 class CompatUIWindowManager extends CompatUIWindowManagerAbstract {
 
-    private final CompatUICallback mCallback;
+    @NonNull
+    private final Consumer<CompatUIEvent> mCallback;
 
+    @NonNull
     private final CompatUIConfiguration mCompatUIConfiguration;
 
+    @NonNull
     private final Consumer<Pair<TaskInfo, ShellTaskOrganizer.TaskListener>> mOnRestartButtonClicked;
 
     // Remember the last reported states in case visibility changes due to keyguard or IME updates.
@@ -65,6 +70,7 @@
     int mCameraCompatControlState = CAMERA_COMPAT_CONTROL_HIDDEN;
 
     @VisibleForTesting
+    @NonNull
     CompatUIHintsState mCompatUIHintsState;
 
     @Nullable
@@ -73,11 +79,15 @@
 
     private final float mHideScmTolerance;
 
-    CompatUIWindowManager(Context context, TaskInfo taskInfo,
-            SyncTransactionQueue syncQueue, CompatUICallback callback,
-            ShellTaskOrganizer.TaskListener taskListener, DisplayLayout displayLayout,
-            CompatUIHintsState compatUIHintsState, CompatUIConfiguration compatUIConfiguration,
-            Consumer<Pair<TaskInfo, ShellTaskOrganizer.TaskListener>> onRestartButtonClicked) {
+    CompatUIWindowManager(@NonNull Context context, @NonNull TaskInfo taskInfo,
+                          @NonNull SyncTransactionQueue syncQueue,
+                          @NonNull Consumer<CompatUIEvent> callback,
+                          @Nullable ShellTaskOrganizer.TaskListener taskListener,
+                          @Nullable DisplayLayout displayLayout,
+                          @NonNull CompatUIHintsState compatUIHintsState,
+                          @NonNull CompatUIConfiguration compatUIConfiguration,
+                          @NonNull Consumer<Pair<TaskInfo, ShellTaskOrganizer.TaskListener>>
+                                  onRestartButtonClicked) {
         super(context, taskInfo, syncQueue, taskListener, displayLayout);
         mCallback = callback;
         mHasSizeCompat = taskInfo.appCompatTaskInfo.topActivityInSizeCompat;
@@ -122,7 +132,7 @@
         updateVisibilityOfViews();
 
         if (mHasSizeCompat) {
-            mCallback.onSizeCompatRestartButtonAppeared(mTaskId);
+            mCallback.accept(new SizeCompatRestartButtonAppeared(mTaskId));
         }
 
         return mLayout;
@@ -177,7 +187,7 @@
                 mCameraCompatControlState == CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED
                         ? CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED
                         : CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
-        mCallback.onCameraControlStateUpdated(mTaskId, mCameraCompatControlState);
+        mCallback.accept(new CameraControlStateUpdated(mTaskId, mCameraCompatControlState));
         mLayout.updateCameraTreatmentButton(mCameraCompatControlState);
     }
 
@@ -188,7 +198,7 @@
             return;
         }
         mCameraCompatControlState = CAMERA_COMPAT_CONTROL_DISMISSED;
-        mCallback.onCameraControlStateUpdated(mTaskId, CAMERA_COMPAT_CONTROL_DISMISSED);
+        mCallback.accept(new CameraControlStateUpdated(mTaskId, CAMERA_COMPAT_CONTROL_DISMISSED));
         mLayout.setCameraControlVisibility(/* show= */ false);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIEvent.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIEvent.kt
new file mode 100644
index 0000000..4a0cf98
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIEvent.kt
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.compatui.api
+
+/**
+ * Abstraction for all the possible Compat UI Component events.
+ */
+interface CompatUIEvent {
+    /**
+     * Unique event identifier
+     */
+    val eventId: Int
+
+    @Suppress("UNCHECKED_CAST")
+    fun <T : CompatUIEvent> asType(): T? = this as? T
+
+    fun <T : CompatUIEvent> asType(clazz: Class<T>): T? {
+        return if (clazz.isInstance(this)) clazz.cast(this) else null
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIHandler.kt
new file mode 100644
index 0000000..817e554
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIHandler.kt
@@ -0,0 +1,34 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.wm.shell.compatui.api
+
+import java.util.function.Consumer
+
+/**
+ * Abstraction for the objects responsible to handle all the CompatUI components and the
+ * communication with the server.
+ */
+interface CompatUIHandler {
+    /**
+     * Invoked when a new model is coming from the server.
+     */
+    fun onCompatInfoChanged(compatUIInfo: CompatUIInfo)
+
+    /**
+     * Optional reference to the object responsible to send {@link CompatUIEvent}
+     */
+    fun setCallback(compatUIEventSender: Consumer<CompatUIEvent>?)
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIInfo.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIInfo.kt
new file mode 100644
index 0000000..dbbf049
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/api/CompatUIInfo.kt
@@ -0,0 +1,25 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.compatui.api
+
+import android.app.TaskInfo
+import com.android.wm.shell.ShellTaskOrganizer
+
+/**
+ * Encapsulate the info of the message from core.
+ */
+data class CompatUIInfo(val taskInfo: TaskInfo, val listener: ShellTaskOrganizer.TaskListener?)
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/CompatUIEvents.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/CompatUIEvents.kt
new file mode 100644
index 0000000..58ce8ed
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/CompatUIEvents.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.compatui.impl
+
+import android.app.AppCompatTaskInfo
+import android.app.CameraCompatTaskInfo
+import com.android.wm.shell.compatui.api.CompatUIEvent
+
+internal const val SIZE_COMPAT_RESTART_BUTTON_APPEARED = 0
+internal const val SIZE_COMPAT_RESTART_BUTTON_CLICKED = 1
+internal const val CAMERA_CONTROL_STATE_UPDATE = 2
+
+/**
+ * All the {@link CompatUIEvent} the Compat UI Framework can handle
+ */
+sealed class CompatUIEvents(override val eventId: Int) : CompatUIEvent {
+    /** Sent when the size compat restart button appears. */
+    data class SizeCompatRestartButtonAppeared(val taskId: Int) :
+            CompatUIEvents(SIZE_COMPAT_RESTART_BUTTON_APPEARED)
+
+    /** Sent when the size compat restart button is clicked. */
+    data class SizeCompatRestartButtonClicked(val taskId: Int) :
+            CompatUIEvents(SIZE_COMPAT_RESTART_BUTTON_CLICKED)
+
+    /** Sent when the camera compat control state is updated. */
+    data class CameraControlStateUpdated(
+            val taskId: Int,
+            @CameraCompatTaskInfo.CameraCompatControlState val state: Int
+    ) : CompatUIEvents(CAMERA_CONTROL_STATE_UPDATE)
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandler.kt
new file mode 100644
index 0000000..a181eaf
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/impl/DefaultCompatUIHandler.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.compatui.impl
+
+import com.android.wm.shell.compatui.api.CompatUIEvent
+import com.android.wm.shell.compatui.api.CompatUIHandler
+import com.android.wm.shell.compatui.api.CompatUIInfo
+import java.util.function.Consumer
+
+/**
+ * Default implementation of {@link CompatUIHandler} to handle CompatUI components
+ */
+class DefaultCompatUIHandler : CompatUIHandler {
+
+    private var compatUIEventSender: Consumer<CompatUIEvent>? = null
+    override fun onCompatInfoChanged(compatUIInfo: CompatUIInfo) {
+        // Empty at the moment
+    }
+
+    override fun setCallback(compatUIEventSender: Consumer<CompatUIEvent>?) {
+        this.compatUIEventSender = compatUIEventSender
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index afaa3c9..9bdc0b2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -71,6 +71,8 @@
 import com.android.wm.shell.compatui.CompatUIConfiguration;
 import com.android.wm.shell.compatui.CompatUIController;
 import com.android.wm.shell.compatui.CompatUIShellCommandHandler;
+import com.android.wm.shell.compatui.api.CompatUIHandler;
+import com.android.wm.shell.compatui.impl.DefaultCompatUIHandler;
 import com.android.wm.shell.desktopmode.DesktopMode;
 import com.android.wm.shell.desktopmode.DesktopModeTaskRepository;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
@@ -211,7 +213,7 @@
             Context context,
             ShellInit shellInit,
             ShellCommandHandler shellCommandHandler,
-            Optional<CompatUIController> compatUI,
+            Optional<CompatUIHandler> compatUI,
             Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasksOptional,
             @ShellMainThread ShellExecutor mainExecutor) {
@@ -230,7 +232,7 @@
 
     @WMSingleton
     @Provides
-    static Optional<CompatUIController> provideCompatUIController(
+    static Optional<CompatUIHandler> provideCompatUIController(
             Context context,
             ShellInit shellInit,
             ShellController shellController,
@@ -247,6 +249,9 @@
         if (!context.getResources().getBoolean(R.bool.config_enableCompatUIController)) {
             return Optional.empty();
         }
+        if (Flags.appCompatUiFramework()) {
+            return Optional.of(new DefaultCompatUIHandler());
+        }
         return Optional.of(
                 new CompatUIController(
                         context,
diff --git a/libs/WindowManager/Shell/tests/OWNERS b/libs/WindowManager/Shell/tests/OWNERS
index b8a19ad..a77fd51 100644
--- a/libs/WindowManager/Shell/tests/OWNERS
+++ b/libs/WindowManager/Shell/tests/OWNERS
@@ -9,7 +9,7 @@
 chenghsiuchang@google.com
 atsjenk@google.com
 jorgegil@google.com
-nmusgrave@google.com
+vaniadesmonda@google.com
 pbdr@google.com
 tkachenkoi@google.com
 mpodolian@google.com
diff --git a/libs/WindowManager/Shell/tests/unittest/Android.bp b/libs/WindowManager/Shell/tests/unittest/Android.bp
index 92be4f9..6b69542 100644
--- a/libs/WindowManager/Shell/tests/unittest/Android.bp
+++ b/libs/WindowManager/Shell/tests/unittest/Android.bp
@@ -59,6 +59,7 @@
         "guava-android-testlib",
         "com.android.window.flags.window-aconfig-java",
         "platform-test-annotations",
+        "flag-junit",
     ],
 
     libs: [
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
index 8303317..e91828b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
@@ -45,6 +45,7 @@
 import static org.mockito.Mockito.verify;
 
 import android.app.ActivityManager.RunningTaskInfo;
+import android.app.TaskInfo;
 import android.content.LocusId;
 import android.content.pm.ParceledListSlice;
 import android.os.Binder;
@@ -63,6 +64,8 @@
 
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.compatui.CompatUIController;
+import com.android.wm.shell.compatui.api.CompatUIInfo;
+import com.android.wm.shell.compatui.impl.CompatUIEvents;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellInit;
@@ -70,6 +73,7 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -371,7 +375,7 @@
         mOrganizer.onTaskAppeared(taskInfo1, /* leash= */ null);
 
         // sizeCompatActivity is null if top activity is not in size compat.
-        verify(mCompatUI).onCompatInfoChanged(taskInfo1, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo1, null /* taskListener */);
 
         // sizeCompatActivity is non-null if top activity is in size compat.
         clearInvocations(mCompatUI);
@@ -381,7 +385,7 @@
         taskInfo2.appCompatTaskInfo.topActivityInSizeCompat = true;
         taskInfo2.isVisible = true;
         mOrganizer.onTaskInfoChanged(taskInfo2);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo2, taskListener);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo2, taskListener);
 
         // Not show size compat UI if task is not visible.
         clearInvocations(mCompatUI);
@@ -391,11 +395,11 @@
         taskInfo3.appCompatTaskInfo.topActivityInSizeCompat = true;
         taskInfo3.isVisible = false;
         mOrganizer.onTaskInfoChanged(taskInfo3);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo3, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo3, null /* taskListener */);
 
         clearInvocations(mCompatUI);
         mOrganizer.onTaskVanished(taskInfo1);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo1, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo1, null /* taskListener */);
     }
 
     @Test
@@ -410,7 +414,7 @@
 
         // Task listener sent to compat UI is null if top activity isn't eligible for letterbox
         // education.
-        verify(mCompatUI).onCompatInfoChanged(taskInfo1, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo1, null /* taskListener */);
 
         // Task listener is non-null if top activity is eligible for letterbox education and task
         // is visible.
@@ -421,7 +425,7 @@
         taskInfo2.appCompatTaskInfo.topActivityEligibleForLetterboxEducation = true;
         taskInfo2.isVisible = true;
         mOrganizer.onTaskInfoChanged(taskInfo2);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo2, taskListener);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo2, taskListener);
 
         // Task listener is null if task is invisible.
         clearInvocations(mCompatUI);
@@ -431,11 +435,11 @@
         taskInfo3.appCompatTaskInfo.topActivityEligibleForLetterboxEducation = true;
         taskInfo3.isVisible = false;
         mOrganizer.onTaskInfoChanged(taskInfo3);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo3, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo3, null /* taskListener */);
 
         clearInvocations(mCompatUI);
         mOrganizer.onTaskVanished(taskInfo1);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo1, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo1, null /* taskListener */);
     }
 
     @Test
@@ -451,7 +455,7 @@
 
         // Task listener sent to compat UI is null if top activity doesn't request a camera
         // compat control.
-        verify(mCompatUI).onCompatInfoChanged(taskInfo1, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo1, null /* taskListener */);
 
         // Task listener is non-null when request a camera compat control for a visible task.
         clearInvocations(mCompatUI);
@@ -462,7 +466,7 @@
                 CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
         taskInfo2.isVisible = true;
         mOrganizer.onTaskInfoChanged(taskInfo2);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo2, taskListener);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo2, taskListener);
 
         // CompatUIController#onCompatInfoChanged is called when requested state for a camera
         // compat control changes for a visible task.
@@ -474,7 +478,7 @@
                 CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED;
         taskInfo3.isVisible = true;
         mOrganizer.onTaskInfoChanged(taskInfo3);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo3, taskListener);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo3, taskListener);
 
         // CompatUIController#onCompatInfoChanged is called when a top activity goes in size compat
         // mode for a visible task that has a compat control.
@@ -487,7 +491,7 @@
                 CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED;
         taskInfo4.isVisible = true;
         mOrganizer.onTaskInfoChanged(taskInfo4);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo4, taskListener);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo4, taskListener);
 
         // Task linster is null when a camera compat control is dimissed for a visible task.
         clearInvocations(mCompatUI);
@@ -498,7 +502,7 @@
                 CAMERA_COMPAT_CONTROL_DISMISSED;
         taskInfo5.isVisible = true;
         mOrganizer.onTaskInfoChanged(taskInfo5);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo5, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo5, null /* taskListener */);
 
         // Task linster is null when request a camera compat control for a invisible task.
         clearInvocations(mCompatUI);
@@ -509,11 +513,11 @@
                 CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
         taskInfo6.isVisible = false;
         mOrganizer.onTaskInfoChanged(taskInfo6);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo6, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo6, null /* taskListener */);
 
         clearInvocations(mCompatUI);
         mOrganizer.onTaskVanished(taskInfo1);
-        verify(mCompatUI).onCompatInfoChanged(taskInfo1, null /* taskListener */);
+        verifyOnCompatInfoChangedInvokedWith(taskInfo1, null /* taskListener */);
     }
 
     @Test
@@ -640,7 +644,8 @@
 
         mOrganizer.onTaskAppeared(task1, /* leash= */ null);
 
-        mOrganizer.onSizeCompatRestartButtonClicked(task1.taskId);
+        mOrganizer.onSizeCompatRestartButtonClicked(
+                new CompatUIEvents.SizeCompatRestartButtonClicked(task1.taskId));
 
         verify(mTaskOrganizerController).restartTaskTopActivityProcessIfVisible(task1.token);
     }
@@ -713,4 +718,13 @@
         taskInfo.isVisible = true;
         return taskInfo;
     }
+
+    private void verifyOnCompatInfoChangedInvokedWith(TaskInfo taskInfo,
+                                                      ShellTaskOrganizer.TaskListener listener) {
+        final ArgumentCaptor<CompatUIInfo> capture = ArgumentCaptor.forClass(CompatUIInfo.class);
+        verify(mCompatUI).onCompatInfoChanged(capture.capture());
+        final CompatUIInfo captureValue = capture.getValue();
+        assertEquals(captureValue.getTaskInfo(), taskInfo);
+        assertEquals(captureValue.getListener(), listener);
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
index 9c00864..fc7a777 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
@@ -38,6 +38,9 @@
 import android.app.TaskInfo;
 import android.content.Context;
 import android.content.res.Configuration;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.view.InsetsSource;
 import android.view.InsetsState;
@@ -45,6 +48,7 @@
 
 import androidx.test.filters.SmallTest;
 
+import com.android.window.flags.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.DisplayController;
@@ -55,6 +59,7 @@
 import com.android.wm.shell.common.DockStateReader;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
+import com.android.wm.shell.compatui.api.CompatUIInfo;
 import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
@@ -63,6 +68,7 @@
 
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -82,6 +88,10 @@
     private static final int DISPLAY_ID = 0;
     private static final int TASK_ID = 12;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     private CompatUIController mController;
     private ShellInit mShellInit;
     @Mock
@@ -168,28 +178,32 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void instantiateController_addInitCallback() {
         verify(mShellInit, times(1)).addInitCallback(any(), any());
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void instantiateController_registerKeyguardChangeListener() {
         verify(mMockShellController, times(1)).addKeyguardChangeListener(any());
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testListenerRegistered() {
         verify(mMockDisplayController).addDisplayWindowListener(mController);
         verify(mMockImeController).addPositionProcessor(mController);
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnCompatInfoChanged() {
         TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true,
                 CAMERA_COMPAT_CONTROL_HIDDEN);
 
         // Verify that the compat controls are added with non-null task listener.
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mController).createCompatUiWindowManager(any(), eq(taskInfo), eq(mMockTaskListener));
         verify(mController).createLetterboxEduWindowManager(any(), eq(taskInfo),
@@ -202,7 +216,7 @@
         clearInvocations(mMockCompatLayout, mMockLetterboxEduLayout, mController);
         taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true,
                 CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockCompatLayout).updateCompatInfo(taskInfo, mMockTaskListener,
                 /* canShow= */ true);
@@ -213,9 +227,9 @@
 
         // Verify that compat controls and letterbox education are removed with null task listener.
         clearInvocations(mMockCompatLayout, mMockLetterboxEduLayout, mController);
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
                 /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN),
-                /* taskListener= */ null);
+                /* taskListener= */ null));
 
         verify(mMockCompatLayout).release();
         verify(mMockLetterboxEduLayout).release();
@@ -223,6 +237,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnCompatInfoChanged_createLayoutReturnsFalse() {
         doReturn(false).when(mMockCompatLayout).createLayout(anyBoolean());
         doReturn(false).when(mMockLetterboxEduLayout).createLayout(anyBoolean());
@@ -230,7 +245,7 @@
 
         TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true,
                 CAMERA_COMPAT_CONTROL_HIDDEN);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mController).createCompatUiWindowManager(any(), eq(taskInfo), eq(mMockTaskListener));
         verify(mController).createLetterboxEduWindowManager(any(), eq(taskInfo),
@@ -240,7 +255,7 @@
 
         // Verify that the layout is created again.
         clearInvocations(mMockCompatLayout, mMockLetterboxEduLayout, mController);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockCompatLayout, never()).updateCompatInfo(any(), any(), anyBoolean());
         verify(mMockLetterboxEduLayout, never()).updateCompatInfo(any(), any(), anyBoolean());
@@ -253,6 +268,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnCompatInfoChanged_updateCompatInfoReturnsFalse() {
         doReturn(false).when(mMockCompatLayout).updateCompatInfo(any(), any(), anyBoolean());
         doReturn(false).when(mMockLetterboxEduLayout).updateCompatInfo(any(), any(), anyBoolean());
@@ -260,7 +276,7 @@
 
         TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true,
                 CAMERA_COMPAT_CONTROL_HIDDEN);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mController).createCompatUiWindowManager(any(), eq(taskInfo), eq(mMockTaskListener));
         verify(mController).createLetterboxEduWindowManager(any(), eq(taskInfo),
@@ -270,7 +286,7 @@
 
         clearInvocations(mMockCompatLayout, mMockLetterboxEduLayout, mMockRestartDialogLayout,
                 mController);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockCompatLayout).updateCompatInfo(taskInfo, mMockTaskListener,
                 /* canShow= */ true);
@@ -282,7 +298,7 @@
         // Verify that the layout is created again.
         clearInvocations(mMockCompatLayout, mMockLetterboxEduLayout, mMockRestartDialogLayout,
                 mController);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockCompatLayout, never()).updateCompatInfo(any(), any(), anyBoolean());
         verify(mMockLetterboxEduLayout, never()).updateCompatInfo(any(), any(), anyBoolean());
@@ -296,6 +312,7 @@
 
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnDisplayAdded() {
         mController.onDisplayAdded(DISPLAY_ID);
         mController.onDisplayAdded(DISPLAY_ID + 1);
@@ -305,11 +322,11 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnDisplayRemoved() {
         mController.onDisplayAdded(DISPLAY_ID);
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
-                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN),
-                mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
+                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener));
 
         mController.onDisplayRemoved(DISPLAY_ID + 1);
 
@@ -328,9 +345,10 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnDisplayConfigurationChanged() {
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
-                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
+                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener));
 
         mController.onDisplayConfigurationChanged(DISPLAY_ID + 1, new Configuration());
 
@@ -346,10 +364,11 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testInsetsChanged() {
         mController.onDisplayAdded(DISPLAY_ID);
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
-                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
+                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener));
         InsetsState insetsState = new InsetsState();
         InsetsSource insetsSource = new InsetsSource(
                 InsetsSource.createId(null, 0, navigationBars()), navigationBars());
@@ -373,9 +392,10 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testChangeLayoutsVisibilityOnImeShowHide() {
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
-                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
+                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener));
 
         // Verify that the restart button is hidden after IME is showing.
         mController.onImeVisibilityChanged(DISPLAY_ID, /* isShowing= */ true);
@@ -387,7 +407,7 @@
         // Verify button remains hidden while IME is showing.
         TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true,
                 CAMERA_COMPAT_CONTROL_HIDDEN);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockCompatLayout).updateCompatInfo(taskInfo, mMockTaskListener,
                 /* canShow= */ false);
@@ -405,9 +425,10 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testChangeLayoutsVisibilityOnKeyguardShowingChanged() {
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
-                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
+                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener));
 
         // Verify that the restart button is hidden after keyguard becomes showing.
         mController.onKeyguardVisibilityChanged(true, false, false);
@@ -419,7 +440,7 @@
         // Verify button remains hidden while keyguard is showing.
         TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true,
                 CAMERA_COMPAT_CONTROL_HIDDEN);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockCompatLayout).updateCompatInfo(taskInfo, mMockTaskListener,
                 /* canShow= */ false);
@@ -437,9 +458,10 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testLayoutsRemainHiddenOnKeyguardShowingFalseWhenImeIsShowing() {
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
-                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
+                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener));
 
         mController.onImeVisibilityChanged(DISPLAY_ID, /* isShowing= */ true);
         mController.onKeyguardVisibilityChanged(true, false, false);
@@ -466,9 +488,10 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testLayoutsRemainHiddenOnImeHideWhenKeyguardIsShowing() {
-        mController.onCompatInfoChanged(createTaskInfo(DISPLAY_ID, TASK_ID,
-                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(createTaskInfo(DISPLAY_ID, TASK_ID,
+                /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN), mMockTaskListener));
 
         mController.onImeVisibilityChanged(DISPLAY_ID, /* isShowing= */ true);
         mController.onKeyguardVisibilityChanged(true, false, false);
@@ -495,6 +518,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testRestartLayoutRecreatedIfNeeded() {
         final TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID,
                 /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN);
@@ -502,14 +526,15 @@
                 .needsToBeRecreated(any(TaskInfo.class),
                         any(ShellTaskOrganizer.TaskListener.class));
 
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockRestartDialogLayout, times(2))
                 .createLayout(anyBoolean());
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testRestartLayoutNotRecreatedIfNotNeeded() {
         final TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID,
                 /* hasSizeCompat= */ true, CAMERA_COMPAT_CONTROL_HIDDEN);
@@ -517,14 +542,15 @@
                 .needsToBeRecreated(any(TaskInfo.class),
                         any(ShellTaskOrganizer.TaskListener.class));
 
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mMockRestartDialogLayout, times(1))
                 .createLayout(anyBoolean());
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateActiveTaskInfo_newTask_visibleAndFocused_updated() {
         // Simulate user aspect ratio button being shown for previous task
         mController.setHasShownUserAspectRatioSettingsButton(true);
@@ -545,6 +571,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateActiveTaskInfo_newTask_notVisibleOrFocused_notUpdated() {
         // Create new task
         final TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID,
@@ -606,6 +633,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateActiveTaskInfo_sameTask_notUpdated() {
         // Create new task
         final TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID,
@@ -634,6 +662,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateActiveTaskInfo_transparentTask_notUpdated() {
         // Create new task
         final TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID,
@@ -674,7 +703,7 @@
                 CAMERA_COMPAT_CONTROL_HIDDEN);
         taskInfo.appCompatTaskInfo.isLetterboxEducationEnabled = false;
 
-        mController.onCompatInfoChanged(taskInfo, mMockTaskListener);
+        mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mController, never()).createLetterboxEduWindowManager(any(), eq(taskInfo),
                 eq(mMockTaskListener));
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java
index cd3e8cb..33d69f5 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUILayoutTest.java
@@ -23,6 +23,7 @@
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.verify;
@@ -31,6 +32,9 @@
 import android.app.CameraCompatTaskInfo.CameraCompatControlState;
 import android.app.TaskInfo;
 import android.graphics.Rect;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 import android.view.LayoutInflater;
@@ -40,16 +44,20 @@
 
 import androidx.test.filters.SmallTest;
 
+import com.android.window.flags.Flags;
 import com.android.wm.shell.R;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.compatui.CompatUIController.CompatUIHintsState;
+import com.android.wm.shell.compatui.api.CompatUIEvent;
+import com.android.wm.shell.compatui.impl.CompatUIEvents;
 
 import junit.framework.Assert;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -70,8 +78,12 @@
 
     private static final int TASK_ID = 1;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Mock private SyncTransactionQueue mSyncTransactionQueue;
-    @Mock private CompatUIController.CompatUICallback mCallback;
+    @Mock private Consumer<CompatUIEvent> mCallback;
     @Mock private Consumer<Pair<TaskInfo, ShellTaskOrganizer.TaskListener>> mOnRestartButtonClicked;
     @Mock private ShellTaskOrganizer.TaskListener mTaskListener;
     @Mock private SurfaceControlViewHost mViewHost;
@@ -101,6 +113,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnClickForRestartButton() {
         final ImageButton button = mLayout.findViewById(R.id.size_compat_restart_button);
         button.performClick();
@@ -117,6 +130,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnLongClickForRestartButton() {
         doNothing().when(mWindowManager).onRestartButtonLongClicked();
 
@@ -127,6 +141,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnClickForSizeCompatHint() {
         mWindowManager.mHasSizeCompat = true;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -137,6 +152,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCameraTreatmentButton_treatmentAppliedByDefault() {
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -145,16 +161,17 @@
         button.performClick();
 
         verify(mWindowManager).onCameraTreatmentButtonClicked();
-        verify(mCallback).onCameraControlStateUpdated(
-                TASK_ID, CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED);
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID,
+                CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED);
 
         button.performClick();
 
-        verify(mCallback).onCameraControlStateUpdated(
-                TASK_ID, CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID,
+                CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCameraTreatmentButton_treatmentSuggestedByDefault() {
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -163,16 +180,17 @@
         button.performClick();
 
         verify(mWindowManager).onCameraTreatmentButtonClicked();
-        verify(mCallback).onCameraControlStateUpdated(
-                TASK_ID, CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID,
+                CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
 
         button.performClick();
 
-        verify(mCallback).onCameraControlStateUpdated(
-                TASK_ID, CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED);
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID,
+                CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED);
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnCameraDismissButtonClicked() {
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -181,12 +199,12 @@
         button.performClick();
 
         verify(mWindowManager).onCameraDismissButtonClicked();
-        verify(mCallback).onCameraControlStateUpdated(
-                TASK_ID, CAMERA_COMPAT_CONTROL_DISMISSED);
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID, CAMERA_COMPAT_CONTROL_DISMISSED);
         verify(mLayout).setCameraControlVisibility(/* show */ false);
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnLongClickForCameraTreatmentButton() {
         doNothing().when(mWindowManager).onCameraButtonLongClicked();
 
@@ -198,6 +216,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnLongClickForCameraDismissButton() {
         doNothing().when(mWindowManager).onCameraButtonLongClicked();
 
@@ -208,6 +227,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnClickForCameraCompatHint() {
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -229,4 +249,15 @@
         taskInfo.configuration.windowConfiguration.setBounds(new Rect(0, 0, 2000, 2000));
         return taskInfo;
     }
+
+    private void verifyOnCameraControlStateUpdatedInvokedWith(int taskId, int state) {
+        final ArgumentCaptor<CompatUIEvent> captureValue = ArgumentCaptor.forClass(
+                CompatUIEvent.class);
+        verify(mCallback).accept(captureValue.capture());
+        final CompatUIEvents.CameraControlStateUpdated compatUIEvent =
+                (CompatUIEvents.CameraControlStateUpdated) captureValue.getValue();
+        Assert.assertEquals((compatUIEvent).getTaskId(), taskId);
+        Assert.assertEquals((compatUIEvent).getState(), state);
+        clearInvocations(mCallback);
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
index 41a81c1..eb3da8f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIWindowManagerTest.java
@@ -25,6 +25,7 @@
 import static android.view.WindowManager.LARGE_SCREEN_SMALLEST_SCREEN_WIDTH_DP;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -42,6 +43,9 @@
 import android.app.TaskInfo;
 import android.content.res.Configuration;
 import android.graphics.Rect;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
@@ -60,6 +64,8 @@
 import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.compatui.CompatUIController.CompatUIHintsState;
+import com.android.wm.shell.compatui.api.CompatUIEvent;
+import com.android.wm.shell.compatui.impl.CompatUIEvents;
 
 import junit.framework.Assert;
 
@@ -85,12 +91,16 @@
     @Rule
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     private static final int TASK_ID = 1;
     private static final int TASK_WIDTH = 2000;
     private static final int TASK_HEIGHT = 2000;
 
     @Mock private SyncTransactionQueue mSyncTransactionQueue;
-    @Mock private CompatUIController.CompatUICallback mCallback;
+    @Mock private Consumer<CompatUIEvent> mCallback;
     @Mock private Consumer<Pair<TaskInfo, ShellTaskOrganizer.TaskListener>> mOnRestartButtonClicked;
     @Mock private ShellTaskOrganizer.TaskListener mTaskListener;
     @Mock private CompatUILayout mLayout;
@@ -129,6 +139,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateSizeCompatButton() {
         // Doesn't create layout if show is false.
         mWindowManager.mHasSizeCompat = true;
@@ -174,6 +185,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateCameraCompatControl() {
         // Doesn't create layout if show is false.
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
@@ -212,6 +224,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testRelease() {
         mWindowManager.mHasSizeCompat = true;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -224,6 +237,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCompatInfo() {
         mWindowManager.mHasSizeCompat = true;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -315,6 +329,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCompatInfoLayoutNotInflatedYet() {
         mWindowManager.createLayout(/* canShow= */ false);
 
@@ -347,6 +362,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateDisplayLayout() {
         final DisplayInfo displayInfo = new DisplayInfo();
         displayInfo.logicalWidth = 1000;
@@ -366,6 +382,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateDisplayLayoutInsets() {
         final DisplayInfo displayInfo = new DisplayInfo();
         displayInfo.logicalWidth = 1000;
@@ -390,6 +407,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateVisibility() {
         // Create button if it is not created.
         mWindowManager.mLayout = null;
@@ -415,6 +433,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testAttachToParentSurface() {
         final SurfaceControl.Builder b = new SurfaceControl.Builder();
         mWindowManager.attachToParentSurface(b);
@@ -423,37 +442,38 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnCameraDismissButtonClicked() {
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
         mWindowManager.createLayout(/* canShow= */ true);
         clearInvocations(mLayout);
         mWindowManager.onCameraDismissButtonClicked();
 
-        verify(mCallback).onCameraControlStateUpdated(TASK_ID, CAMERA_COMPAT_CONTROL_DISMISSED);
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID, CAMERA_COMPAT_CONTROL_DISMISSED);
         verify(mLayout).setCameraControlVisibility(/* show= */ false);
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnCameraTreatmentButtonClicked() {
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
         mWindowManager.createLayout(/* canShow= */ true);
         clearInvocations(mLayout);
         mWindowManager.onCameraTreatmentButtonClicked();
 
-        verify(mCallback).onCameraControlStateUpdated(
-                TASK_ID, CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
-        verify(mLayout).updateCameraTreatmentButton(
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID,
                 CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
+        verify(mLayout).updateCameraTreatmentButton(CAMERA_COMPAT_CONTROL_TREATMENT_APPLIED);
 
         mWindowManager.onCameraTreatmentButtonClicked();
 
-        verify(mCallback).onCameraControlStateUpdated(
-                TASK_ID, CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED);
-        verify(mLayout).updateCameraTreatmentButton(
+        verifyOnCameraControlStateUpdatedInvokedWith(TASK_ID,
                 CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED);
+        verify(mLayout).updateCameraTreatmentButton(CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED);
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnRestartButtonClicked() {
         mWindowManager.onRestartButtonClicked();
 
@@ -468,8 +488,9 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnRestartButtonLongClicked_showHint() {
-       // Not create hint popup.
+        // Not create hint popup.
         mWindowManager.mHasSizeCompat = true;
         mWindowManager.mCompatUIHintsState.mHasShownSizeCompatHint = true;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -483,6 +504,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnCameraControlLongClicked_showHint() {
        // Not create hint popup.
         mWindowManager.mCameraCompatControlState = CAMERA_COMPAT_CONTROL_TREATMENT_SUGGESTED;
@@ -498,6 +520,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testWhenDockedStateHasChanged_needsToBeRecreated() {
         ActivityManager.RunningTaskInfo newTaskInfo = new ActivityManager.RunningTaskInfo();
         newTaskInfo.configuration.uiMode |= Configuration.UI_MODE_TYPE_DESK;
@@ -506,6 +529,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testShouldShowSizeCompatRestartButton() {
         mSetFlagsRule.enableFlags(Flags.FLAG_ALLOW_HIDE_SCM_BUTTON);
         doReturn(85).when(mCompatUIConfiguration).getHideSizeCompatRestartButtonTolerance();
@@ -558,4 +582,15 @@
         taskInfo.configuration.smallestScreenWidthDp = LARGE_SCREEN_SMALLEST_SCREEN_WIDTH_DP;
         return taskInfo;
     }
+
+    private void verifyOnCameraControlStateUpdatedInvokedWith(int taskId, int state) {
+        final ArgumentCaptor<CompatUIEvent> captureValue = ArgumentCaptor.forClass(
+                CompatUIEvent.class);
+        verify(mCallback).accept(captureValue.capture());
+        final CompatUIEvents.CameraControlStateUpdated compatUIEvent =
+                (CompatUIEvents.CameraControlStateUpdated) captureValue.getValue();
+        Assert.assertEquals((compatUIEvent).getTaskId(), taskId);
+        Assert.assertEquals((compatUIEvent).getState(), state);
+        clearInvocations(mCallback);
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java
index 172c263..e8191db 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduDialogLayoutTest.java
@@ -16,12 +16,17 @@
 
 package com.android.wm.shell.compatui;
 
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.view.LayoutInflater;
 import android.view.View;
@@ -32,6 +37,7 @@
 import com.android.wm.shell.ShellTestCase;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -54,6 +60,10 @@
     private View mDismissButton;
     private View mDialogContainer;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -66,6 +76,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnFinishInflate() {
         assertEquals(mLayout.getDialogContainerView(),
                 mLayout.findViewById(R.id.letterbox_education_dialog_container));
@@ -76,6 +87,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnDismissButtonClicked() {
         assertTrue(mDismissButton.performClick());
 
@@ -83,6 +95,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnBackgroundClicked() {
         assertTrue(mLayout.performClick());
 
@@ -90,6 +103,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnDialogContainerClicked() {
         assertTrue(mDialogContainer.performClick());
 
@@ -97,6 +111,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testSetDismissOnClickListenerNull() {
         mLayout.setDismissOnClickListener(null);
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java
index a60a1cb..b5664ac 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/LetterboxEduWindowManagerTest.java
@@ -19,6 +19,7 @@
 import static android.content.res.Configuration.UI_MODE_NIGHT_YES;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -37,6 +38,9 @@
 import android.app.TaskInfo;
 import android.graphics.Insets;
 import android.graphics.Rect;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 import android.view.DisplayCutout;
@@ -61,6 +65,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -116,6 +121,10 @@
     private CompatUIConfiguration mCompatUIConfiguration;
     private TestShellExecutor mExecutor;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -153,6 +162,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_notEligible_doesNotCreateLayout() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ false);
 
@@ -162,6 +172,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_eligibleAndDocked_doesNotCreateLayout() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */
                 true, /* isDocked */ true);
@@ -172,6 +183,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_taskBarEducationIsShowing_doesNotCreateLayout() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true,
                 USER_ID_1, /* isTaskbarEduShowing= */ true);
@@ -182,6 +194,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_canShowFalse_returnsTrueButDoesNotCreateLayout() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
 
@@ -192,6 +205,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_canShowTrue_createsLayoutCorrectly() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
 
@@ -238,6 +252,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_alreadyShownToUser_createsLayoutForOtherUserOnly() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true,
                 USER_ID_1, /* isTaskbarEduShowing= */ false);
@@ -271,6 +286,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_windowManagerReleasedBeforeTransitionsIsIdle_doesNotStartAnim() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
 
@@ -288,6 +304,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCompatInfo_updatesLayoutCorrectly() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
 
@@ -316,6 +333,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCompatInfo_notEligibleUntilUpdate_createsLayoutAfterUpdate() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ false);
 
@@ -329,6 +347,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCompatInfo_canShowFalse_doesNothing() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
 
@@ -343,6 +362,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateDisplayLayout_updatesLayoutCorrectly() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
 
@@ -364,6 +384,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testRelease_animationIsCancelled() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
 
@@ -374,6 +395,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testDeviceThemeChange_educationDialogUnseen_recreated() {
         LetterboxEduWindowManager windowManager = createWindowManager(/* eligible= */ true);
         ActivityManager.RunningTaskInfo newTaskInfo = new ActivityManager.RunningTaskInfo();
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java
index 4f71b83..0da14d6 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduLayoutTest.java
@@ -16,6 +16,8 @@
 
 package com.android.wm.shell.compatui;
 
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
+
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertNotNull;
 
@@ -23,6 +25,9 @@
 import static org.mockito.Mockito.verify;
 
 import android.app.TaskInfo;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.LayoutInflater;
@@ -34,6 +39,7 @@
 import com.android.wm.shell.ShellTestCase;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -62,6 +68,10 @@
     @Mock
     private TaskInfo mTaskInfo;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -74,6 +84,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnFinishInflate() {
         assertNotNull(mMoveUpButton);
         assertNotNull(mMoveDownButton);
@@ -82,6 +93,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void handleVisibility_educationNotEnabled_buttonsAreHidden() {
         mLayout.handleVisibility(/* horizontalEnabled */ false, /* verticalEnabled */
                 false, /* letterboxVerticalPosition */
@@ -94,6 +106,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void handleVisibility_horizontalEducationEnableduiConfigurationIsUpdated() {
         mLayout.handleVisibility(/* horizontalEnabled */ true, /* verticalEnabled */
                 false, /* letterboxVerticalPosition */ -1, /* letterboxHorizontalPosition */
@@ -106,6 +119,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void handleVisibility_verticalEducationEnabled_uiConfigurationIsUpdated() {
         mLayout.handleVisibility(/* horizontalEnabled */ false, /* verticalEnabled */
                 true, /* letterboxVerticalPosition */ 0, /* letterboxHorizontalPosition */
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java
index 5867a85..eafb414 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/ReachabilityEduWindowManagerTest.java
@@ -16,12 +16,17 @@
 
 package com.android.wm.shell.compatui;
 
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
+
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 
 import android.app.ActivityManager;
 import android.app.TaskInfo;
 import android.content.res.Configuration;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 
 import androidx.test.filters.SmallTest;
@@ -35,6 +40,7 @@
 import junit.framework.Assert;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -65,6 +71,10 @@
     private TaskInfo mTaskInfo;
     private ReachabilityEduWindowManager mWindowManager;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -80,6 +90,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateLayout_notEligible_doesNotCreateLayout() {
         assertFalse(mWindowManager.createLayout(/* canShow= */ true));
 
@@ -87,6 +98,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testWhenDockedStateHasChanged_needsToBeRecreated() {
         ActivityManager.RunningTaskInfo newTaskInfo = new ActivityManager.RunningTaskInfo();
         newTaskInfo.configuration.uiMode =
@@ -97,6 +109,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testWhenDarkLightThemeHasChanged_needsToBeRecreated() {
         ActivityManager.RunningTaskInfo newTaskInfo = new ActivityManager.RunningTaskInfo();
         mTaskInfo.configuration.uiMode =
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java
index e2dcdb0..6b0c5dd 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogLayoutTest.java
@@ -16,6 +16,8 @@
 
 package com.android.wm.shell.compatui;
 
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -23,6 +25,9 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.view.LayoutInflater;
 import android.view.View;
@@ -34,6 +39,7 @@
 import com.android.wm.shell.ShellTestCase;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -60,6 +66,10 @@
     private View mDialogContainer;
     private CheckBox mDontRepeatCheckBox;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -76,6 +86,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnFinishInflate() {
         assertEquals(mLayout.getDialogContainerView(),
                 mLayout.findViewById(R.id.letterbox_restart_dialog_container));
@@ -86,6 +97,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnDismissButtonClicked() {
         assertTrue(mDismissButton.performClick());
 
@@ -93,6 +105,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnRestartButtonClickedWithoutCheckbox() {
         mDontRepeatCheckBox.setChecked(false);
         assertTrue(mRestartButton.performClick());
@@ -101,6 +114,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnRestartButtonClickedWithCheckbox() {
         mDontRepeatCheckBox.setChecked(true);
         assertTrue(mRestartButton.performClick());
@@ -109,6 +123,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnBackgroundClickedDoesntDismiss() {
         assertFalse(mLayout.performClick());
 
@@ -116,6 +131,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnDialogContainerClicked() {
         assertTrue(mDialogContainer.performClick());
 
@@ -124,6 +140,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testSetDismissOnClickListenerNull() {
         mLayout.setDismissOnClickListener(null);
 
@@ -135,6 +152,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testSetRestartOnClickListenerNull() {
         mLayout.setRestartOnClickListener(null);
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java
index 9f109a1..cfeef90 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/RestartDialogWindowManagerTest.java
@@ -16,9 +16,14 @@
 
 package com.android.wm.shell.compatui;
 
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
+
 import android.app.ActivityManager;
 import android.app.TaskInfo;
 import android.content.res.Configuration;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 
@@ -33,6 +38,7 @@
 import junit.framework.Assert;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -60,6 +66,10 @@
     private RestartDialogWindowManager mWindowManager;
     private TaskInfo mTaskInfo;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -76,6 +86,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testWhenDockedStateHasChanged_needsToBeRecreated() {
         ActivityManager.RunningTaskInfo newTaskInfo = new ActivityManager.RunningTaskInfo();
         newTaskInfo.configuration.uiMode =
@@ -86,6 +97,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testWhenDarkLightThemeHasChanged_needsToBeRecreated() {
         ActivityManager.RunningTaskInfo newTaskInfo = new ActivityManager.RunningTaskInfo();
         mTaskInfo.configuration.uiMode =
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java
index 0231612..3fa21ce 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsLayoutTest.java
@@ -19,6 +19,7 @@
 import static android.app.CameraCompatTaskInfo.CAMERA_COMPAT_CONTROL_HIDDEN;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
 
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
@@ -28,6 +29,9 @@
 import android.app.CameraCompatTaskInfo.CameraCompatControlState;
 import android.app.TaskInfo;
 import android.content.ComponentName;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.util.Pair;
 import android.view.LayoutInflater;
@@ -47,6 +51,7 @@
 import junit.framework.Assert;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -86,6 +91,10 @@
     private UserAspectRatioSettingsLayout mLayout;
     private TaskInfo mTaskInfo;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -107,6 +116,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnClickForUserAspectRatioSettingsButton() {
         final ImageButton button = mLayout.findViewById(R.id.user_aspect_ratio_settings_button);
         button.performClick();
@@ -123,6 +133,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnLongClickForUserAspectRatioButton() {
         doNothing().when(mWindowManager).onUserAspectRatioSettingsButtonLongClicked();
 
@@ -133,6 +144,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnClickForUserAspectRatioSettingsHint() {
         mWindowManager.mHasUserAspectRatioSettingsButton = true;
         mWindowManager.createLayout(/* canShow= */ true);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java
index 94e168e..9f288cc 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/UserAspectRatioSettingsWindowManagerTest.java
@@ -22,6 +22,7 @@
 import static android.view.WindowInsets.Type.navigationBars;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.window.flags.Flags.FLAG_APP_COMPAT_UI_FRAMEWORK;
 
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -39,6 +40,9 @@
 import android.content.Intent;
 import android.content.res.Configuration;
 import android.graphics.Rect;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper.RunWithLooper;
 import android.util.Pair;
@@ -61,6 +65,7 @@
 import junit.framework.Assert;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -107,6 +112,10 @@
 
     private TestShellExecutor mExecutor;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -138,6 +147,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testCreateUserAspectRatioButton() {
         // Doesn't create layout if show is false.
         mWindowManager.mHasUserAspectRatioSettingsButton = true;
@@ -178,6 +188,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testRelease() {
         mWindowManager.mHasUserAspectRatioSettingsButton = true;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -190,6 +201,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCompatInfo() {
         mWindowManager.mHasUserAspectRatioSettingsButton = true;
         mWindowManager.createLayout(/* canShow= */ true);
@@ -242,6 +254,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateCompatInfoLayoutNotInflatedYet() {
         mWindowManager.mHasUserAspectRatioSettingsButton = true;
         mWindowManager.createLayout(/* canShow= */ false);
@@ -267,6 +280,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testEligibleButtonHiddenIfLetterboxBoundsEqualToStableBounds() {
         TaskInfo taskInfo = createTaskInfo(/* eligibleForUserAspectRatioButton= */
                 true, /* topActivityBoundsLetterboxed */ true, ACTION_MAIN, CATEGORY_LAUNCHER);
@@ -292,6 +306,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUserFullscreenOverrideEnabled_buttonAlwaysShown() {
         TaskInfo taskInfo = createTaskInfo(/* eligibleForUserAspectRatioButton= */
                 true, /* topActivityBoundsLetterboxed */ true, ACTION_MAIN, CATEGORY_LAUNCHER);
@@ -310,6 +325,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateDisplayLayout() {
         final DisplayInfo displayInfo = new DisplayInfo();
         displayInfo.logicalWidth = 1000;
@@ -329,6 +345,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateDisplayLayoutInsets() {
         final DisplayInfo displayInfo = new DisplayInfo();
         displayInfo.logicalWidth = 1000;
@@ -353,6 +370,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testUpdateVisibility() {
         // Create button if it is not created.
         mWindowManager.removeLayout();
@@ -378,6 +396,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testLayoutHasUserAspectRatioSettingsButton() {
         clearInvocations(mWindowManager);
         spyOn(mWindowManager);
@@ -411,6 +430,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testAttachToParentSurface() {
         final SurfaceControl.Builder b = new SurfaceControl.Builder();
         mWindowManager.attachToParentSurface(b);
@@ -419,6 +439,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnUserAspectRatioButtonClicked() {
         mWindowManager.onUserAspectRatioSettingsButtonClicked();
 
@@ -433,6 +454,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testOnUserAspectRatioButtonLongClicked_showHint() {
        // Not create hint popup.
         mWindowManager.mHasUserAspectRatioSettingsButton = true;
@@ -448,6 +470,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled(FLAG_APP_COMPAT_UI_FRAMEWORK)
     public void testWhenDockedStateHasChanged_needsToBeRecreated() {
         ActivityManager.RunningTaskInfo newTaskInfo = new ActivityManager.RunningTaskInfo();
         newTaskInfo.configuration.uiMode |= Configuration.UI_MODE_TYPE_DESK;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlagsTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlagsTest.kt
index 4661042..00f37da 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlagsTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlagsTest.kt
@@ -427,15 +427,10 @@
   }
 
   private fun resetCache() {
-    val cachedToggleOverrideDesktopMode =
-        DESKTOP_WINDOWING_MODE::class.java.getDeclaredField("cachedToggleOverride")
-    cachedToggleOverrideDesktopMode.isAccessible = true
-    cachedToggleOverrideDesktopMode.set(DESKTOP_WINDOWING_MODE, null)
-
-    val cachedToggleOverrideWallpaperActivity =
-      WALLPAPER_ACTIVITY::class.java.getDeclaredField("cachedToggleOverride")
-    cachedToggleOverrideWallpaperActivity.isAccessible = true
-    cachedToggleOverrideWallpaperActivity.set(WALLPAPER_ACTIVITY, null)
+    val cachedToggleOverride =
+      DesktopModeFlags::class.java.getDeclaredField("cachedToggleOverride")
+    cachedToggleOverride.isAccessible = true
+    cachedToggleOverride.set(null, null)
 
     // Clear override cache stored in System property
     System.clearProperty(SYSTEM_PROPERTY_OVERRIDE_KEY)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
index 316cefa1..b1803e9 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
@@ -360,7 +360,7 @@
                 isTopActivityStyleFloating = true
                 numActivities = 1
             }
-            doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+            doReturn(true).`when` { DesktopModeStatus.canEnterDesktopMode(any()) }
             setUpMockDecorationsForTasks(task)
 
             onTaskOpening(task)
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExt.kt b/packages/SettingsLib/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExt.kt
new file mode 100644
index 0000000..02d684d
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExt.kt
@@ -0,0 +1,100 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.media.data.repository
+
+import android.media.AudioManager
+import android.media.IVolumeController
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.buffer
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.launch
+
+/** Returns [AudioManager.setVolumeController] events as a [Flow] */
+fun AudioManager.volumeControllerEvents(): Flow<VolumeControllerEvent> =
+    callbackFlow {
+            volumeController =
+                object : IVolumeController.Stub() {
+                    override fun displaySafeVolumeWarning(flags: Int) {
+                        launch { send(VolumeControllerEvent.DisplaySafeVolumeWarning(flags)) }
+                    }
+
+                    override fun volumeChanged(streamType: Int, flags: Int) {
+                        launch { send(VolumeControllerEvent.VolumeChanged(streamType, flags)) }
+                    }
+
+                    override fun masterMuteChanged(flags: Int) {
+                        launch { send(VolumeControllerEvent.MasterMuteChanged(flags)) }
+                    }
+
+                    override fun setLayoutDirection(layoutDirection: Int) {
+                        launch { send(VolumeControllerEvent.SetLayoutDirection(layoutDirection)) }
+                    }
+
+                    override fun dismiss() {
+                        launch { send(VolumeControllerEvent.Dismiss) }
+                    }
+
+                    override fun setA11yMode(mode: Int) {
+                        launch { send(VolumeControllerEvent.SetA11yMode(mode)) }
+                    }
+
+                    override fun displayCsdWarning(
+                        csdWarning: Int,
+                        displayDurationMs: Int,
+                    ) {
+                        launch {
+                            send(
+                                VolumeControllerEvent.DisplayCsdWarning(
+                                    csdWarning,
+                                    displayDurationMs,
+                                )
+                            )
+                        }
+                    }
+                }
+            awaitClose { volumeController = null }
+        }
+        .buffer()
+
+/** Models events received via [IVolumeController] */
+sealed interface VolumeControllerEvent {
+
+    /** @see [IVolumeController.displaySafeVolumeWarning] */
+    data class DisplaySafeVolumeWarning(val flags: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.volumeChanged] */
+    data class VolumeChanged(val streamType: Int, val flags: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.masterMuteChanged] */
+    data class MasterMuteChanged(val flags: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.setLayoutDirection] */
+    data class SetLayoutDirection(val layoutDirection: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.setA11yMode] */
+    data class SetA11yMode(val mode: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.displayCsdWarning] */
+    data class DisplayCsdWarning(
+        val csdWarning: Int,
+        val displayDurationMs: Int,
+    ) : VolumeControllerEvent
+
+    /** @see [IVolumeController.dismiss] */
+    data object Dismiss : VolumeControllerEvent
+}
\ No newline at end of file
diff --git a/packages/SettingsLib/tests/integ/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExtTest.kt b/packages/SettingsLib/tests/integ/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExtTest.kt
new file mode 100644
index 0000000..83b612d
--- /dev/null
+++ b/packages/SettingsLib/tests/integ/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExtTest.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settingslib.media.data.repository
+
+import android.media.AudioManager
+import android.media.IVolumeController
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class AudioManagerVolumeControllerExtTest {
+
+    private val testScope = TestScope()
+
+    @Captor private lateinit var volumeControllerCaptor: ArgumentCaptor<IVolumeController>
+    @Mock private lateinit var audioManager: AudioManager
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+    }
+
+    @Test
+    fun displaySafeVolumeWarning_emitsEvent() =
+        testEvent(VolumeControllerEvent.DisplaySafeVolumeWarning(1)) { displaySafeVolumeWarning(1) }
+
+    @Test
+    fun volumeChanged_emitsEvent() =
+        testEvent(VolumeControllerEvent.VolumeChanged(1, 2)) { volumeChanged(1, 2) }
+
+    @Test
+    fun masterMuteChanged_emitsEvent() =
+        testEvent(VolumeControllerEvent.MasterMuteChanged(1)) { masterMuteChanged(1) }
+
+    @Test
+    fun setLayoutDirection_emitsEvent() =
+        testEvent(VolumeControllerEvent.SetLayoutDirection(1)) { setLayoutDirection(1) }
+
+    @Test
+    fun setA11yMode_emitsEvent() =
+        testEvent(VolumeControllerEvent.SetA11yMode(1)) { setA11yMode(1) }
+
+    @Test
+    fun displayCsdWarning_emitsEvent() =
+        testEvent(VolumeControllerEvent.DisplayCsdWarning(1, 2)) { displayCsdWarning(1, 2) }
+
+    @Test fun dismiss_emitsEvent() = testEvent(VolumeControllerEvent.Dismiss) { dismiss() }
+
+    private fun testEvent(
+        expectedEvent: VolumeControllerEvent,
+        emit: IVolumeController.() -> Unit,
+    ) =
+        testScope.runTest {
+            var event: VolumeControllerEvent? = null
+            audioManager.volumeControllerEvents().onEach { event = it }.launchIn(backgroundScope)
+            runCurrent()
+            verify(audioManager).volumeController = volumeControllerCaptor.capture()
+
+            volumeControllerCaptor.value.emit()
+            runCurrent()
+
+            assertThat(event).isEqualTo(expectedEvent)
+        }
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
index 97ed74f..b7d6e09 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
@@ -658,14 +658,18 @@
                 )
                 .onSizeChanged { setToolbarSize(it) },
     ) {
+        val addWidgetText = stringResource(R.string.hub_mode_add_widget_button_text)
         ToolbarButton(
             isPrimary = !removeEnabled,
-            modifier = Modifier.align(Alignment.CenterStart),
+            modifier =
+                Modifier.align(Alignment.CenterStart).semantics {
+                    contentDescription = addWidgetText
+                },
             onClick = onOpenWidgetPicker,
         ) {
-            Icon(Icons.Default.Add, stringResource(R.string.hub_mode_add_widget_button_text))
+            Icon(Icons.Default.Add, null)
             Text(
-                text = stringResource(R.string.hub_mode_add_widget_button_text),
+                text = addWidgetText,
             )
         }
 
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index ca5fc12..0637d39 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -341,13 +341,17 @@
     <string name="share_to_app_stop_dialog_button">Stop sharing</string>
 
     <!-- Content description for the status bar chip shown to the user when they're casting their screen to a different device [CHAR LIMIT=NONE] -->
-    <string name="cast_to_other_device_chip_accessibility_label">Casting screen</string>
+    <string name="cast_screen_to_other_device_chip_accessibility_label">Casting screen</string>
     <!-- Title for a dialog shown to the user that will let them stop casting their screen to a different device [CHAR LIMIT=50] -->
-    <string name="cast_to_other_device_stop_dialog_title">Stop casting screen?</string>
+    <string name="cast_screen_to_other_device_stop_dialog_title">Stop casting screen?</string>
+    <!-- Title for a dialog shown to the user that will let them stop casting to a different device [CHAR LIMIT=50] -->
+    <string name="cast_to_other_device_stop_dialog_title">Stop casting?</string>
     <!-- Text telling a user that they will stop casting their screen to a different device if they click the "Stop casting" button [CHAR LIMIT=100] -->
-    <string name="cast_to_other_device_stop_dialog_message">You will stop casting your screen</string>
+    <string name="cast_screen_to_other_device_stop_dialog_message">You will stop casting your screen</string>
     <!-- Text telling a user that they will stop casting the contents of the specified [app_name] to a different device if they click the "Stop casting" button. Note that the app name will appear in bold.  [CHAR LIMIT=100] -->
-    <string name="cast_to_other_device_stop_dialog_message_specific_app">You will stop casting &lt;b><xliff:g id="app_name" example="Photos App">%1$s</xliff:g>&lt;/b></string>
+    <string name="cast_screen_to_other_device_stop_dialog_message_specific_app">You will stop casting &lt;b><xliff:g id="app_name" example="Photos App">%1$s</xliff:g>&lt;/b></string>
+    <!-- Text telling a user that they're currently casting to a different device [CHAR LIMIT=100] -->
+    <string name="cast_to_other_device_stop_dialog_message">You\'re currently casting</string>
     <!-- Button to stop screen casting to a different device [CHAR LIMIT=35] -->
     <string name="cast_to_other_device_stop_dialog_button">Stop casting</string>
 
diff --git a/packages/SystemUI/src/com/android/systemui/assist/AssistManager.java b/packages/SystemUI/src/com/android/systemui/assist/AssistManager.java
index d30f33f..a67dcdb 100644
--- a/packages/SystemUI/src/com/android/systemui/assist/AssistManager.java
+++ b/packages/SystemUI/src/com/android/systemui/assist/AssistManager.java
@@ -515,7 +515,7 @@
     }
 
     @Nullable
-    private ComponentName getAssistInfo() {
+    public ComponentName getAssistInfo() {
         return getAssistInfoForUser(mSelectedUserInteractor.getSelectedUserId());
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index 4b2fb44..572283a 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -74,6 +74,7 @@
 import com.android.systemui.mediaprojection.MediaProjectionModule;
 import com.android.systemui.mediaprojection.appselector.MediaProjectionActivitiesModule;
 import com.android.systemui.mediaprojection.taskswitcher.MediaProjectionTaskSwitcherModule;
+import com.android.systemui.mediarouter.MediaRouterModule;
 import com.android.systemui.model.SceneContainerPlugin;
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.motiontool.MotionToolModule;
@@ -220,6 +221,7 @@
         MediaProjectionActivitiesModule.class,
         MediaProjectionModule.class,
         MediaProjectionTaskSwitcherModule.class,
+        MediaRouterModule.class,
         MotionToolModule.class,
         NotificationIconAreaControllerModule.class,
         PeopleHubModule.class,
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
index 5b9d30e..b703892 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
@@ -227,7 +227,7 @@
     get() =
         when (type) {
             ShortcutCategoryType.SYSTEM -> R.string.shortcut_helper_category_system
-            ShortcutCategoryType.MULTI_TASKING -> R.string.shortcut_helper_category_system
+            ShortcutCategoryType.MULTI_TASKING -> R.string.shortcut_helper_category_multitasking
             ShortcutCategoryType.IME -> R.string.shortcut_helper_category_input
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
index 646db40..d0e3ab4 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/view/ShortcutHelperActivity.kt
@@ -34,7 +34,6 @@
 import androidx.lifecycle.lifecycleScope
 import com.android.compose.theme.PlatformTheme
 import com.android.systemui.keyboard.shortcut.ui.composable.ShortcutHelper
-import com.android.systemui.keyboard.shortcut.ui.model.ShortcutsUiState
 import com.android.systemui.keyboard.shortcut.ui.viewmodel.ShortcutHelperViewModel
 import com.android.systemui.res.R
 import com.google.android.material.bottomsheet.BottomSheetBehavior
@@ -81,10 +80,7 @@
         requireViewById<ComposeView>(R.id.shortcut_helper_compose_container).apply {
             setContent {
                 PlatformTheme {
-                    val shortcutsUiState by
-                        viewModel.shortcutsUiState.collectAsStateWithLifecycle(
-                            initialValue = ShortcutsUiState.Inactive
-                        )
+                    val shortcutsUiState by viewModel.shortcutsUiState.collectAsStateWithLifecycle()
                     ShortcutHelper(
                         shortcutsUiState = shortcutsUiState,
                         onKeyboardSettingsClicked = ::onKeyboardSettingsClicked,
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
index 3759b0c..e602cad 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/viewmodel/ShortcutHelperViewModel.kt
@@ -23,13 +23,17 @@
 import com.android.systemui.keyboard.shortcut.ui.model.ShortcutsUiState
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 
 class ShortcutHelperViewModel
 @Inject
 constructor(
+    @Background private val backgroundScope: CoroutineScope,
     @Background private val backgroundDispatcher: CoroutineDispatcher,
     private val stateInteractor: ShortcutHelperStateInteractor,
     categoriesInteractor: ShortcutHelperCategoriesInteractor,
@@ -42,16 +46,22 @@
             .flowOn(backgroundDispatcher)
 
     val shortcutsUiState =
-        categoriesInteractor.shortcutCategories.map {
-            if (it.isEmpty()) {
-                ShortcutsUiState.Inactive
-            } else {
-                ShortcutsUiState.Active(
-                    shortcutCategories = it,
-                    defaultSelectedCategory = it.first().type,
-                )
+        categoriesInteractor.shortcutCategories
+            .map {
+                if (it.isEmpty()) {
+                    ShortcutsUiState.Inactive
+                } else {
+                    ShortcutsUiState.Active(
+                        shortcutCategories = it,
+                        defaultSelectedCategory = it.first().type,
+                    )
+                }
             }
-        }
+            .stateIn(
+                scope = backgroundScope,
+                started = SharingStarted.Lazily,
+                initialValue = ShortcutsUiState.Inactive
+            )
 
     fun onViewClosed() {
         stateInteractor.onViewClosed()
diff --git a/packages/SystemUI/src/com/android/systemui/mediarouter/MediaRouterModule.kt b/packages/SystemUI/src/com/android/systemui/mediarouter/MediaRouterModule.kt
new file mode 100644
index 0000000..c07e3a0
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/mediarouter/MediaRouterModule.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.mediarouter
+
+import com.android.systemui.mediarouter.data.repository.MediaRouterRepository
+import com.android.systemui.mediarouter.data.repository.MediaRouterRepositoryImpl
+import dagger.Binds
+import dagger.Module
+
+@Module
+interface MediaRouterModule {
+    @Binds fun mediaRouterRepository(impl: MediaRouterRepositoryImpl): MediaRouterRepository
+}
diff --git a/packages/SystemUI/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepository.kt b/packages/SystemUI/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepository.kt
new file mode 100644
index 0000000..998d76c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepository.kt
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.mediarouter.data.repository
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.statusbar.policy.CastController
+import com.android.systemui.statusbar.policy.CastDevice
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.stateIn
+
+/** A repository for data coming from MediaRouter APIs. */
+interface MediaRouterRepository {
+    /** A list of the cast devices that MediaRouter is currently aware of. */
+    val castDevices: StateFlow<List<CastDevice>>
+
+    /** Stops the cast to the given device. */
+    fun stopCasting(device: CastDevice)
+}
+
+@SysUISingleton
+class MediaRouterRepositoryImpl
+@Inject
+constructor(
+    @Application private val scope: CoroutineScope,
+    private val castController: CastController,
+) : MediaRouterRepository {
+    override val castDevices: StateFlow<List<CastDevice>> =
+        conflatedCallbackFlow {
+                val callback =
+                    CastController.Callback {
+                        val mediaRouterCastDevices =
+                            castController.castDevices.filter {
+                                it.origin == CastDevice.CastOrigin.MediaRouter
+                            }
+                        trySend(mediaRouterCastDevices)
+                    }
+                castController.addCallback(callback)
+                awaitClose { castController.removeCallback(callback) }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), emptyList())
+
+    override fun stopCasting(device: CastDevice) {
+        castController.stopCasting(device)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/interactor/MediaRouterChipInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/interactor/MediaRouterChipInteractor.kt
new file mode 100644
index 0000000..21f301c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/interactor/MediaRouterChipInteractor.kt
@@ -0,0 +1,63 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.chips.casttootherdevice.domain.interactor
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.mediarouter.data.repository.MediaRouterRepository
+import com.android.systemui.statusbar.chips.casttootherdevice.domain.model.MediaRouterCastModel
+import com.android.systemui.statusbar.policy.CastDevice
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+
+/**
+ * Interactor for MediaRouter events, used to show the cast-audio-to-other-device chip in the status
+ * bar.
+ */
+@SysUISingleton
+class MediaRouterChipInteractor
+@Inject
+constructor(
+    @Application private val scope: CoroutineScope,
+    private val mediaRouterRepository: MediaRouterRepository,
+) {
+    private val activeCastDevice: StateFlow<CastDevice?> =
+        mediaRouterRepository.castDevices
+            .map { allDevices -> allDevices.firstOrNull { it.isCasting } }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), null)
+
+    /** The current casting state, according to MediaRouter APIs. */
+    val mediaRouterCastingState: StateFlow<MediaRouterCastModel> =
+        activeCastDevice
+            .map {
+                if (it != null) {
+                    MediaRouterCastModel.Casting
+                } else {
+                    MediaRouterCastModel.DoingNothing
+                }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), MediaRouterCastModel.DoingNothing)
+
+    /** Stops the currently active MediaRouter cast. */
+    fun stopCasting() {
+        activeCastDevice.value?.let { mediaRouterRepository.stopCasting(it) }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/model/MediaRouterCastModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/model/MediaRouterCastModel.kt
new file mode 100644
index 0000000..b228922
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/model/MediaRouterCastModel.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.chips.casttootherdevice.domain.model
+
+/** Represents the current casting state, according to MediaRouter APIs. */
+sealed interface MediaRouterCastModel {
+    /** MediaRouter isn't aware of any active cast. */
+    data object DoingNothing : MediaRouterCastModel
+
+    /** MediaRouter has an active cast. */
+    data object Casting : MediaRouterCastModel
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegate.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastScreenToOtherDeviceDialogDelegate.kt
similarity index 85%
rename from packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegate.kt
rename to packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastScreenToOtherDeviceDialogDelegate.kt
index bc0f492..ffb20a7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegate.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastScreenToOtherDeviceDialogDelegate.kt
@@ -24,7 +24,7 @@
 import com.android.systemui.statusbar.phone.SystemUIDialog
 
 /** A dialog that lets the user stop an ongoing cast-screen-to-other-device event. */
-class EndCastToOtherDeviceDialogDelegate(
+class EndCastScreenToOtherDeviceDialogDelegate(
     private val endMediaProjectionDialogHelper: EndMediaProjectionDialogHelper,
     private val stopAction: () -> Unit,
     private val state: ProjectionChipModel.Projecting,
@@ -36,13 +36,14 @@
     override fun beforeCreate(dialog: SystemUIDialog, savedInstanceState: Bundle?) {
         with(dialog) {
             setIcon(CAST_TO_OTHER_DEVICE_ICON)
-            setTitle(R.string.cast_to_other_device_stop_dialog_title)
+            setTitle(R.string.cast_screen_to_other_device_stop_dialog_title)
+            // TODO(b/332662551): Include device name in this string.
             setMessage(
                 endMediaProjectionDialogHelper.getDialogMessage(
                     state.projectionState,
-                    genericMessageResId = R.string.cast_to_other_device_stop_dialog_message,
+                    genericMessageResId = R.string.cast_screen_to_other_device_stop_dialog_message,
                     specificAppMessageResId =
-                        R.string.cast_to_other_device_stop_dialog_message_specific_app,
+                        R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
                 )
             )
             // No custom on-click, because the dialog will automatically be dismissed when the
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegate.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndGenericCastToOtherDeviceDialogDelegate.kt
similarity index 74%
copy from packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegate.kt
copy to packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndGenericCastToOtherDeviceDialogDelegate.kt
index bc0f492..afe67b4 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegate.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndGenericCastToOtherDeviceDialogDelegate.kt
@@ -19,15 +19,17 @@
 import android.os.Bundle
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.chips.casttootherdevice.ui.viewmodel.CastToOtherDeviceChipViewModel.Companion.CAST_TO_OTHER_DEVICE_ICON
-import com.android.systemui.statusbar.chips.mediaprojection.domain.model.ProjectionChipModel
 import com.android.systemui.statusbar.chips.mediaprojection.ui.view.EndMediaProjectionDialogHelper
 import com.android.systemui.statusbar.phone.SystemUIDialog
 
-/** A dialog that lets the user stop an ongoing cast-screen-to-other-device event. */
-class EndCastToOtherDeviceDialogDelegate(
+/**
+ * A dialog that lets the user stop an ongoing cast-to-other-device event. The user could be casting
+ * their screen, or just casting their audio. This dialog uses generic strings to handle both cases
+ * well.
+ */
+class EndGenericCastToOtherDeviceDialogDelegate(
     private val endMediaProjectionDialogHelper: EndMediaProjectionDialogHelper,
     private val stopAction: () -> Unit,
-    private val state: ProjectionChipModel.Projecting,
 ) : SystemUIDialog.Delegate {
     override fun createDialog(): SystemUIDialog {
         return endMediaProjectionDialogHelper.createDialog(this)
@@ -37,14 +39,8 @@
         with(dialog) {
             setIcon(CAST_TO_OTHER_DEVICE_ICON)
             setTitle(R.string.cast_to_other_device_stop_dialog_title)
-            setMessage(
-                endMediaProjectionDialogHelper.getDialogMessage(
-                    state.projectionState,
-                    genericMessageResId = R.string.cast_to_other_device_stop_dialog_message,
-                    specificAppMessageResId =
-                        R.string.cast_to_other_device_stop_dialog_message_specific_app,
-                )
-            )
+            // TODO(b/332662551): Include device name in this string.
+            setMessage(R.string.cast_to_other_device_stop_dialog_message)
             // No custom on-click, because the dialog will automatically be dismissed when the
             // button is clicked anyway.
             setNegativeButton(R.string.close_dialog_button, /* onClick= */ null)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
index 73ccaab..2eff336 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
@@ -23,7 +23,10 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.res.R
-import com.android.systemui.statusbar.chips.casttootherdevice.ui.view.EndCastToOtherDeviceDialogDelegate
+import com.android.systemui.statusbar.chips.casttootherdevice.domain.interactor.MediaRouterChipInteractor
+import com.android.systemui.statusbar.chips.casttootherdevice.domain.model.MediaRouterCastModel
+import com.android.systemui.statusbar.chips.casttootherdevice.ui.view.EndCastScreenToOtherDeviceDialogDelegate
+import com.android.systemui.statusbar.chips.casttootherdevice.ui.view.EndGenericCastToOtherDeviceDialogDelegate
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.MediaProjectionChipInteractor
 import com.android.systemui.statusbar.chips.mediaprojection.domain.model.ProjectionChipModel
 import com.android.systemui.statusbar.chips.mediaprojection.ui.view.EndMediaProjectionDialogHelper
@@ -36,12 +39,14 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 
 /**
- * View model for the cast-to-other-device chip, shown when sharing your phone screen content to a
- * different device. (Triggered from the Quick Settings Cast tile or from the Settings app.)
+ * View model for the cast-to-other-device chip, shown when a user is sharing content to a different
+ * device. (Triggered from the Quick Settings Cast tile or from the Settings app.) The content could
+ * either be the user's screen, or just the user's audio.
  */
 @SysUISingleton
 class CastToOtherDeviceChipViewModel
@@ -49,14 +54,19 @@
 constructor(
     @Application private val scope: CoroutineScope,
     private val mediaProjectionChipInteractor: MediaProjectionChipInteractor,
+    private val mediaRouterChipInteractor: MediaRouterChipInteractor,
     private val systemClock: SystemClock,
     private val dialogTransitionAnimator: DialogTransitionAnimator,
     private val endMediaProjectionDialogHelper: EndMediaProjectionDialogHelper,
 ) : OngoingActivityChipViewModel {
-    override val chip: StateFlow<OngoingActivityChipModel> =
-        // TODO(b/342169876): The MediaProjection APIs are not invoked for certain
-        // cast-to-other-device events, like audio-only casting. We should also listen to
-        // MediaRouter APIs to cover all cast events.
+    /**
+     * The cast chip to show, based only on MediaProjection API events.
+     *
+     * This chip will only be [OngoingActivityChipModel.Shown] when the user is casting their
+     * *screen*. If the user is only casting audio, this chip will be
+     * [OngoingActivityChipModel.Hidden].
+     */
+    private val projectionChip: StateFlow<OngoingActivityChipModel> =
         mediaProjectionChipInteractor.projection
             .map { projectionModel ->
                 when (projectionModel) {
@@ -65,7 +75,7 @@
                         if (projectionModel.type != ProjectionChipModel.Type.CAST_TO_OTHER_DEVICE) {
                             OngoingActivityChipModel.Hidden
                         } else {
-                            createCastToOtherDeviceChip(projectionModel)
+                            createCastScreenToOtherDeviceChip(projectionModel)
                         }
                     }
                 }
@@ -73,43 +83,132 @@
             // See b/347726238.
             .stateIn(scope, SharingStarted.Lazily, OngoingActivityChipModel.Hidden)
 
+    /**
+     * The cast chip to show, based only on MediaRouter API events.
+     *
+     * This chip will be [OngoingActivityChipModel.Shown] when the user is casting their screen *or*
+     * their audio.
+     *
+     * The MediaProjection APIs are not invoked for casting *only audio* to another device because
+     * MediaProjection is only concerned with *screen* sharing (see b/342169876). We listen to
+     * MediaRouter APIs here to cover audio-only casting.
+     *
+     * Note that this means we will start showing the cast chip before the casting actually starts,
+     * for **both** audio-only casting and screen casting. MediaRouter is aware of all
+     * cast-to-other-device events, and MediaRouter immediately marks a device as "connecting" once
+     * a user selects what device they'd like to cast to, even if they haven't hit "Start casting"
+     * yet. All of SysUI considers "connecting" devices to be casting (see
+     * [com.android.systemui.statusbar.policy.CastDevice.isCasting]), so the chip will follow the
+     * same convention and start showing once a device is selected. See b/269975671.
+     */
+    private val routerChip =
+        mediaRouterChipInteractor.mediaRouterCastingState
+            .map { routerModel ->
+                when (routerModel) {
+                    is MediaRouterCastModel.DoingNothing -> OngoingActivityChipModel.Hidden
+                    is MediaRouterCastModel.Casting -> {
+                        // A consequence of b/269975671 is that MediaRouter will mark a device as
+                        // casting before casting has actually started. To alleviate this bug a bit,
+                        // we won't show a timer for MediaRouter events. That way, we won't show a
+                        // timer if cast hasn't actually started.
+                        //
+                        // This does mean that the audio-only casting chip will *never* show a
+                        // timer, because audio-only casting never activates the MediaProjection
+                        // APIs and those are the only cast APIs that show a timer.
+                        createIconOnlyCastChip()
+                    }
+                }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), OngoingActivityChipModel.Hidden)
+
+    override val chip: StateFlow<OngoingActivityChipModel> =
+        combine(projectionChip, routerChip) { projection, router ->
+                // A consequence of b/269975671 is that MediaRouter and MediaProjection APIs fire at
+                // different times when *screen* casting:
+                //
+                // 1. When the user chooses what device to cast to, the MediaRouter APIs mark the
+                // device as casting (even though casting hasn't actually started yet). At this
+                // point, `routerChip` is [OngoingActivityChipModel.Shown] but `projectionChip` is
+                // [OngoingActivityChipModel.Hidden], and we'll show the router chip.
+                //
+                // 2. Once casting has actually started, the MediaProjection APIs become aware of
+                // the device. At this point, both `routerChip` and `projectionChip` are
+                // [OngoingActivityChipModel.Shown].
+                //
+                // Because the MediaProjection APIs have activated, we know that the user is screen
+                // casting (not audio casting). We need to switch to using `projectionChip` because
+                // that chip will show information specific to screen casting. The `projectionChip`
+                // will also show a timer, as opposed to `routerChip`'s icon-only display.
+                if (projection is OngoingActivityChipModel.Shown) {
+                    projection
+                } else {
+                    router
+                }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), OngoingActivityChipModel.Hidden)
+
     /** Stops the currently active projection. */
     private fun stopProjecting() {
         mediaProjectionChipInteractor.stopProjecting()
     }
 
-    private fun createCastToOtherDeviceChip(
+    private fun stopMediaRouterCasting() {
+        mediaRouterChipInteractor.stopCasting()
+    }
+
+    private fun createCastScreenToOtherDeviceChip(
         state: ProjectionChipModel.Projecting,
     ): OngoingActivityChipModel.Shown {
         return OngoingActivityChipModel.Shown.Timer(
             icon =
                 Icon.Resource(
                     CAST_TO_OTHER_DEVICE_ICON,
-                    // Note: This string is "Casting screen", which is okay right now because this
-                    // chip does not currently support audio-only casting. If the chip starts
-                    // supporting audio-only casting (see b/342169876), update the content
-                    // description to just "Casting".
+                    // This string is "Casting screen"
                     ContentDescription.Resource(
-                        R.string.cast_to_other_device_chip_accessibility_label,
+                        R.string.cast_screen_to_other_device_chip_accessibility_label,
                     ),
                 ),
             colors = ColorsModel.Red,
             // TODO(b/332662551): Maybe use a MediaProjection API to fetch this time.
             startTimeMs = systemClock.elapsedRealtime(),
             createDialogLaunchOnClickListener(
-                createCastToOtherDeviceDialogDelegate(state),
+                createCastScreenToOtherDeviceDialogDelegate(state),
                 dialogTransitionAnimator,
             ),
         )
     }
 
-    private fun createCastToOtherDeviceDialogDelegate(state: ProjectionChipModel.Projecting) =
-        EndCastToOtherDeviceDialogDelegate(
+    private fun createIconOnlyCastChip(): OngoingActivityChipModel.Shown {
+        return OngoingActivityChipModel.Shown.IconOnly(
+            icon =
+                Icon.Resource(
+                    CAST_TO_OTHER_DEVICE_ICON,
+                    // This string is just "Casting"
+                    ContentDescription.Resource(R.string.accessibility_casting),
+                ),
+            colors = ColorsModel.Red,
+            createDialogLaunchOnClickListener(
+                createGenericCastToOtherDeviceDialogDelegate(),
+                dialogTransitionAnimator,
+            ),
+        )
+    }
+
+    private fun createCastScreenToOtherDeviceDialogDelegate(
+        state: ProjectionChipModel.Projecting,
+    ) =
+        EndCastScreenToOtherDeviceDialogDelegate(
             endMediaProjectionDialogHelper,
             stopAction = this::stopProjecting,
             state,
         )
 
+    private fun createGenericCastToOtherDeviceDialogDelegate() =
+        EndGenericCastToOtherDeviceDialogDelegate(
+            endMediaProjectionDialogHelper,
+            stopAction = this::stopMediaRouterCasting,
+        )
+
     companion object {
         @DrawableRes val CAST_TO_OTHER_DEVICE_ICON = R.drawable.ic_cast_connected
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinator.kt
index 5a1146d..1434dc0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinator.kt
@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-@file:OptIn(ExperimentalCoroutinesApi::class)
-
 package com.android.systemui.statusbar.notification.collection.coordinator
 
 import android.app.NotificationManager
@@ -27,9 +25,10 @@
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.keyguard.data.repository.KeyguardRepository
-import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.plugins.statusbar.StatusBarStateController
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.StatusBarState
 import com.android.systemui.statusbar.expansionChanges
 import com.android.systemui.statusbar.notification.collection.GroupEntry
@@ -58,7 +57,6 @@
 import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.delay
@@ -89,7 +87,7 @@
     private val headsUpManager: HeadsUpManager,
     private val keyguardNotificationVisibilityProvider: KeyguardNotificationVisibilityProvider,
     private val keyguardRepository: KeyguardRepository,
-    private val keyguardTransitionRepository: KeyguardTransitionRepository,
+    private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
     private val logger: KeyguardCoordinatorLogger,
     @Application private val scope: CoroutineScope,
     private val sectionHeaderVisibilityProvider: SectionHeaderVisibilityProvider,
@@ -126,8 +124,9 @@
     private suspend fun trackSeenNotifications() {
         // Whether or not keyguard is visible (or occluded).
         val isKeyguardPresent: Flow<Boolean> =
-            keyguardTransitionRepository.transitions
-                .map { step -> step.to != KeyguardState.GONE }
+            keyguardTransitionInteractor
+                .transitionValue(Scenes.Gone, stateWithoutSceneContainer = KeyguardState.GONE)
+                .map { it == 0f }
                 .distinctUntilChanged()
                 .onEach { trackingUnseen -> logger.logTrackingUnseen(trackingUnseen) }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java
index 6012ecd..775f34d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/BatteryControllerImpl.java
@@ -250,6 +250,7 @@
             mLevel = (int) (100f
                     * intent.getIntExtra(BatteryManager.EXTRA_LEVEL, 0)
                     / intent.getIntExtra(BatteryManager.EXTRA_SCALE, 100));
+            int previousPluggedChargingSource = mPluggedChargingSource;
             mPluggedChargingSource = intent.getIntExtra(BatteryManager.EXTRA_PLUGGED, 0);
             mPluggedIn = mPluggedChargingSource != 0;
             final int status = intent.getIntExtra(BatteryManager.EXTRA_STATUS,
@@ -276,7 +277,9 @@
                 mIsBatteryDefender = isBatteryDefender;
                 fireIsBatteryDefenderChanged();
             }
-
+            if (mPluggedChargingSource != previousPluggedChargingSource) {
+                updatePowerSave();
+            }
             fireBatteryLevelChanged();
         } else if (action.equals(PowerManager.ACTION_POWER_SAVE_MODE_CHANGED)) {
             updatePowerSave();
diff --git a/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java b/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java
index 9c8a481..501fee6 100644
--- a/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java
+++ b/packages/SystemUI/src/com/android/systemui/util/dagger/UtilModule.java
@@ -20,15 +20,15 @@
 import com.android.systemui.util.RingerModeTrackerImpl;
 import com.android.systemui.util.animation.data.repository.AnimationStatusRepository;
 import com.android.systemui.util.animation.data.repository.AnimationStatusRepositoryImpl;
+import com.android.systemui.util.icons.AppCategoryIconProvider;
+import com.android.systemui.util.icons.AppCategoryIconProviderImpl;
 import com.android.systemui.util.wrapper.UtilWrapperModule;
 
 import dagger.Binds;
 import dagger.Module;
 
 /** Dagger Module for code in the util package. */
-@Module(includes = {
-                UtilWrapperModule.class
-        })
+@Module(includes = {UtilWrapperModule.class})
 public interface UtilModule {
     /** */
     @Binds
@@ -37,4 +37,8 @@
     @Binds
     AnimationStatusRepository provideAnimationStatus(
             AnimationStatusRepositoryImpl ringerModeTrackerImpl);
+
+    /** */
+    @Binds
+    AppCategoryIconProvider appCategoryIconProvider(AppCategoryIconProviderImpl impl);
 }
diff --git a/packages/SystemUI/src/com/android/systemui/util/icons/AppCategoryIconProvider.kt b/packages/SystemUI/src/com/android/systemui/util/icons/AppCategoryIconProvider.kt
new file mode 100644
index 0000000..6e3f8f1
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/util/icons/AppCategoryIconProvider.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.util.icons
+
+import android.content.Intent
+import android.content.pm.PackageInfo
+import android.content.pm.PackageManager
+import android.graphics.drawable.Icon
+import android.os.RemoteException
+import android.util.Log
+import com.android.systemui.assist.AssistManager
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.shared.system.PackageManagerWrapper
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+interface AppCategoryIconProvider {
+    /** Returns the [Icon] of the default assistant app, if it exists. */
+    suspend fun assistantAppIcon(): Icon?
+
+    /**
+     * Returns the [Icon] of the default app of [category], if it exists. Category can be for
+     * example [Intent.CATEGORY_APP_EMAIL] or [Intent.CATEGORY_APP_CALCULATOR].
+     */
+    suspend fun categoryAppIcon(category: String): Icon?
+}
+
+class AppCategoryIconProviderImpl
+@Inject
+constructor(
+    @Background private val backgroundDispatcher: CoroutineDispatcher,
+    private val assistManager: AssistManager,
+    private val packageManager: PackageManager,
+    private val packageManagerWrapper: PackageManagerWrapper,
+) : AppCategoryIconProvider {
+
+    override suspend fun assistantAppIcon(): Icon? {
+        val assistInfo = assistManager.assistInfo ?: return null
+        return getPackageIcon(assistInfo.packageName)
+    }
+
+    override suspend fun categoryAppIcon(category: String): Icon? {
+        val intent = Intent(Intent.ACTION_MAIN).apply { addCategory(category) }
+        val packageInfo = getPackageInfo(intent) ?: return null
+        return getPackageIcon(packageInfo.packageName)
+    }
+
+    private suspend fun getPackageInfo(intent: Intent): PackageInfo? =
+        withContext(backgroundDispatcher) {
+            val packageName =
+                packageManagerWrapper
+                    .resolveActivity(/* intent= */ intent, /* flags= */ 0)
+                    ?.activityInfo
+                    ?.packageName ?: return@withContext null
+            return@withContext getPackageInfo(packageName)
+        }
+
+    private suspend fun getPackageIcon(packageName: String): Icon? {
+        val appInfo = getPackageInfo(packageName)?.applicationInfo ?: return null
+        return if (appInfo.icon != 0) {
+            Icon.createWithResource(appInfo.packageName, appInfo.icon)
+        } else {
+            null
+        }
+    }
+
+    private suspend fun getPackageInfo(packageName: String): PackageInfo? =
+        withContext(backgroundDispatcher) {
+            try {
+                return@withContext packageManager.getPackageInfo(packageName, /* flags= */ 0)
+            } catch (e: RemoteException) {
+                Log.e(TAG, "Failed to retrieve package info for $packageName")
+                return@withContext null
+            }
+        }
+
+    companion object {
+        private const val TAG = "DefaultAppsIconProvider"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt b/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt
index 160ae86..fe54044 100644
--- a/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/settings/SettingsProxy.kt
@@ -19,6 +19,7 @@
 import android.database.ContentObserver
 import android.net.Uri
 import android.provider.Settings.SettingNotFoundException
+import androidx.annotation.WorkerThread
 import com.android.app.tracing.TraceUtils.trace
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
@@ -59,10 +60,13 @@
     fun getUriFor(name: String): Uri
 
     /**
-     * Convenience wrapper around [ContentResolver.registerContentObserver].'
-     *
+     * Registers listener for a given content observer <b>while blocking the current thread</b>.
      * Implicitly calls [getUriFor] on the passed in name.
+     *
+     * This should not be called from the main thread, use [registerContentObserver] or
+     * [registerContentObserverAsync] instead.
      */
+    @WorkerThread
     fun registerContentObserverSync(name: String, settingsObserver: ContentObserver) {
         registerContentObserverSync(getUriFor(name), settingsObserver)
     }
@@ -90,7 +94,13 @@
             registerContentObserverSync(getUriFor(name), settingsObserver)
         }
 
-    /** Convenience wrapper around [ContentResolver.registerContentObserver].' */
+    /**
+     * Registers listener for a given content observer <b>while blocking the current thread</b>.
+     *
+     * This should not be called from the main thread, use [registerContentObserver] or
+     * [registerContentObserverAsync] instead.
+     */
+    @WorkerThread
     fun registerContentObserverSync(uri: Uri, settingsObserver: ContentObserver) =
         registerContentObserverSync(uri, false, settingsObserver)
 
@@ -157,7 +167,13 @@
             registerContentObserverSync(getUriFor(name), notifyForDescendants, settingsObserver)
         }
 
-    /** Convenience wrapper around [ContentResolver.registerContentObserver].' */
+    /**
+     * Registers listener for a given content observer <b>while blocking the current thread</b>.
+     *
+     * This should not be called from the main thread, use [registerContentObserver] or
+     * [registerContentObserverAsync] instead.
+     */
+    @WorkerThread
     fun registerContentObserverSync(
         uri: Uri,
         notifyForDescendants: Boolean,
@@ -200,7 +216,13 @@
             registerContentObserverSync(uri, notifyForDescendants, settingsObserver)
         }
 
-    /** See [ContentResolver.unregisterContentObserver]. */
+    /**
+     * Unregisters the given content observer <b>while blocking the current thread</b>.
+     *
+     * This should not be called from the main thread, use [unregisterContentObserver] or
+     * [unregisterContentObserverAsync] instead.
+     */
+    @WorkerThread
     fun unregisterContentObserverSync(settingsObserver: ContentObserver) {
         trace({ "SP#unregisterObserver" }) {
             getContentResolver().unregisterContentObserver(settingsObserver)
diff --git a/packages/SystemUI/src/com/android/systemui/volume/VolumeControllerCollector.kt b/packages/SystemUI/src/com/android/systemui/volume/VolumeControllerCollector.kt
new file mode 100644
index 0000000..6859191
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/volume/VolumeControllerCollector.kt
@@ -0,0 +1,59 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.volume
+
+import android.media.IVolumeController
+import com.android.settingslib.media.data.repository.VolumeControllerEvent
+import com.android.systemui.dagger.qualifiers.Application
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.launch
+
+/**
+ * This class is a bridge between
+ * [com.android.settingslib.volume.data.repository.AudioRepository.volumeControllerEvents] and the
+ * old code that uses [IVolumeController] interface directly.
+ */
+class VolumeControllerCollector
+@Inject
+constructor(@Application private val coroutineScope: CoroutineScope) {
+
+    /** Collects [Flow] of [VolumeControllerEvent] into [IVolumeController]. */
+    fun collectToController(
+        eventsFlow: Flow<VolumeControllerEvent>,
+        controller: IVolumeController
+    ) =
+        coroutineScope.launch {
+            eventsFlow.collect { event ->
+                when (event) {
+                    is VolumeControllerEvent.VolumeChanged ->
+                        controller.volumeChanged(event.streamType, event.flags)
+                    VolumeControllerEvent.Dismiss -> controller.dismiss()
+                    is VolumeControllerEvent.DisplayCsdWarning ->
+                        controller.displayCsdWarning(event.csdWarning, event.displayDurationMs)
+                    is VolumeControllerEvent.DisplaySafeVolumeWarning ->
+                        controller.displaySafeVolumeWarning(event.flags)
+                    is VolumeControllerEvent.MasterMuteChanged ->
+                        controller.masterMuteChanged(event.flags)
+                    is VolumeControllerEvent.SetA11yMode -> controller.setA11yMode(event.mode)
+                    is VolumeControllerEvent.SetLayoutDirection ->
+                        controller.setLayoutDirection(event.layoutDirection)
+                }
+            }
+        }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepositoryTest.kt
new file mode 100644
index 0000000..84ec1a5
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepositoryTest.kt
@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.mediarouter.data.repository
+
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.statusbar.policy.CastDevice
+import com.android.systemui.statusbar.policy.fakeCastController
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.Test
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+
+@SmallTest
+@OptIn(ExperimentalCoroutinesApi::class)
+class MediaRouterRepositoryTest : SysuiTestCase() {
+    val kosmos = Kosmos()
+    val testScope = kosmos.testScope
+    val castController = kosmos.fakeCastController
+
+    val underTest = kosmos.realMediaRouterRepository
+
+    @Test
+    fun castDevices_empty_isEmpty() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.castDevices)
+            // Required to let the listener attach before the devices get set
+            runCurrent()
+
+            castController.castDevices = emptyList()
+
+            assertThat(latest).isEmpty()
+        }
+
+    @Test
+    fun castDevices_onlyIncludesMediaRouterOriginDevices() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.castDevices)
+            // Required to let the listener attach before the devices get set
+            runCurrent()
+
+            val projectionDevice =
+                CastDevice(
+                    id = "idProjection",
+                    name = "name",
+                    description = "desc",
+                    state = CastDevice.CastState.Connected,
+                    origin = CastDevice.CastOrigin.MediaProjection,
+                )
+            val routerDevice1 =
+                CastDevice(
+                    id = "idRouter1",
+                    name = "name",
+                    description = "desc",
+                    state = CastDevice.CastState.Connected,
+                    origin = CastDevice.CastOrigin.MediaRouter,
+                )
+
+            val routerDevice2 =
+                CastDevice(
+                    id = "idRouter2",
+                    name = "name",
+                    description = "desc",
+                    state = CastDevice.CastState.Connected,
+                    origin = CastDevice.CastOrigin.MediaRouter,
+                )
+            castController.setCastDevices(listOf(projectionDevice, routerDevice1, routerDevice2))
+
+            assertThat(latest).containsExactly(routerDevice1, routerDevice2).inOrder()
+        }
+
+    @Test
+    fun stopCasting_notifiesCastController() {
+        val device =
+            CastDevice(
+                id = "id",
+                name = "name",
+                description = "desc",
+                state = CastDevice.CastState.Connected,
+                origin = CastDevice.CastOrigin.MediaRouter,
+            )
+
+        underTest.stopCasting(device)
+
+        assertThat(castController.lastStoppedDevice).isEqualTo(device)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/domian/interactor/MediaRouterChipInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/domian/interactor/MediaRouterChipInteractorTest.kt
new file mode 100644
index 0000000..8a6a50d
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/domian/interactor/MediaRouterChipInteractorTest.kt
@@ -0,0 +1,220 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.chips.casttootherdevice.domian.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.mediarouter.data.repository.fakeMediaRouterRepository
+import com.android.systemui.statusbar.chips.casttootherdevice.domain.interactor.mediaRouterChipInteractor
+import com.android.systemui.statusbar.chips.casttootherdevice.domain.model.MediaRouterCastModel
+import com.android.systemui.statusbar.policy.CastDevice
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.Test
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+
+@SmallTest
+@OptIn(ExperimentalCoroutinesApi::class)
+class MediaRouterChipInteractorTest : SysuiTestCase() {
+    val kosmos = Kosmos()
+    val testScope = kosmos.testScope
+    val mediaRouterRepository = kosmos.fakeMediaRouterRepository
+
+    val underTest = kosmos.mediaRouterChipInteractor
+
+    @Test
+    fun mediaRouterCastingState_noDevices_doingNothing() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.mediaRouterCastingState)
+
+            mediaRouterRepository.castDevices.value = emptyList()
+
+            assertThat(latest).isEqualTo(MediaRouterCastModel.DoingNothing)
+        }
+
+    @Test
+    fun mediaRouterCastingState_disconnectedDevice_doingNothing() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.mediaRouterCastingState)
+
+            mediaRouterRepository.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Disconnected,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+
+            assertThat(latest).isEqualTo(MediaRouterCastModel.DoingNothing)
+        }
+
+    @Test
+    fun mediaRouterCastingState_connectingDevice_casting() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.mediaRouterCastingState)
+
+            mediaRouterRepository.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Connecting,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+
+            assertThat(latest).isEqualTo(MediaRouterCastModel.Casting)
+        }
+
+    @Test
+    fun mediaRouterCastingState_connectedDevice_casting() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.mediaRouterCastingState)
+
+            mediaRouterRepository.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Connected,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+
+            assertThat(latest).isEqualTo(MediaRouterCastModel.Casting)
+        }
+
+    @Test
+    fun stopCasting_noDevices_doesNothing() =
+        testScope.runTest {
+            collectLastValue(underTest.mediaRouterCastingState)
+
+            mediaRouterRepository.castDevices.value = emptyList()
+            // Let the interactor catch up to the repo value
+            runCurrent()
+
+            underTest.stopCasting()
+
+            assertThat(mediaRouterRepository.lastStoppedDevice).isNull()
+        }
+
+    @Test
+    fun stopCasting_disconnectedDevice_doesNothing() =
+        testScope.runTest {
+            collectLastValue(underTest.mediaRouterCastingState)
+
+            mediaRouterRepository.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Disconnected,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+            // Let the interactor catch up to the repo value
+            runCurrent()
+
+            underTest.stopCasting()
+
+            assertThat(mediaRouterRepository.lastStoppedDevice).isNull()
+        }
+
+    @Test
+    fun stopCasting_connectingDevice_notifiesRepo() =
+        testScope.runTest {
+            collectLastValue(underTest.mediaRouterCastingState)
+
+            val device =
+                CastDevice(
+                    state = CastDevice.CastState.Connecting,
+                    id = "id",
+                    name = "name",
+                    description = "desc",
+                    origin = CastDevice.CastOrigin.MediaRouter,
+                )
+            mediaRouterRepository.castDevices.value = listOf(device)
+            // Let the interactor catch up to the repo value
+            runCurrent()
+
+            underTest.stopCasting()
+
+            assertThat(mediaRouterRepository.lastStoppedDevice).isEqualTo(device)
+        }
+
+    @Test
+    fun stopCasting_connectedDevice_notifiesRepo() =
+        testScope.runTest {
+            collectLastValue(underTest.mediaRouterCastingState)
+
+            val device =
+                CastDevice(
+                    state = CastDevice.CastState.Connected,
+                    id = "id",
+                    name = "name",
+                    description = "desc",
+                    origin = CastDevice.CastOrigin.MediaRouter,
+                )
+            mediaRouterRepository.castDevices.value = listOf(device)
+            // Let the interactor catch up to the repo value
+            runCurrent()
+
+            underTest.stopCasting()
+
+            assertThat(mediaRouterRepository.lastStoppedDevice).isEqualTo(device)
+        }
+
+    @Test
+    fun stopCasting_multipleConnectedDevices_notifiesRepoOfFirst() =
+        testScope.runTest {
+            collectLastValue(underTest.mediaRouterCastingState)
+
+            val device1 =
+                CastDevice(
+                    state = CastDevice.CastState.Connected,
+                    id = "id1",
+                    name = "name",
+                    description = "desc",
+                    origin = CastDevice.CastOrigin.MediaRouter,
+                )
+            val device2 =
+                CastDevice(
+                    state = CastDevice.CastState.Connected,
+                    id = "id2",
+                    name = "name",
+                    description = "desc",
+                    origin = CastDevice.CastOrigin.MediaRouter,
+                )
+            mediaRouterRepository.castDevices.value = listOf(device1, device2)
+            // Let the interactor catch up to the repo value
+            runCurrent()
+
+            underTest.stopCasting()
+
+            assertThat(mediaRouterRepository.lastStoppedDevice).isEqualTo(device1)
+        }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastScreenToOtherDeviceDialogDelegateTest.kt
similarity index 92%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegateTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastScreenToOtherDeviceDialogDelegateTest.kt
index c6fb481..e9d6f0e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastToOtherDeviceDialogDelegateTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndCastScreenToOtherDeviceDialogDelegateTest.kt
@@ -48,10 +48,10 @@
 
 @SmallTest
 @OptIn(ExperimentalCoroutinesApi::class)
-class EndCastToOtherDeviceDialogDelegateTest : SysuiTestCase() {
+class EndCastScreenToOtherDeviceDialogDelegateTest : SysuiTestCase() {
     private val kosmos = Kosmos().also { it.testCase = this }
     private val sysuiDialog = mock<SystemUIDialog>()
-    private lateinit var underTest: EndCastToOtherDeviceDialogDelegate
+    private lateinit var underTest: EndCastScreenToOtherDeviceDialogDelegate
 
     @Test
     fun icon() {
@@ -68,7 +68,7 @@
 
         underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
 
-        verify(sysuiDialog).setTitle(R.string.cast_to_other_device_stop_dialog_title)
+        verify(sysuiDialog).setTitle(R.string.cast_screen_to_other_device_stop_dialog_title)
     }
 
     @Test
@@ -78,7 +78,7 @@
         underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
 
         verify(sysuiDialog)
-            .setMessage(context.getString(R.string.cast_to_other_device_stop_dialog_message))
+            .setMessage(context.getString(R.string.cast_screen_to_other_device_stop_dialog_message))
     }
 
     @Test
@@ -99,7 +99,7 @@
 
         underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
 
-        // It'd be nice to use R.string.cast_to_other_device_stop_dialog_message_specific_app
+        // It'd be nice to use R.string.cast_screen_to_other_device_stop_dialog_message_specific_app
         // directly, but it includes the <b> tags which aren't in the returned string.
         val result = argumentCaptor<CharSequence>()
         verify(sysuiDialog).setMessage(result.capture())
@@ -142,7 +142,7 @@
 
     private fun createAndSetDelegate(state: MediaProjectionState.Projecting) {
         underTest =
-            EndCastToOtherDeviceDialogDelegate(
+            EndCastScreenToOtherDeviceDialogDelegate(
                 kosmos.endMediaProjectionDialogHelper,
                 stopAction = kosmos.mediaProjectionChipInteractor::stopProjecting,
                 ProjectionChipModel.Projecting(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndGenericCastToOtherDeviceDialogDelegateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndGenericCastToOtherDeviceDialogDelegateTest.kt
new file mode 100644
index 0000000..0af423d
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/view/EndGenericCastToOtherDeviceDialogDelegateTest.kt
@@ -0,0 +1,132 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.chips.casttootherdevice.ui.view
+
+import android.content.DialogInterface
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testCase
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.mediarouter.data.repository.fakeMediaRouterRepository
+import com.android.systemui.res.R
+import com.android.systemui.statusbar.chips.casttootherdevice.domain.interactor.mediaRouterChipInteractor
+import com.android.systemui.statusbar.chips.mediaprojection.ui.view.endMediaProjectionDialogHelper
+import com.android.systemui.statusbar.phone.SystemUIDialog
+import com.android.systemui.statusbar.policy.CastDevice
+import com.google.common.truth.Truth.assertThat
+import kotlin.test.Test
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+
+@SmallTest
+@OptIn(ExperimentalCoroutinesApi::class)
+class EndGenericCastToOtherDeviceDialogDelegateTest : SysuiTestCase() {
+    private val kosmos = Kosmos().also { it.testCase = this }
+    private val sysuiDialog = mock<SystemUIDialog>()
+    private lateinit var underTest: EndGenericCastToOtherDeviceDialogDelegate
+
+    @Test
+    fun icon() {
+        createAndSetDelegate()
+
+        underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
+
+        verify(sysuiDialog).setIcon(R.drawable.ic_cast_connected)
+    }
+
+    @Test
+    fun title() {
+        createAndSetDelegate()
+
+        underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
+
+        verify(sysuiDialog).setTitle(R.string.cast_to_other_device_stop_dialog_title)
+    }
+
+    @Test
+    fun message() {
+        createAndSetDelegate()
+
+        underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
+
+        verify(sysuiDialog).setMessage(R.string.cast_to_other_device_stop_dialog_message)
+    }
+
+    @Test
+    fun negativeButton() {
+        createAndSetDelegate()
+
+        underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
+
+        verify(sysuiDialog).setNegativeButton(R.string.close_dialog_button, null)
+    }
+
+    @Test
+    fun positiveButton() =
+        kosmos.testScope.runTest {
+            createAndSetDelegate()
+
+            // Set up a real device so the stop works correctly
+            collectLastValue(kosmos.mediaRouterChipInteractor.mediaRouterCastingState)
+            val device =
+                CastDevice(
+                    state = CastDevice.CastState.Connected,
+                    id = "id",
+                    name = "name",
+                    description = "desc",
+                    origin = CastDevice.CastOrigin.MediaRouter,
+                )
+            kosmos.fakeMediaRouterRepository.castDevices.value = listOf(device)
+            // Let everything catch up to the repo value
+            runCurrent()
+            runCurrent()
+
+            underTest.beforeCreate(sysuiDialog, /* savedInstanceState= */ null)
+
+            val clickListener = argumentCaptor<DialogInterface.OnClickListener>()
+
+            // Verify the button has the right text
+            verify(sysuiDialog)
+                .setPositiveButton(
+                    eq(R.string.cast_to_other_device_stop_dialog_button),
+                    clickListener.capture()
+                )
+
+            // Verify that clicking the button stops the recording
+            assertThat(kosmos.fakeMediaRouterRepository.lastStoppedDevice).isNull()
+
+            clickListener.firstValue.onClick(mock<DialogInterface>(), 0)
+            runCurrent()
+
+            assertThat(kosmos.fakeMediaRouterRepository.lastStoppedDevice).isEqualTo(device)
+        }
+
+    private fun createAndSetDelegate() {
+        underTest =
+            EndGenericCastToOtherDeviceDialogDelegate(
+                kosmos.endMediaProjectionDialogHelper,
+                stopAction = kosmos.mediaRouterChipInteractor::stopCasting,
+            )
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
index 74b6ae2..fe29140 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
@@ -20,6 +20,7 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.mockDialogTransitionAnimator
+import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.Kosmos
@@ -28,8 +29,10 @@
 import com.android.systemui.mediaprojection.data.model.MediaProjectionState
 import com.android.systemui.mediaprojection.data.repository.fakeMediaProjectionRepository
 import com.android.systemui.mediaprojection.taskswitcher.FakeActivityTaskManager.Companion.createTask
+import com.android.systemui.mediarouter.data.repository.fakeMediaRouterRepository
 import com.android.systemui.res.R
-import com.android.systemui.statusbar.chips.casttootherdevice.ui.view.EndCastToOtherDeviceDialogDelegate
+import com.android.systemui.statusbar.chips.casttootherdevice.ui.view.EndCastScreenToOtherDeviceDialogDelegate
+import com.android.systemui.statusbar.chips.casttootherdevice.ui.view.EndGenericCastToOtherDeviceDialogDelegate
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.MediaProjectionChipInteractorTest.Companion.CAST_TO_OTHER_DEVICES_PACKAGE
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.MediaProjectionChipInteractorTest.Companion.NORMAL_PACKAGE
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.MediaProjectionChipInteractorTest.Companion.setUpPackageManagerForMediaProjection
@@ -38,6 +41,7 @@
 import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.phone.mockSystemUIDialogFactory
+import com.android.systemui.statusbar.policy.CastDevice
 import com.android.systemui.util.time.fakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.Test
@@ -55,9 +59,11 @@
     private val kosmos = Kosmos().also { it.testCase = this }
     private val testScope = kosmos.testScope
     private val mediaProjectionRepo = kosmos.fakeMediaProjectionRepository
+    private val mediaRouterRepo = kosmos.fakeMediaRouterRepository
     private val systemClock = kosmos.fakeSystemClock
 
-    private val mockCastDialog = mock<SystemUIDialog>()
+    private val mockScreenCastDialog = mock<SystemUIDialog>()
+    private val mockGenericCastDialog = mock<SystemUIDialog>()
 
     private val chipBackgroundView = mock<ChipBackgroundContainer>()
     private val chipView =
@@ -76,14 +82,25 @@
     fun setUp() {
         setUpPackageManagerForMediaProjection(kosmos)
 
-        whenever(kosmos.mockSystemUIDialogFactory.create(any<EndCastToOtherDeviceDialogDelegate>()))
-            .thenReturn(mockCastDialog)
+        whenever(
+                kosmos.mockSystemUIDialogFactory.create(
+                    any<EndCastScreenToOtherDeviceDialogDelegate>()
+                )
+            )
+            .thenReturn(mockScreenCastDialog)
+        whenever(
+                kosmos.mockSystemUIDialogFactory.create(
+                    any<EndGenericCastToOtherDeviceDialogDelegate>()
+                )
+            )
+            .thenReturn(mockGenericCastDialog)
     }
 
     @Test
     fun chip_notProjectingState_isHidden() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
+            mediaRouterRepo.castDevices.value = emptyList()
 
             mediaProjectionRepo.mediaProjectionState.value = MediaProjectionState.NotProjecting
 
@@ -91,9 +108,10 @@
         }
 
     @Test
-    fun chip_singleTaskState_otherDevicesPackage_isShownAsTimer() =
+    fun chip_projectionIsSingleTaskState_otherDevicesPackage_isShownAsTimer_forScreen() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
+            mediaRouterRepo.castDevices.value = emptyList()
 
             mediaProjectionRepo.mediaProjectionState.value =
                 MediaProjectionState.Projecting.SingleTask(
@@ -104,13 +122,15 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon = (latest as OngoingActivityChipModel.Shown).icon
             assertThat((icon as Icon.Resource).res).isEqualTo(R.drawable.ic_cast_connected)
-            assertThat(icon.contentDescription).isNotNull()
+            assertThat((icon.contentDescription as ContentDescription.Resource).res)
+                .isEqualTo(R.string.cast_screen_to_other_device_chip_accessibility_label)
         }
 
     @Test
-    fun chip_entireScreenState_otherDevicesPackage_isShownAsTimer() =
+    fun chip_projectionIsEntireScreenState_otherDevicesPackage_isShownAsTimer_forScreen() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
+            mediaRouterRepo.castDevices.value = emptyList()
 
             mediaProjectionRepo.mediaProjectionState.value =
                 MediaProjectionState.Projecting.EntireScreen(CAST_TO_OTHER_DEVICES_PACKAGE)
@@ -118,7 +138,72 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon = (latest as OngoingActivityChipModel.Shown).icon
             assertThat((icon as Icon.Resource).res).isEqualTo(R.drawable.ic_cast_connected)
-            assertThat(icon.contentDescription).isNotNull()
+            assertThat((icon.contentDescription as ContentDescription.Resource).res)
+                .isEqualTo(R.string.cast_screen_to_other_device_chip_accessibility_label)
+        }
+
+    @Test
+    fun chip_routerStateDoingNothing_isHidden() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+            mediaProjectionRepo.mediaProjectionState.value = MediaProjectionState.NotProjecting
+
+            mediaRouterRepo.castDevices.value = emptyList()
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Hidden::class.java)
+        }
+
+    @Test
+    fun chip_routerStateCasting_isShownAsGenericIconOnly() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+            mediaProjectionRepo.mediaProjectionState.value = MediaProjectionState.NotProjecting
+
+            mediaRouterRepo.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Connected,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.IconOnly::class.java)
+            val icon = (latest as OngoingActivityChipModel.Shown).icon
+            assertThat((icon as Icon.Resource).res).isEqualTo(R.drawable.ic_cast_connected)
+            // This content description is just generic "Casting", not "Casting screen"
+            assertThat((icon.contentDescription as ContentDescription.Resource).res)
+                .isEqualTo(R.string.accessibility_casting)
+        }
+
+    @Test
+    fun chip_projectingAndRouterCasting_projectionInfoShown() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            mediaProjectionRepo.mediaProjectionState.value =
+                MediaProjectionState.Projecting.EntireScreen(CAST_TO_OTHER_DEVICES_PACKAGE)
+            mediaRouterRepo.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Connected,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+
+            // Only the projection info will show a timer
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
+            val icon = (latest as OngoingActivityChipModel.Shown).icon
+            assertThat((icon as Icon.Resource).res).isEqualTo(R.drawable.ic_cast_connected)
+            // MediaProjection == screen casting, so this content description reflects that we're
+            // using the MediaProjection information.
+            assertThat((icon.contentDescription as ContentDescription.Resource).res)
+                .isEqualTo(R.string.cast_screen_to_other_device_chip_accessibility_label)
         }
 
     @Test
@@ -133,7 +218,7 @@
         }
 
     @Test
-    fun chip_singleTaskState_normalPackage_isHidden() =
+    fun chip_projectionIsSingleTaskState_normalPackage_isHidden() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
 
@@ -144,7 +229,7 @@
         }
 
     @Test
-    fun chip_entireScreenState_normalPackage_isHidden() =
+    fun chip_projectionIsEntireScreenState_normalPackage_isHidden() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
 
@@ -155,7 +240,7 @@
         }
 
     @Test
-    fun chip_timeResetsOnEachNewShare() =
+    fun chip_projectionOnly_timeResetsOnEachNewShare() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
 
@@ -181,7 +266,38 @@
         }
 
     @Test
-    fun chip_entireScreen_clickListenerShowsCastDialog() =
+    fun chip_routerInfoThenProjectionInfo_switchesToTimer() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            // First, set only MediaRouter to have information and verify we just show the icon
+            systemClock.setElapsedRealtime(1234)
+            mediaProjectionRepo.mediaProjectionState.value = MediaProjectionState.NotProjecting
+            mediaRouterRepo.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Connected,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.IconOnly::class.java)
+
+            // Later, set MediaProjection to also have information
+            systemClock.setElapsedRealtime(5678)
+            mediaProjectionRepo.mediaProjectionState.value =
+                MediaProjectionState.Projecting.EntireScreen(CAST_TO_OTHER_DEVICES_PACKAGE)
+
+            // Verify the new time is used
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
+            assertThat((latest as OngoingActivityChipModel.Shown.Timer).startTimeMs).isEqualTo(5678)
+        }
+
+    @Test
+    fun chip_projectionStateEntireScreen_clickListenerShowsScreenCastDialog() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
             mediaProjectionRepo.mediaProjectionState.value =
@@ -193,7 +309,7 @@
             clickListener!!.onClick(chipView)
             verify(kosmos.mockDialogTransitionAnimator)
                 .showFromView(
-                    eq(mockCastDialog),
+                    eq(mockScreenCastDialog),
                     eq(chipBackgroundView),
                     eq(null),
                     ArgumentMatchers.anyBoolean(),
@@ -201,7 +317,7 @@
         }
 
     @Test
-    fun chip_singleTask_clickListenerShowsCastDialog() =
+    fun chip_projectionStateSingleTask_clickListenerShowsScreenCastDialog() =
         testScope.runTest {
             val latest by collectLastValue(underTest.chip)
 
@@ -217,7 +333,36 @@
             clickListener!!.onClick(chipView)
             verify(kosmos.mockDialogTransitionAnimator)
                 .showFromView(
-                    eq(mockCastDialog),
+                    eq(mockScreenCastDialog),
+                    eq(chipBackgroundView),
+                    eq(null),
+                    ArgumentMatchers.anyBoolean(),
+                )
+        }
+
+    @Test
+    fun chip_routerStateCasting_clickListenerShowsGenericCastDialog() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            mediaRouterRepo.castDevices.value =
+                listOf(
+                    CastDevice(
+                        state = CastDevice.CastState.Connected,
+                        id = "id",
+                        name = "name",
+                        description = "desc",
+                        origin = CastDevice.CastOrigin.MediaRouter,
+                    )
+                )
+
+            val clickListener = ((latest as OngoingActivityChipModel.Shown).onClickListener)
+            assertThat(clickListener).isNotNull()
+
+            clickListener!!.onClick(chipView)
+            verify(kosmos.mockDialogTransitionAnimator)
+                .showFromView(
+                    eq(mockGenericCastDialog),
                     eq(chipBackgroundView),
                     eq(null),
                     ArgumentMatchers.anyBoolean(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/mediaprojection/ui/view/EndMediaProjectionDialogHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/mediaprojection/ui/view/EndMediaProjectionDialogHelperTest.kt
index ee929ae..f9ad5ac 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/mediaprojection/ui/view/EndMediaProjectionDialogHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/mediaprojection/ui/view/EndMediaProjectionDialogHelperTest.kt
@@ -61,7 +61,7 @@
             underTest.getDialogMessage(
                 MediaProjectionState.Projecting.EntireScreen("host.package"),
                 R.string.accessibility_home,
-                R.string.cast_to_other_device_stop_dialog_message_specific_app
+                R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
             )
 
         assertThat(result).isEqualTo(context.getString(R.string.accessibility_home))
@@ -84,7 +84,7 @@
             underTest.getDialogMessage(
                 projectionState,
                 R.string.accessibility_home,
-                R.string.cast_to_other_device_stop_dialog_message_specific_app
+                R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
             )
 
         assertThat(result).isEqualTo(context.getString(R.string.accessibility_home))
@@ -109,7 +109,7 @@
             underTest.getDialogMessage(
                 projectionState,
                 R.string.accessibility_home,
-                R.string.cast_to_other_device_stop_dialog_message_specific_app
+                R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
             )
 
         // It'd be nice to use the R.string resources directly, but they include the <b> tags which
@@ -123,7 +123,7 @@
             underTest.getDialogMessage(
                 specificTaskInfo = null,
                 R.string.accessibility_home,
-                R.string.cast_to_other_device_stop_dialog_message_specific_app,
+                R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
             )
 
         assertThat(result).isEqualTo(context.getString(R.string.accessibility_home))
@@ -141,7 +141,7 @@
             underTest.getDialogMessage(
                 task,
                 R.string.accessibility_home,
-                R.string.cast_to_other_device_stop_dialog_message_specific_app
+                R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
             )
 
         assertThat(result).isEqualTo(context.getString(R.string.accessibility_home))
@@ -161,7 +161,7 @@
             underTest.getDialogMessage(
                 task,
                 R.string.accessibility_home,
-                R.string.cast_to_other_device_stop_dialog_message_specific_app
+                R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
             )
 
         assertThat(result.toString()).isEqualTo("You will stop casting Fake Package")
@@ -186,7 +186,7 @@
             underTest.getDialogMessage(
                 projectionState,
                 R.string.accessibility_home,
-                R.string.cast_to_other_device_stop_dialog_message_specific_app
+                R.string.cast_screen_to_other_device_stop_dialog_message_specific_app,
             )
 
         assertThat(result.toString()).isEqualTo("You will stop casting Fake & Package <Here>")
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt
index 25533d8..d87b3e2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/coordinator/KeyguardCoordinatorTest.kt
@@ -19,16 +19,23 @@
 
 import android.app.Notification
 import android.os.UserHandle
+import android.platform.test.flag.junit.FlagsParameterization
 import android.provider.Settings
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.log.logcatLogBuffer
+import com.android.systemui.flags.andSceneContainer
 import com.android.systemui.keyguard.data.repository.FakeKeyguardRepository
-import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.TransitionStep
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.log.logcatLogBuffer
 import com.android.systemui.plugins.statusbar.StatusBarStateController
+import com.android.systemui.scene.data.repository.Idle
+import com.android.systemui.scene.data.repository.setTransition
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.StatusBarState
 import com.android.systemui.statusbar.notification.collection.GroupEntryBuilder
 import com.android.systemui.statusbar.notification.collection.NotifPipeline
@@ -49,6 +56,8 @@
 import com.android.systemui.util.mockito.withArgCaptor
 import com.android.systemui.util.settings.FakeSettings
 import com.google.common.truth.Truth.assertThat
+import java.util.function.Consumer
+import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestCoroutineScheduler
@@ -62,22 +71,28 @@
 import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
-import java.util.function.Consumer
-import kotlin.time.Duration.Companion.seconds
 import org.mockito.Mockito.`when` as whenever
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
 
 @SmallTest
-@RunWith(AndroidJUnit4::class)
-class KeyguardCoordinatorTest : SysuiTestCase() {
+@RunWith(ParameterizedAndroidJunit4::class)
+class KeyguardCoordinatorTest(flags: FlagsParameterization) : SysuiTestCase() {
+
+    private val kosmos = Kosmos()
 
     private val headsUpManager: HeadsUpManager = mock()
     private val keyguardNotifVisibilityProvider: KeyguardNotificationVisibilityProvider = mock()
     private val keyguardRepository = FakeKeyguardRepository()
-    private val keyguardTransitionRepository = FakeKeyguardTransitionRepository()
+    private val keyguardTransitionRepository = kosmos.fakeKeyguardTransitionRepository
     private val notifPipeline: NotifPipeline = mock()
     private val sectionHeaderVisibilityProvider: SectionHeaderVisibilityProvider = mock()
     private val statusBarStateController: StatusBarStateController = mock()
 
+    init {
+        mSetFlagsRule.setFlagsParameterization(flags)
+    }
+
     @Test
     fun testSetSectionHeadersVisibleInShade() = runKeyguardCoordinatorTest {
         clearInvocations(sectionHeaderVisibilityProvider)
@@ -147,10 +162,9 @@
         keyguardRepository.setKeyguardShowing(false)
         whenever(statusBarStateController.isExpanded).thenReturn(false)
         runKeyguardCoordinatorTest {
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Gone),
+                stateTransition = TransitionStep(KeyguardState.LOCKSCREEN, KeyguardState.GONE)
             )
 
             // WHEN: A notification is posted
@@ -163,24 +177,20 @@
 
             // WHEN: The keyguard is now showing
             keyguardRepository.setKeyguardShowing(true)
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.AOD,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition = TransitionStep(KeyguardState.GONE, KeyguardState.AOD)
             )
-            testScheduler.runCurrent()
 
             // THEN: The notification is recognized as "seen" and is filtered out.
             assertThat(unseenFilter.shouldFilterOut(fakeEntry, 0L)).isTrue()
 
             // WHEN: The keyguard goes away
             keyguardRepository.setKeyguardShowing(false)
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.AOD,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Gone),
+                stateTransition = TransitionStep(KeyguardState.AOD, KeyguardState.GONE)
             )
-            testScheduler.runCurrent()
 
             // THEN: The notification is shown regardless
             assertThat(unseenFilter.shouldFilterOut(fakeEntry, 0L)).isFalse()
@@ -344,9 +354,9 @@
             val fakeEntry = NotificationEntryBuilder().build()
             collectionListener.onEntryAdded(fakeEntry)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.AOD,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.AOD,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
@@ -356,21 +366,17 @@
 
             // WHEN: Keyguard is no longer showing
             keyguardRepository.setKeyguardShowing(false)
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Gone),
+                stateTransition = TransitionStep(KeyguardState.LOCKSCREEN, KeyguardState.GONE)
             )
-            testScheduler.runCurrent()
 
             // WHEN: Keyguard is shown again
             keyguardRepository.setKeyguardShowing(true)
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.AOD,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition = TransitionStep(KeyguardState.GONE, KeyguardState.AOD)
             )
-            testScheduler.runCurrent()
 
             // THEN: The notification is now recognized as "seen" and is filtered out.
             assertThat(unseenFilter.shouldFilterOut(fakeEntry, 0L)).isTrue()
@@ -383,9 +389,9 @@
         keyguardRepository.setKeyguardShowing(true)
         runKeyguardCoordinatorTest {
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.GONE,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             val fakeEntry = NotificationEntryBuilder().build()
             collectionListener.onEntryAdded(fakeEntry)
@@ -393,9 +399,9 @@
             // WHEN: Keyguard is no longer showing
             keyguardRepository.setKeyguardShowing(false)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.GONE,
+                this.testScheduler,
             )
 
             // WHEN: Keyguard is shown again
@@ -413,10 +419,9 @@
         keyguardRepository.setKeyguardShowing(true)
         keyguardRepository.setIsDozing(false)
         runKeyguardCoordinatorTest {
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition = TransitionStep(KeyguardState.GONE, KeyguardState.LOCKSCREEN)
             )
             val firstEntry = NotificationEntryBuilder().setId(1).build()
             collectionListener.onEntryAdded(firstEntry)
@@ -437,21 +442,17 @@
 
             // WHEN: the keyguard is no longer showing
             keyguardRepository.setKeyguardShowing(false)
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Gone),
+                stateTransition = TransitionStep(KeyguardState.LOCKSCREEN, KeyguardState.GONE)
             )
-            testScheduler.runCurrent()
 
             // WHEN: Keyguard is shown again
             keyguardRepository.setKeyguardShowing(true)
-            keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition = TransitionStep(KeyguardState.GONE, KeyguardState.LOCKSCREEN)
             )
-            testScheduler.runCurrent()
 
             // THEN: The first notification is considered seen and is filtered out.
             assertThat(unseenFilter.shouldFilterOut(firstEntry, 0L)).isTrue()
@@ -468,9 +469,9 @@
         keyguardRepository.setIsDozing(false)
         runKeyguardCoordinatorTest {
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.GONE,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
@@ -498,18 +499,18 @@
             // WHEN: the keyguard is no longer showing
             keyguardRepository.setKeyguardShowing(false)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.GONE,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
             // WHEN: Keyguard is shown again
             keyguardRepository.setKeyguardShowing(true)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.GONE,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
@@ -525,9 +526,9 @@
         keyguardRepository.setIsDozing(false)
         runKeyguardCoordinatorTest {
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.GONE,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
@@ -555,18 +556,18 @@
             // WHEN: the keyguard is no longer showing
             keyguardRepository.setKeyguardShowing(false)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.GONE,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
             // WHEN: Keyguard is shown again
             keyguardRepository.setKeyguardShowing(true)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.GONE,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
@@ -582,9 +583,9 @@
         keyguardRepository.setIsDozing(false)
         runKeyguardCoordinatorTest {
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.GONE,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
@@ -608,18 +609,18 @@
             // WHEN: the keyguard is no longer showing
             keyguardRepository.setKeyguardShowing(false)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    this.testScheduler,
+                from = KeyguardState.LOCKSCREEN,
+                to = KeyguardState.GONE,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
             // WHEN: Keyguard is shown again
             keyguardRepository.setKeyguardShowing(true)
             keyguardTransitionRepository.sendTransitionSteps(
-                    from = KeyguardState.GONE,
-                    to = KeyguardState.LOCKSCREEN,
-                    this.testScheduler,
+                from = KeyguardState.GONE,
+                to = KeyguardState.LOCKSCREEN,
+                this.testScheduler,
             )
             testScheduler.runCurrent()
 
@@ -646,7 +647,7 @@
                 headsUpManager,
                 keyguardNotifVisibilityProvider,
                 keyguardRepository,
-                keyguardTransitionRepository,
+                kosmos.keyguardTransitionInteractor,
                 KeyguardCoordinatorLogger(logcatLogBuffer()),
                 testScope.backgroundScope,
                 sectionHeaderVisibilityProvider,
@@ -706,4 +707,12 @@
                 )
             }
     }
+
+    companion object {
+        @JvmStatic
+        @Parameters(name = "{0}")
+        fun getParams(): List<FlagsParameterization> {
+            return FlagsParameterization.allCombinationsOf().andSceneContainer()
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java
index ed8843b..db829a2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/BatteryControllerTest.java
@@ -23,6 +23,8 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.staticMockMarker;
 import static com.android.settingslib.fuelgauge.BatterySaverLogging.SAVER_ENABLED_QS;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -308,6 +310,52 @@
         mBatteryController.fireBatteryLevelChanged();
     }
 
+    @Test
+    public void plugAndUnplugWhenInBatterySaver_stateUpdatedWithoutBatterySaverBroadcast() {
+        PowerSaveState state = new PowerSaveState.Builder()
+                .setBatterySaverEnabled(false)
+                .build();
+        when(mPowerManager.getPowerSaveState(PowerManager.ServiceType.AOD)).thenReturn(state);
+
+        // Set up on power save and not charging
+        when(mPowerManager.isPowerSaveMode()).thenReturn(true);
+        mBatteryController.onReceive(
+                getContext(), new Intent(PowerManager.ACTION_POWER_SAVE_MODE_CHANGED));
+        mBatteryController.onReceive(getContext(), createChargingIntent(false));
+
+        TestCallback callback = new TestCallback();
+        mBatteryController.addCallback(callback);
+
+        assertThat(callback.pluggedIn).isFalse();
+        assertThat(callback.powerSaverOn).isTrue();
+
+        // Plug in (battery saver turns off)
+        when(mPowerManager.isPowerSaveMode()).thenReturn(false);
+        mBatteryController.onReceive(getContext(), createChargingIntent(true));
+
+        assertThat(callback.pluggedIn).isTrue();
+        assertThat(callback.powerSaverOn).isFalse();
+
+        // Unplug (battery saver turns back on)
+        when(mPowerManager.isPowerSaveMode()).thenReturn(true);
+        mBatteryController.onReceive(getContext(), createChargingIntent(false));
+
+        assertThat(callback.pluggedIn).isFalse();
+        assertThat(callback.powerSaverOn).isTrue();
+    }
+
+    private Intent createChargingIntent(boolean charging) {
+        Intent intent = new Intent(Intent.ACTION_BATTERY_CHANGED);
+        if (charging) {
+            return intent
+                .putExtra(BatteryManager.EXTRA_STATUS, BatteryManager.BATTERY_STATUS_CHARGING)
+                .putExtra(BatteryManager.EXTRA_PLUGGED, BatteryManager.BATTERY_PLUGGED_AC);
+        } else {
+            return intent
+                .putExtra(BatteryManager.EXTRA_STATUS, BatteryManager.BATTERY_STATUS_DISCHARGING);
+        }
+    }
+
     private void setupIncompatibleCharging() {
         final List<UsbPort> usbPorts = new ArrayList<>();
         usbPorts.add(mUsbPort);
@@ -318,4 +366,19 @@
         when(mUsbPortStatus.getComplianceWarnings())
                 .thenReturn(new int[]{UsbPortStatus.COMPLIANCE_WARNING_DEBUG_ACCESSORY});
     }
+
+    private static class TestCallback
+        implements BatteryController.BatteryStateChangeCallback {
+        boolean pluggedIn = false;
+        boolean powerSaverOn = false;
+        @Override
+        public void onBatteryLevelChanged(int level, boolean pluggedIn, boolean charging) {
+            this.pluggedIn = pluggedIn;
+        }
+
+        @Override
+        public void onPowerSaveChanged(boolean isPowerSave) {
+            this.powerSaverOn = isPowerSave;
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/icons/AppCategoryIconProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/util/icons/AppCategoryIconProviderTest.kt
new file mode 100644
index 0000000..ef41b6ee
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/icons/AppCategoryIconProviderTest.kt
@@ -0,0 +1,154 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.util.icons
+
+import android.app.role.RoleManager.ROLE_ASSISTANT
+import android.content.ComponentName
+import android.content.Intent
+import android.content.Intent.CATEGORY_APP_BROWSER
+import android.content.Intent.CATEGORY_APP_CONTACTS
+import android.content.Intent.CATEGORY_APP_EMAIL
+import android.content.mockPackageManager
+import android.content.mockPackageManagerWrapper
+import android.content.pm.ActivityInfo
+import android.content.pm.ApplicationInfo
+import android.content.pm.PackageInfo
+import android.content.pm.ResolveInfo
+import android.graphics.drawable.Icon
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.assist.mockAssistManager
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class AppCategoryIconProviderTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val packageManagerWrapper = kosmos.mockPackageManagerWrapper
+    private val packageManager = kosmos.mockPackageManager
+    private val assistManager = kosmos.mockAssistManager
+    private val provider = kosmos.appCategoryIconProvider
+
+    @Before
+    fun setUp() {
+        whenever(packageManagerWrapper.resolveActivity(any<Intent>(), any<Int>())).thenAnswer {
+            invocation ->
+            val category = (invocation.arguments[0] as Intent).categories.first()
+            val categoryAppIcon =
+                categoryAppIcons.firstOrNull { it.category == category } ?: return@thenAnswer null
+            val activityInfo = ActivityInfo().also { it.packageName = categoryAppIcon.packageName }
+            return@thenAnswer ResolveInfo().also { it.activityInfo = activityInfo }
+        }
+        whenever(packageManager.getPackageInfo(any<String>(), any<Int>())).thenAnswer { invocation
+            ->
+            val packageName = invocation.arguments[0] as String
+            val categoryAppIcon =
+                categoryAppIcons.firstOrNull { it.packageName == packageName }
+                    ?: return@thenAnswer null
+            val applicationInfo =
+                ApplicationInfo().also {
+                    it.packageName = packageName
+                    it.icon = categoryAppIcon.iconResId
+                }
+            return@thenAnswer PackageInfo().also {
+                it.packageName = packageName
+                it.applicationInfo = applicationInfo
+            }
+        }
+    }
+
+    @Test
+    fun assistantAppIcon_defaultAssistantSet_returnsIcon() =
+        testScope.runTest {
+            whenever(assistManager.assistInfo)
+                .thenReturn(ComponentName(ASSISTANT_PACKAGE, ASSISTANT_CLASS))
+
+            val icon = provider.assistantAppIcon() as Icon
+
+            assertThat(icon.resPackage).isEqualTo(ASSISTANT_PACKAGE)
+            assertThat(icon.resId).isEqualTo(ASSISTANT_ICON_RES_ID)
+        }
+
+    @Test
+    fun assistantAppIcon_defaultAssistantNotSet_returnsNull() =
+        testScope.runTest {
+            whenever(assistManager.assistInfo).thenReturn(null)
+
+            assertThat(provider.assistantAppIcon()).isNull()
+        }
+
+    @Test
+    fun categoryAppIcon_returnsIconOfKnownBrowserApp() {
+        testScope.runTest {
+            val icon = provider.categoryAppIcon(CATEGORY_APP_BROWSER) as Icon
+
+            assertThat(icon.resPackage).isEqualTo(BROWSER_PACKAGE)
+            assertThat(icon.resId).isEqualTo(BROWSER_ICON_RES_ID)
+        }
+    }
+
+    @Test
+    fun categoryAppIcon_returnsIconOfKnownContactsApp() {
+        testScope.runTest {
+            val icon = provider.categoryAppIcon(CATEGORY_APP_CONTACTS) as Icon
+
+            assertThat(icon.resPackage).isEqualTo(CONTACTS_PACKAGE)
+            assertThat(icon.resId).isEqualTo(CONTACTS_ICON_RES_ID)
+        }
+    }
+
+    @Test
+    fun categoryAppIcon_noDefaultAppForCategoryEmail_returnsNull() {
+        testScope.runTest {
+            val icon = provider.categoryAppIcon(CATEGORY_APP_EMAIL)
+
+            assertThat(icon).isNull()
+        }
+    }
+
+    private companion object {
+        private const val ASSISTANT_PACKAGE = "the.assistant.app"
+        private const val ASSISTANT_CLASS = "the.assistant.app.class"
+        private const val ASSISTANT_ICON_RES_ID = 123
+
+        private const val BROWSER_PACKAGE = "com.test.browser"
+        private const val BROWSER_ICON_RES_ID = 1
+
+        private const val CONTACTS_PACKAGE = "app.test.contacts"
+        private const val CONTACTS_ICON_RES_ID = 234
+
+        private val categoryAppIcons =
+            listOf(
+                App(ROLE_ASSISTANT, ASSISTANT_PACKAGE, ASSISTANT_ICON_RES_ID),
+                App(CATEGORY_APP_BROWSER, BROWSER_PACKAGE, BROWSER_ICON_RES_ID),
+                App(CATEGORY_APP_CONTACTS, CONTACTS_PACKAGE, CONTACTS_ICON_RES_ID),
+            )
+    }
+
+    private class App(val category: String, val packageName: String, val iconResId: Int)
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerCollectorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerCollectorTest.kt
new file mode 100644
index 0000000..dd78e4a
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerCollectorTest.kt
@@ -0,0 +1,100 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.volume
+
+import android.media.IVolumeController
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.settingslib.media.data.repository.VolumeControllerEvent
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class VolumeControllerCollectorTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val eventsFlow = MutableStateFlow<VolumeControllerEvent?>(null)
+    private val underTest = VolumeControllerCollector(kosmos.applicationCoroutineScope)
+
+    private val volumeController = mock<IVolumeController> {}
+
+    @Test
+    fun volumeControllerEvent_volumeChanged_callsMethod() =
+        testEvent(VolumeControllerEvent.VolumeChanged(3, 0)) {
+            verify(volumeController) { 1 * { volumeController.volumeChanged(eq(3), eq(0)) } }
+        }
+
+    @Test
+    fun volumeControllerEvent_dismiss_callsMethod() =
+        testEvent(VolumeControllerEvent.Dismiss) {
+            verify(volumeController) { 1 * { volumeController.dismiss() } }
+        }
+
+    @Test
+    fun volumeControllerEvent_displayCsdWarning_callsMethod() =
+        testEvent(VolumeControllerEvent.DisplayCsdWarning(0, 1)) {
+            verify(volumeController) { 1 * { volumeController.displayCsdWarning(eq(0), eq(1)) } }
+        }
+
+    @Test
+    fun volumeControllerEvent_displaySafeVolumeWarning_callsMethod() =
+        testEvent(VolumeControllerEvent.DisplaySafeVolumeWarning(1)) {
+            verify(volumeController) { 1 * { volumeController.displaySafeVolumeWarning(eq(1)) } }
+        }
+
+    @Test
+    fun volumeControllerEvent_masterMuteChanged_callsMethod() =
+        testEvent(VolumeControllerEvent.MasterMuteChanged(1)) {
+            verify(volumeController) { 1 * { volumeController.masterMuteChanged(1) } }
+        }
+
+    @Test
+    fun volumeControllerEvent_setA11yMode_callsMethod() =
+        testEvent(VolumeControllerEvent.SetA11yMode(1)) {
+            verify(volumeController) { 1 * { volumeController.setA11yMode(1) } }
+        }
+
+    @Test
+    fun volumeControllerEvent_SetLayoutDirection_callsMethod() =
+        testEvent(VolumeControllerEvent.SetLayoutDirection(1)) {
+            verify(volumeController) { 1 * { volumeController.setLayoutDirection(eq(1)) } }
+        }
+
+    private fun testEvent(event: VolumeControllerEvent, verify: () -> Unit) =
+        kosmos.testScope.runTest {
+            underTest.collectToController(eventsFlow.filterNotNull(), volumeController)
+
+            eventsFlow.value = event
+            runCurrent()
+
+            verify()
+        }
+}
diff --git a/packages/SystemUI/tests/utils/src/android/content/PackageManagerKosmos.kt b/packages/SystemUI/tests/utils/src/android/content/PackageManagerKosmos.kt
index 8901314..9d7d916 100644
--- a/packages/SystemUI/tests/utils/src/android/content/PackageManagerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/android/content/PackageManagerKosmos.kt
@@ -17,6 +17,13 @@
 
 import android.content.pm.PackageManager
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.shared.system.PackageManagerWrapper
 import com.android.systemui.util.mockito.mock
 
-val Kosmos.packageManager by Kosmos.Fixture { mock<PackageManager>() }
+val Kosmos.mockPackageManager by Kosmos.Fixture { mock<PackageManager>() }
+
+var Kosmos.packageManager by Kosmos.Fixture { mockPackageManager }
+
+val Kosmos.mockPackageManagerWrapper by Kosmos.Fixture { mock<PackageManagerWrapper>() }
+
+var Kosmos.packageManagerWrapper by Kosmos.Fixture { mockPackageManagerWrapper }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/assist/AssistManagerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/assist/AssistManagerKosmos.kt
index b7d6f3a..22eb646 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/assist/AssistManagerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/assist/AssistManagerKosmos.kt
@@ -19,4 +19,6 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.util.mockito.mock
 
-var Kosmos.assistManager by Kosmos.Fixture { mock<AssistManager>() }
+val Kosmos.mockAssistManager by Kosmos.Fixture { mock<AssistManager>() }
+
+var Kosmos.assistManager by Kosmos.Fixture { mockAssistManager }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
index 55c803a..a1021f6 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyboard/shortcut/KeyboardShortcutHelperKosmos.kt
@@ -100,6 +100,7 @@
 val Kosmos.shortcutHelperViewModel by
     Kosmos.Fixture {
         ShortcutHelperViewModel(
+            applicationCoroutineScope,
             testDispatcher,
             shortcutHelperStateInteractor,
             shortcutHelperCategoriesInteractor
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/mediarouter/data/repository/FakeMediaRouterRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/mediarouter/data/repository/FakeMediaRouterRepository.kt
new file mode 100644
index 0000000..8aa7a03
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/mediarouter/data/repository/FakeMediaRouterRepository.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.mediarouter.data.repository
+
+import com.android.systemui.statusbar.policy.CastDevice
+import kotlinx.coroutines.flow.MutableStateFlow
+
+class FakeMediaRouterRepository : MediaRouterRepository {
+    override val castDevices: MutableStateFlow<List<CastDevice>> = MutableStateFlow(emptyList())
+
+    var lastStoppedDevice: CastDevice? = null
+        private set
+
+    override fun stopCasting(device: CastDevice) {
+        lastStoppedDevice = device
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepositoryKosmos.kt
new file mode 100644
index 0000000..eec9920
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/mediarouter/data/repository/MediaRouterRepositoryKosmos.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.mediarouter.data.repository
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.statusbar.policy.fakeCastController
+
+val Kosmos.realMediaRouterRepository by
+    Kosmos.Fixture {
+        MediaRouterRepositoryImpl(
+            scope = applicationCoroutineScope,
+            castController = fakeCastController,
+        )
+    }
+
+val Kosmos.fakeMediaRouterRepository by Kosmos.Fixture { FakeMediaRouterRepository() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/interactor/MediaRouterChipInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/interactor/MediaRouterChipInteractorKosmos.kt
new file mode 100644
index 0000000..cb18b68
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/domain/interactor/MediaRouterChipInteractorKosmos.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.chips.casttootherdevice.domain.interactor
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.mediarouter.data.repository.fakeMediaRouterRepository
+
+val Kosmos.mediaRouterChipInteractor by
+    Kosmos.Fixture {
+        MediaRouterChipInteractor(
+            scope = applicationCoroutineScope,
+            mediaRouterRepository = fakeMediaRouterRepository,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt
index 4baa8d0..a8de460 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelKosmos.kt
@@ -19,6 +19,7 @@
 import com.android.systemui.animation.mockDialogTransitionAnimator
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.statusbar.chips.casttootherdevice.domain.interactor.mediaRouterChipInteractor
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.mediaProjectionChipInteractor
 import com.android.systemui.statusbar.chips.mediaprojection.ui.view.endMediaProjectionDialogHelper
 import com.android.systemui.util.time.fakeSystemClock
@@ -28,6 +29,7 @@
         CastToOtherDeviceChipViewModel(
             scope = applicationCoroutineScope,
             mediaProjectionChipInteractor = mediaProjectionChipInteractor,
+            mediaRouterChipInteractor = mediaRouterChipInteractor,
             systemClock = fakeSystemClock,
             endMediaProjectionDialogHelper = endMediaProjectionDialogHelper,
             dialogTransitionAnimator = mockDialogTransitionAnimator,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/CastControllerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/CastControllerKosmos.kt
new file mode 100644
index 0000000..8e77437
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/CastControllerKosmos.kt
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.policy
+
+import com.android.systemui.kosmos.Kosmos
+
+val Kosmos.fakeCastController: FakeCastController by Kosmos.Fixture { FakeCastController() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/FakeCastController.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/FakeCastController.kt
new file mode 100644
index 0000000..2df0c7a5
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/FakeCastController.kt
@@ -0,0 +1,61 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.policy
+
+import java.io.PrintWriter
+
+class FakeCastController : CastController {
+    private var listeners = mutableListOf<CastController.Callback>()
+
+    private var castDevices = emptyList<CastDevice>()
+
+    var lastStoppedDevice: CastDevice? = null
+        private set
+
+    override fun addCallback(listener: CastController.Callback) {
+        listeners += listener
+    }
+
+    override fun removeCallback(listener: CastController.Callback) {
+        listeners -= listener
+    }
+
+    override fun getCastDevices(): List<CastDevice> {
+        return castDevices
+    }
+
+    fun setCastDevices(devices: List<CastDevice>) {
+        castDevices = devices
+        listeners.forEach { it.onCastDevicesChanged() }
+    }
+
+    override fun startCasting(device: CastDevice?) {}
+
+    override fun stopCasting(device: CastDevice?) {
+        lastStoppedDevice = device
+    }
+
+    override fun hasConnectedCastDevice(): Boolean {
+        return castDevices.any { it.state == CastDevice.CastState.Connected }
+    }
+
+    override fun dump(pw: PrintWriter, args: Array<out String>) {}
+
+    override fun setDiscovering(request: Boolean) {}
+
+    override fun setCurrentUserId(currentUserId: Int) {}
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/icons/AppCategoryIconProviderKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/icons/AppCategoryIconProviderKosmos.kt
new file mode 100644
index 0000000..7987185
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/icons/AppCategoryIconProviderKosmos.kt
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.util.icons
+
+import android.content.mockPackageManager
+import android.content.mockPackageManagerWrapper
+import com.android.systemui.assist.mockAssistManager
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testDispatcher
+
+var Kosmos.fakeAppCategoryIconProvider by Kosmos.Fixture { FakeAppCategoryIconProvider() }
+
+var Kosmos.appCategoryIconProvider: AppCategoryIconProvider by
+    Kosmos.Fixture {
+        AppCategoryIconProviderImpl(
+            testDispatcher,
+            mockAssistManager,
+            mockPackageManager,
+            mockPackageManagerWrapper
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/icons/FakeAppCategoryIconProvider.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/icons/FakeAppCategoryIconProvider.kt
new file mode 100644
index 0000000..3e7bf21
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/icons/FakeAppCategoryIconProvider.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.util.icons
+
+import android.app.role.RoleManager.ROLE_ASSISTANT
+import android.graphics.drawable.Icon
+
+class FakeAppCategoryIconProvider : AppCategoryIconProvider {
+
+    private val installedApps = mutableMapOf<String, App>()
+
+    fun installCategoryApp(category: String, packageName: String, iconResId: Int) {
+        installedApps[category] = App(packageName, iconResId)
+    }
+
+    fun installAssistantApp(packageName: String, iconResId: Int) {
+        installedApps[ROLE_ASSISTANT] = App(packageName, iconResId)
+    }
+
+    override suspend fun assistantAppIcon() = categoryAppIcon(ROLE_ASSISTANT)
+
+    override suspend fun categoryAppIcon(category: String): Icon? {
+        val app = installedApps[category] ?: return null
+        return Icon.createWithResource(app.packageName, app.iconResId)
+    }
+
+    private class App(val packageName: String, val iconResId: Int)
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/LeakCheckedTest.java b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/LeakCheckedTest.java
index 5d21ddd..372a7c7 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/LeakCheckedTest.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/LeakCheckedTest.java
@@ -109,7 +109,7 @@
                 } else if (cls == ZenModeController.class) {
                     obj = new FakeZenModeController(this);
                 } else if (cls == CastController.class) {
-                    obj = new FakeCastController(this);
+                    obj = new LeakCheckerCastController(this);
                 } else if (cls == HotspotController.class) {
                     obj = new FakeHotspotController(this);
                 } else if (cls == FlashlightController.class) {
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeCastController.java b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/LeakCheckerCastController.java
similarity index 67%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeCastController.java
rename to packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/LeakCheckerCastController.java
index 5fae38f..2249bc0 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeCastController.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/LeakCheckerCastController.java
@@ -1,15 +1,17 @@
 /*
  * Copyright (C) 2016 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
+ * 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.
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 package com.android.systemui.utils.leaks;
@@ -23,8 +25,8 @@
 import java.util.ArrayList;
 import java.util.List;
 
-public class FakeCastController extends BaseLeakChecker<Callback> implements CastController {
-    public FakeCastController(LeakCheck test) {
+public class LeakCheckerCastController extends BaseLeakChecker<Callback> implements CastController {
+    public LeakCheckerCastController(LeakCheck test) {
         super(test, "cast");
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt
new file mode 100644
index 0000000..d60f14c
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.volume
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+
+val Kosmos.volumeControllerCollector by
+    Kosmos.Fixture { VolumeControllerCollector(applicationCoroutineScope) }
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index fbb6ccf..b0d734d 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -70,7 +70,6 @@
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.ContentProvider;
-import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
@@ -80,7 +79,6 @@
 import android.content.pm.ResolveInfo;
 import android.content.pm.ServiceInfo;
 import android.content.res.Resources;
-import android.database.ContentObserver;
 import android.hardware.input.InputManager;
 import android.inputmethodservice.InputMethodService;
 import android.media.AudioManagerInternal;
@@ -195,7 +193,6 @@
 import java.security.InvalidParameterException;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
@@ -349,8 +346,6 @@
     @GuardedBy("ImfLock.class")
     private UserDataRepository mUserDataRepository;
 
-    @MultiUserUnawareField
-    final SettingsObserver mSettingsObserver;
     final WindowManagerInternal mWindowManagerInternal;
     private final ActivityManagerInternal mActivityManagerInternal;
     final PackageManagerInternal mPackageManagerInternal;
@@ -570,82 +565,52 @@
     @NonNull
     private final ImeTrackerService mImeTrackerService;
 
-    class SettingsObserver extends ContentObserver {
-
-        /**
-         * <em>This constructor must be called within the lock.</em>
-         */
-        SettingsObserver(Handler handler) {
-            super(handler);
+    @GuardedBy("ImfLock.class")
+    private void onSecureSettingsChangedLocked(@NonNull String key, @UserIdInt int userId) {
+        if (!mConcurrentMultiUserModeEnabled && userId != mCurrentUserId) {
+            return;
         }
-
-        void registerContentObserverForAllUsers() {
-            ContentResolver resolver = mContext.getContentResolver();
-            resolver.registerContentObserverAsUser(Settings.Secure.getUriFor(
-                    Settings.Secure.DEFAULT_INPUT_METHOD), false, this, UserHandle.ALL);
-            resolver.registerContentObserverAsUser(Settings.Secure.getUriFor(
-                    Settings.Secure.ENABLED_INPUT_METHODS), false, this, UserHandle.ALL);
-            resolver.registerContentObserverAsUser(Settings.Secure.getUriFor(
-                    Settings.Secure.SELECTED_INPUT_METHOD_SUBTYPE), false, this, UserHandle.ALL);
-            resolver.registerContentObserverAsUser(Settings.Secure.getUriFor(
-                    Settings.Secure.SHOW_IME_WITH_HARD_KEYBOARD), false, this, UserHandle.ALL);
-            resolver.registerContentObserverAsUser(Settings.Secure.getUriFor(
-                    Settings.Secure.ACCESSIBILITY_SOFT_KEYBOARD_MODE), false, this, UserHandle.ALL);
-            resolver.registerContentObserverAsUser(Settings.Secure.getUriFor(
-                    STYLUS_HANDWRITING_ENABLED), false, this, UserHandle.ALL);
-        }
-
-        @Override
-        public void onChange(boolean selfChange, @NonNull Collection<Uri> uris, int flags,
-                @UserIdInt int userId) {
-            uris.forEach(uri -> onChangeInternal(uri, userId));
-        }
-
-        private void onChangeInternal(@NonNull Uri uri, @UserIdInt int userId) {
-            final Uri showImeUri = Settings.Secure.getUriFor(
-                    Settings.Secure.SHOW_IME_WITH_HARD_KEYBOARD);
-            final Uri accessibilityRequestingNoImeUri = Settings.Secure.getUriFor(
-                    Settings.Secure.ACCESSIBILITY_SOFT_KEYBOARD_MODE);
-            final Uri stylusHandwritingEnabledUri = Settings.Secure.getUriFor(
-                    STYLUS_HANDWRITING_ENABLED);
-            synchronized (ImfLock.class) {
-                if (!mConcurrentMultiUserModeEnabled && mCurrentUserId != userId) {
-                    return;
+        switch (key) {
+            case Settings.Secure.SHOW_IME_WITH_HARD_KEYBOARD: {
+                mMenuController.updateKeyboardFromSettingsLocked();
+                break;
+            }
+            case Settings.Secure.ACCESSIBILITY_SOFT_KEYBOARD_MODE: {
+                final int accessibilitySoftKeyboardSetting = Settings.Secure.getIntForUser(
+                        mContext.getContentResolver(),
+                        Settings.Secure.ACCESSIBILITY_SOFT_KEYBOARD_MODE, 0, userId);
+                mVisibilityStateComputer.getImePolicy().setA11yRequestNoSoftKeyboard(
+                        accessibilitySoftKeyboardSetting);
+                final var userData = getUserData(userId);
+                if (mVisibilityStateComputer.getImePolicy().isA11yRequestNoSoftKeyboard()) {
+                    hideCurrentInputLocked(userData.mImeBindingState.mFocusedWindow,
+                            0 /* flags */, SoftInputShowHideReason.HIDE_SETTINGS_ON_CHANGE, userId);
+                } else if (isShowRequestedForCurrentWindow(userId)) {
+                    showCurrentInputLocked(userData.mImeBindingState.mFocusedWindow,
+                            InputMethodManager.SHOW_IMPLICIT,
+                            SoftInputShowHideReason.SHOW_SETTINGS_ON_CHANGE, userId);
                 }
-
-                if (showImeUri.equals(uri)) {
-                    mMenuController.updateKeyboardFromSettingsLocked();
-                } else if (accessibilityRequestingNoImeUri.equals(uri)) {
-                    final int accessibilitySoftKeyboardSetting = Settings.Secure.getIntForUser(
-                            mContext.getContentResolver(),
-                            Settings.Secure.ACCESSIBILITY_SOFT_KEYBOARD_MODE, 0, userId);
-                    mVisibilityStateComputer.getImePolicy().setA11yRequestNoSoftKeyboard(
-                            accessibilitySoftKeyboardSetting);
-                    final var userData = getUserData(userId);
-                    if (mVisibilityStateComputer.getImePolicy().isA11yRequestNoSoftKeyboard()) {
-                        hideCurrentInputLocked(userData.mImeBindingState.mFocusedWindow,
-                                0 /* flags */, SoftInputShowHideReason.HIDE_SETTINGS_ON_CHANGE,
-                                userId);
-                    } else if (isShowRequestedForCurrentWindow(userId)) {
-                        showCurrentInputLocked(userData.mImeBindingState.mFocusedWindow,
-                                InputMethodManager.SHOW_IMPLICIT,
-                                SoftInputShowHideReason.SHOW_SETTINGS_ON_CHANGE, userId);
-                    }
-                } else if (stylusHandwritingEnabledUri.equals(uri)) {
-                    InputMethodManager.invalidateLocalStylusHandwritingAvailabilityCaches();
-                    InputMethodManager
-                            .invalidateLocalConnectionlessStylusHandwritingAvailabilityCaches();
-                } else {
-                    boolean enabledChanged = false;
-                    String newEnabled = InputMethodSettingsRepository.get(userId)
-                            .getEnabledInputMethodsStr();
-                    final var userData = getUserData(userId);
-                    if (!userData.mLastEnabledInputMethodsStr.equals(newEnabled)) {
-                        userData.mLastEnabledInputMethodsStr = newEnabled;
-                        enabledChanged = true;
-                    }
-                    updateInputMethodsFromSettingsLocked(enabledChanged, userId);
+                break;
+            }
+            case STYLUS_HANDWRITING_ENABLED: {
+                InputMethodManager.invalidateLocalStylusHandwritingAvailabilityCaches();
+                InputMethodManager
+                        .invalidateLocalConnectionlessStylusHandwritingAvailabilityCaches();
+                break;
+            }
+            case Settings.Secure.DEFAULT_INPUT_METHOD:
+            case Settings.Secure.ENABLED_INPUT_METHODS:
+            case Settings.Secure.SELECTED_INPUT_METHOD_SUBTYPE: {
+                boolean enabledChanged = false;
+                String newEnabled = InputMethodSettingsRepository.get(userId)
+                        .getEnabledInputMethodsStr();
+                final var userData = getUserData(userId);
+                if (!userData.mLastEnabledInputMethodsStr.equals(newEnabled)) {
+                    userData.mLastEnabledInputMethodsStr = newEnabled;
+                    enabledChanged = true;
                 }
+                updateInputMethodsFromSettingsLocked(enabledChanged, userId);
+                break;
             }
         }
     }
@@ -1118,8 +1083,6 @@
             }
             SystemLocaleWrapper.onStart(context, this::onActionLocaleChanged, mHandler);
             mImeTrackerService = new ImeTrackerService(mHandler);
-            // Note: SettingsObserver doesn't register observers in its constructor.
-            mSettingsObserver = new SettingsObserver(mHandler);
             mWindowManagerInternal = LocalServices.getService(WindowManagerInternal.class);
             mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
             mPackageManagerInternal = LocalServices.getService(PackageManagerInternal.class);
@@ -1389,7 +1352,19 @@
                 }, "Lazily initialize IMMS#mImeDrawsImeNavBarRes");
 
                 mMyPackageMonitor.register(mContext, UserHandle.ALL, mIoHandler);
-                mSettingsObserver.registerContentObserverForAllUsers();
+                SecureSettingsChangeCallback.register(mHandler, mContext.getContentResolver(),
+                        new String[] {
+                                Settings.Secure.ACCESSIBILITY_SOFT_KEYBOARD_MODE,
+                                Settings.Secure.DEFAULT_INPUT_METHOD,
+                                Settings.Secure.ENABLED_INPUT_METHODS,
+                                Settings.Secure.SELECTED_INPUT_METHOD_SUBTYPE,
+                                Settings.Secure.SHOW_IME_WITH_HARD_KEYBOARD,
+                                Settings.Secure.STYLUS_HANDWRITING_ENABLED,
+                        }, (key, flags, userId) -> {
+                            synchronized (ImfLock.class) {
+                                onSecureSettingsChangedLocked(key, userId);
+                            }
+                        });
 
                 final IntentFilter broadcastFilterForAllUsers = new IntentFilter();
                 broadcastFilterForAllUsers.addAction(Intent.ACTION_CLOSE_SYSTEM_DIALOGS);
diff --git a/services/core/java/com/android/server/inputmethod/SecureSettingsChangeCallback.java b/services/core/java/com/android/server/inputmethod/SecureSettingsChangeCallback.java
new file mode 100644
index 0000000..328d7c6
--- /dev/null
+++ b/services/core/java/com/android/server/inputmethod/SecureSettingsChangeCallback.java
@@ -0,0 +1,78 @@
+/*
+ * 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.inputmethod;
+
+import android.annotation.NonNull;
+import android.annotation.UserIdInt;
+import android.content.ContentResolver;
+import android.database.ContentObserver;
+import android.net.Uri;
+import android.os.Handler;
+import android.os.UserHandle;
+import android.provider.Settings;
+import android.util.ArrayMap;
+
+import java.util.Collection;
+
+/**
+ * A wrapper interface to monitor the given set of {@link Settings.Secure}.
+ */
+@FunctionalInterface
+interface SecureSettingsChangeCallback {
+    /**
+     * Called back when the value associated with {@code key} is updated.
+     *
+     * @param key a key defined in {@link Settings.Secure}
+     * @param flags flags defined in {@link ContentResolver.NotifyFlags}
+     * @param userId the user ID with which the value is associated
+     */
+    void onChange(@NonNull String key, @ContentResolver.NotifyFlags int flags,
+            @UserIdInt int userId);
+
+    /**
+     * Registers {@link SecureSettingsChangeCallback} to the given set of {@link Settings.Secure}.
+     *
+     * @param handler  {@link Handler} to be used to call back {@link #onChange(String, int, int)}
+     * @param resolver {@link ContentResolver} with which {@link Settings.Secure} will be retrieved
+     * @param keys     A set of {@link Settings.Secure} to be monitored
+     * @param callback {@link SecureSettingsChangeCallback} to be called back
+     */
+    @NonNull
+    static void register(@NonNull Handler handler, @NonNull ContentResolver resolver,
+            @NonNull String[] keys, @NonNull SecureSettingsChangeCallback callback) {
+        final ArrayMap<Uri, String> uriMapper = new ArrayMap<>();
+        for (String key : keys) {
+            uriMapper.put(Settings.Secure.getUriFor(key), key);
+        }
+        final ContentObserver observer = new ContentObserver(handler) {
+            @Override
+            public void onChange(boolean selfChange, @NonNull Collection<Uri> uris, int flags,
+                    @UserIdInt int userId) {
+                uris.forEach(uri -> {
+                    final String key = uriMapper.get(uri);
+                    if (key != null) {
+                        callback.onChange(key, flags, userId);
+                    }
+                });
+            }
+        };
+        for (Uri uri : uriMapper.keySet()) {
+            resolver.registerContentObserverAsUser(uri, false /* notifyForDescendants */, observer,
+                    UserHandle.ALL);
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/notification/PreferencesHelper.java b/services/core/java/com/android/server/notification/PreferencesHelper.java
index 3a0c1d0..c09504f 100644
--- a/services/core/java/com/android/server/notification/PreferencesHelper.java
+++ b/services/core/java/com/android/server/notification/PreferencesHelper.java
@@ -196,9 +196,12 @@
         int USER_LOCKED_BUBBLE = 0x00000002;
     }
 
+    private final Object mLock = new Object();
     // pkg|uid => PackagePreferences
+    @GuardedBy("mLock")
     private final ArrayMap<String, PackagePreferences> mPackagePreferences = new ArrayMap<>();
     // pkg|userId => PackagePreferences
+    @GuardedBy("mLock")
     private final ArrayMap<String, PackagePreferences> mRestoredWithoutUids = new ArrayMap<>();
 
     private final Context mContext;
@@ -270,7 +273,7 @@
                     Settings.Global.REVIEW_PERMISSIONS_NOTIFICATION_STATE,
                     NotificationManagerService.REVIEW_NOTIF_STATE_SHOULD_SHOW);
         }
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             while ((type = parser.next()) != XmlPullParser.END_DOCUMENT) {
                 tag = parser.getName();
                 if (type == XmlPullParser.END_TAG && TAG_RANKING.equals(tag)) {
@@ -492,6 +495,7 @@
                 DEFAULT_BUBBLE_PREFERENCE, mClock.millis());
     }
 
+    @GuardedBy("mLock")
     private PackagePreferences getOrCreatePackagePreferencesLocked(String pkg,
             @UserIdInt int userId, int uid, int importance, int priority, int visibility,
             boolean showBadge, int bubblePreference, long creationTime) {
@@ -661,7 +665,7 @@
             notifPermissions = mPermissionHelper.getNotificationPermissionValues(userId);
         }
 
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             final int N = mPackagePreferences.size();
             for (int i = 0; i < N; i++) {
                 final PackagePreferences r = mPackagePreferences.valueAt(i);
@@ -670,11 +674,10 @@
                 }
                 writePackageXml(r, out, notifPermissions, forBackup);
             }
-        }
-        if (Flags.persistIncompleteRestoreData() && !forBackup) {
-            synchronized (mRestoredWithoutUids) {
-                final int N = mRestoredWithoutUids.size();
-                for (int i = 0; i < N; i++) {
+
+            if (Flags.persistIncompleteRestoreData() && !forBackup) {
+                final int M = mRestoredWithoutUids.size();
+                for (int i = 0; i < M; i++) {
                     final PackagePreferences r = mRestoredWithoutUids.valueAt(i);
                     writePackageXml(r, out, notifPermissions, false);
                 }
@@ -777,7 +780,7 @@
      */
     public void setBubblesAllowed(String pkg, int uid, int bubblePreference) {
         boolean changed;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences p = getOrCreatePackagePreferencesLocked(pkg, uid);
             changed = p.bubblePreference != bubblePreference;
             p.bubblePreference = bubblePreference;
@@ -797,20 +800,20 @@
      */
     @Override
     public int getBubblePreference(String pkg, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             return getOrCreatePackagePreferencesLocked(pkg, uid).bubblePreference;
         }
     }
 
     public int getAppLockedFields(String pkg, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             return getOrCreatePackagePreferencesLocked(pkg, uid).lockedAppFields;
         }
     }
 
     @Override
     public boolean canShowBadge(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             return getOrCreatePackagePreferencesLocked(packageName, uid).showBadge;
         }
     }
@@ -818,7 +821,7 @@
     @Override
     public void setShowBadge(String packageName, int uid, boolean showBadge) {
         boolean changed = false;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences pkgPrefs = getOrCreatePackagePreferencesLocked(packageName, uid);
             if (pkgPrefs.showBadge != showBadge) {
                 pkgPrefs.showBadge = showBadge;
@@ -831,28 +834,28 @@
     }
 
     public boolean isInInvalidMsgState(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             return r.hasSentInvalidMessage && !r.hasSentValidMessage;
         }
     }
 
     public boolean hasUserDemotedInvalidMsgApp(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             return isInInvalidMsgState(packageName, uid) ? r.userDemotedMsgApp : false;
         }
     }
 
     public void setInvalidMsgAppDemoted(String packageName, int uid, boolean isDemoted) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             r.userDemotedMsgApp = isDemoted;
         }
     }
 
     public boolean setInvalidMessageSent(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             boolean valueChanged = r.hasSentInvalidMessage == false;
             r.hasSentInvalidMessage = true;
@@ -862,7 +865,7 @@
     }
 
     public boolean setValidMessageSent(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             boolean valueChanged = r.hasSentValidMessage == false;
             r.hasSentValidMessage = true;
@@ -873,7 +876,7 @@
 
     @VisibleForTesting
     boolean hasSentInvalidMsg(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             return r.hasSentInvalidMessage;
         }
@@ -881,7 +884,7 @@
 
     @VisibleForTesting
     boolean hasSentValidMsg(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             return r.hasSentValidMessage;
         }
@@ -889,7 +892,7 @@
 
     @VisibleForTesting
     boolean didUserEverDemoteInvalidMsgApp(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             return r.userDemotedMsgApp;
         }
@@ -897,7 +900,7 @@
 
     /** Sets whether this package has sent a notification with valid bubble metadata. */
     public boolean setValidBubbleSent(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             boolean valueChanged = !r.hasSentValidBubble;
             r.hasSentValidBubble = true;
@@ -906,14 +909,14 @@
     }
 
     boolean hasSentValidBubble(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             return r.hasSentValidBubble;
         }
     }
 
     boolean isImportanceLocked(String pkg, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             return r.fixedImportance || r.defaultAppLockedImportance;
         }
@@ -924,7 +927,7 @@
         if (groupId == null) {
             return false;
         }
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(packageName, uid);
             NotificationChannelGroup group = r.groups.get(groupId);
             if (group == null) {
@@ -935,13 +938,13 @@
     }
 
     int getPackagePriority(String pkg, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             return getOrCreatePackagePreferencesLocked(pkg, uid).priority;
         }
     }
 
     int getPackageVisibility(String pkg, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             return getOrCreatePackagePreferencesLocked(pkg, uid).visibility;
         }
     }
@@ -956,7 +959,7 @@
             throw new IllegalArgumentException("group.getName() can't be empty");
         }
         boolean needsDndChange = false;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 throw new IllegalArgumentException("Invalid package");
@@ -1010,7 +1013,7 @@
                 && channel.getImportance() <= IMPORTANCE_MAX, "Invalid importance level");
 
         boolean needsPolicyFileChange = false, wasUndeleted = false, needsDndChange = false;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 throw new IllegalArgumentException("Invalid package");
@@ -1154,7 +1157,7 @@
 
     void unlockNotificationChannelImportance(String pkg, int uid, String updatedChannelId) {
         Objects.requireNonNull(updatedChannelId);
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 throw new IllegalArgumentException("Invalid package");
@@ -1176,7 +1179,7 @@
         Objects.requireNonNull(updatedChannel.getId());
         boolean changed = false;
         boolean needsDndChange = false;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 throw new IllegalArgumentException("Invalid package");
@@ -1351,7 +1354,7 @@
             String channelId, String conversationId, boolean returnParentIfNoConversationChannel,
             boolean includeDeleted) {
         Preconditions.checkNotNull(pkg);
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return null;
@@ -1392,7 +1395,7 @@
         Preconditions.checkNotNull(pkg);
         Preconditions.checkNotNull(conversationId);
         List<NotificationChannel> channels = new ArrayList<>();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return channels;
@@ -1412,7 +1415,7 @@
             int callingUid, boolean fromSystemOrSystemUi) {
         boolean deletedChannel = false;
         boolean channelBypassedDnd = false;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return false;
@@ -1448,7 +1451,7 @@
     public void permanentlyDeleteNotificationChannel(String pkg, int uid, String channelId) {
         Objects.requireNonNull(pkg);
         Objects.requireNonNull(channelId);
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return;
@@ -1460,7 +1463,7 @@
     @Override
     public void permanentlyDeleteNotificationChannels(String pkg, int uid) {
         Objects.requireNonNull(pkg);
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return;
@@ -1491,7 +1494,7 @@
                 boolean fixed = mPermissionHelper.isPermissionFixed(
                         pi.packageName, user.getUserHandle().getIdentifier());
                 if (fixed) {
-                    synchronized (mPackagePreferences) {
+                    synchronized (mLock) {
                         PackagePreferences p = getOrCreatePackagePreferencesLocked(
                                 pi.packageName, pi.applicationInfo.uid);
                         p.fixedImportance = true;
@@ -1506,7 +1509,7 @@
 
     public void updateDefaultApps(int userId, ArraySet<String> toRemove,
             ArraySet<Pair<String, Integer>> toAdd) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             for (PackagePreferences p : mPackagePreferences.values()) {
                 if (userId == UserHandle.getUserId(p.uid)) {
                     if (toRemove != null && toRemove.contains(p.pkg)) {
@@ -1536,7 +1539,7 @@
     public NotificationChannelGroup getNotificationChannelGroupWithChannels(String pkg,
             int uid, String groupId, boolean includeDeleted) {
         Objects.requireNonNull(pkg);
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null || groupId == null || !r.groups.containsKey(groupId)) {
                 return null;
@@ -1559,7 +1562,7 @@
     public NotificationChannelGroup getNotificationChannelGroup(String groupId, String pkg,
             int uid) {
         Objects.requireNonNull(pkg);
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return null;
@@ -1573,7 +1576,7 @@
             boolean includeBlocked, Set<String> activeChannelFilter) {
         Objects.requireNonNull(pkg);
         Map<String, NotificationChannelGroup> groups = new ArrayMap<>();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return ParceledListSlice.emptyList();
@@ -1624,7 +1627,7 @@
             String groupId, int callingUid, boolean fromSystemOrSystemUi) {
         List<NotificationChannel> deletedChannels = new ArrayList<>();
         boolean groupBypassedDnd = false;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null || TextUtils.isEmpty(groupId)) {
                 return deletedChannels;
@@ -1656,7 +1659,7 @@
     public Collection<NotificationChannelGroup> getNotificationChannelGroups(String pkg,
             int uid) {
         List<NotificationChannelGroup> groups = new ArrayList<>();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return groups;
@@ -1667,7 +1670,7 @@
     }
 
     public NotificationChannelGroup getGroupForChannel(String pkg, int uid, String channelId) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences p = getPackagePreferencesLocked(pkg, uid);
             if (p != null) {
                 NotificationChannel nc = p.channels.get(channelId);
@@ -1686,7 +1689,7 @@
 
     public ArrayList<ConversationChannelWrapper> getConversations(IntArray userIds,
             boolean onlyImportant) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             ArrayList<ConversationChannelWrapper> conversations = new ArrayList<>();
             for (PackagePreferences p : mPackagePreferences.values()) {
                 if (userIds.binarySearch(UserHandle.getUserId(p.uid)) >= 0) {
@@ -1730,7 +1733,7 @@
 
     public ArrayList<ConversationChannelWrapper> getConversations(String pkg, int uid) {
         Objects.requireNonNull(pkg);
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return new ArrayList<>();
@@ -1772,7 +1775,7 @@
     public @NonNull List<String> deleteConversations(String pkg, int uid,
             Set<String> conversationIds, int callingUid, boolean fromSystemOrSystemUi) {
         List<String> deletedChannelIds = new ArrayList<>();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return deletedChannelIds;
@@ -1805,7 +1808,7 @@
             boolean includeDeleted) {
         Objects.requireNonNull(pkg);
         List<NotificationChannel> channels = new ArrayList<>();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return ParceledListSlice.emptyList();
@@ -1827,7 +1830,7 @@
     public ParceledListSlice<NotificationChannel> getNotificationChannelsBypassingDnd(String pkg,
             int uid) {
         List<NotificationChannel> channels = new ArrayList<>();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             final PackagePreferences r = mPackagePreferences.get(
                     packagePreferencesKey(pkg, uid));
             if (r != null) {
@@ -1848,7 +1851,7 @@
      * upgrades.
      */
     public boolean onlyHasDefaultChannel(String pkg, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getOrCreatePackagePreferencesLocked(pkg, uid);
             if (r.channels.size() == (notificationClassification() ? 5 : 1)
                     && r.channels.containsKey(NotificationChannel.DEFAULT_CHANNEL_ID)) {
@@ -1861,7 +1864,7 @@
     public int getDeletedChannelCount(String pkg, int uid) {
         Objects.requireNonNull(pkg);
         int deletedCount = 0;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return deletedCount;
@@ -1880,7 +1883,7 @@
     public int getBlockedChannelCount(String pkg, int uid) {
         Objects.requireNonNull(pkg);
         int blockedCount = 0;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences r = getPackagePreferencesLocked(pkg, uid);
             if (r == null) {
                 return blockedCount;
@@ -1923,7 +1926,7 @@
         ArraySet<Pair<String, Integer>> candidatePkgs = new ArraySet<>();
 
         final IntArray currentUserIds = mUserProfiles.getCurrentProfileIds();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             final int numPackagePreferences = mPackagePreferences.size();
             for (int i = 0; i < numPackagePreferences; i++) {
                 final PackagePreferences r = mPackagePreferences.valueAt(i);
@@ -1992,7 +1995,7 @@
      * considered for sentiment adjustments (and thus never show a blocking helper).
      */
     public void setAppImportanceLocked(String packageName, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences prefs = getOrCreatePackagePreferencesLocked(packageName, uid);
             if ((prefs.lockedAppFields & LockableAppFields.USER_LOCKED_IMPORTANCE) != 0) {
                 return;
@@ -2008,7 +2011,7 @@
      * Returns the delegate for a given package, if it's allowed by the package and the user.
      */
     public @Nullable String getNotificationDelegate(String sourcePkg, int sourceUid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences prefs = getPackagePreferencesLocked(sourcePkg, sourceUid);
 
             if (prefs == null || prefs.delegate == null) {
@@ -2026,7 +2029,7 @@
      */
     public void setNotificationDelegate(String sourcePkg, int sourceUid,
             String delegatePkg, int delegateUid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences prefs = getOrCreatePackagePreferencesLocked(sourcePkg, sourceUid);
             prefs.delegate = new Delegate(delegatePkg, delegateUid, true);
         }
@@ -2036,7 +2039,7 @@
      * Used by an app to turn off its notification delegate.
      */
     public void revokeNotificationDelegate(String sourcePkg, int sourceUid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences prefs = getPackagePreferencesLocked(sourcePkg, sourceUid);
             if (prefs != null && prefs.delegate != null) {
                 prefs.delegate.mEnabled = false;
@@ -2050,7 +2053,7 @@
      */
     public boolean isDelegateAllowed(String sourcePkg, int sourceUid,
             String potentialDelegatePkg, int potentialDelegateUid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences prefs = getPackagePreferencesLocked(sourcePkg, sourceUid);
 
             return prefs != null && prefs.isValidDelegate(potentialDelegatePkg,
@@ -2096,24 +2099,25 @@
         pw.println("per-package config version: " + XML_VERSION);
 
         pw.println("PackagePreferences:");
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             dumpPackagePreferencesLocked(pw, prefix, filter, mPackagePreferences, pkgPermissions);
+            pw.println("Restored without uid:");
+            dumpPackagePreferencesLocked(pw, prefix, filter, mRestoredWithoutUids, null);
         }
-        pw.println("Restored without uid:");
-        dumpPackagePreferencesLocked(pw, prefix, filter, mRestoredWithoutUids, null);
     }
 
     public void dump(ProtoOutputStream proto,
             @NonNull NotificationManagerService.DumpFilter filter,
             ArrayMap<Pair<Integer, String>, Pair<Boolean, Boolean>> pkgPermissions) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             dumpPackagePreferencesLocked(proto, RankingHelperProto.RECORDS, filter,
                     mPackagePreferences, pkgPermissions);
+            dumpPackagePreferencesLocked(proto, RankingHelperProto.RECORDS_RESTORED_WITHOUT_UID,
+                    filter, mRestoredWithoutUids, null);
         }
-        dumpPackagePreferencesLocked(proto, RankingHelperProto.RECORDS_RESTORED_WITHOUT_UID, filter,
-                mRestoredWithoutUids, null);
     }
 
+    @GuardedBy("mLock")
     private void dumpPackagePreferencesLocked(PrintWriter pw, String prefix,
             @NonNull NotificationManagerService.DumpFilter filter,
             ArrayMap<String, PackagePreferences> packagePreferences,
@@ -2298,7 +2302,7 @@
             pkgsWithPermissionsToHandle = pkgPermissions.keySet();
         }
         int pulledEvents = 0;
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             for (int i = 0; i < mPackagePreferences.size(); i++) {
                 if (pulledEvents > NOTIFICATION_PREFERENCES_PULL_LIMIT) {
                     break;
@@ -2378,7 +2382,7 @@
      * {@link StatsEvent}.
      */
     public void pullPackageChannelPreferencesStats(List<StatsEvent> events) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             int totalChannelsPulled = 0;
             for (int i = 0; i < mPackagePreferences.size(); i++) {
                 if (totalChannelsPulled > NOTIFICATION_CHANNEL_PULL_LIMIT) {
@@ -2414,7 +2418,7 @@
      * {@link StatsEvent}.
      */
     public void pullPackageChannelGroupPreferencesStats(List<StatsEvent> events) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             int totalGroupsPulled = 0;
             for (int i = 0; i < mPackagePreferences.size(); i++) {
                 if (totalGroupsPulled > NOTIFICATION_CHANNEL_GROUP_PULL_LIMIT) {
@@ -2443,10 +2447,12 @@
             ArrayMap<Pair<Integer, String>, Pair<Boolean, Boolean>> pkgPermissions) {
         JSONObject ranking = new JSONObject();
         JSONArray PackagePreferencess = new JSONArray();
-        try {
-            ranking.put("noUid", mRestoredWithoutUids.size());
-        } catch (JSONException e) {
-            // pass
+        synchronized (mLock) {
+            try {
+                ranking.put("noUid", mRestoredWithoutUids.size());
+            } catch (JSONException e) {
+                // pass
+            }
         }
 
         // Track data that we've handled from the permissions-based list
@@ -2455,7 +2461,7 @@
             pkgsWithPermissionsToHandle = pkgPermissions.keySet();
         }
 
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             final int N = mPackagePreferences.size();
             for (int i = 0; i < N; i++) {
                 final PackagePreferences r = mPackagePreferences.valueAt(i);
@@ -2561,7 +2567,7 @@
     }
 
     public Map<Integer, String> getPackageBans() {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             final int N = mPackagePreferences.size();
             ArrayMap<Integer, String> packageBans = new ArrayMap<>(N);
             for (int i = 0; i < N; i++) {
@@ -2620,7 +2626,7 @@
 
     private Map<String, Integer> getPackageChannels() {
         ArrayMap<String, Integer> packageChannels = new ArrayMap<>();
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             for (int i = 0; i < mPackagePreferences.size(); i++) {
                 final PackagePreferences r = mPackagePreferences.valueAt(i);
                 int channelCount = 0;
@@ -2636,7 +2642,7 @@
     }
 
     public void onUserRemoved(int userId) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             int N = mPackagePreferences.size();
             for (int i = N - 1; i >= 0; i--) {
                 PackagePreferences PackagePreferences = mPackagePreferences.valueAt(i);
@@ -2648,7 +2654,7 @@
     }
 
     protected void onLocaleChanged(Context context, int userId) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             int N = mPackagePreferences.size();
             for (int i = 0; i < N; i++) {
                 PackagePreferences PackagePreferences = mPackagePreferences.valueAt(i);
@@ -2678,22 +2684,24 @@
             for (int i = 0; i < size; i++) {
                 final String pkg = pkgList[i];
                 final int uid = uidList[i];
-                synchronized (mPackagePreferences) {
+                synchronized (mLock) {
                     mPackagePreferences.remove(packagePreferencesKey(pkg, uid));
+                    mRestoredWithoutUids.remove(unrestoredPackageKey(pkg, changeUserId));
                 }
-                mRestoredWithoutUids.remove(unrestoredPackageKey(pkg, changeUserId));
                 updated = true;
             }
         } else {
             for (String pkg : pkgList) {
-                // Package install
-                final PackagePreferences r =
-                        mRestoredWithoutUids.get(unrestoredPackageKey(pkg, changeUserId));
-                if (r != null) {
-                    try {
-                        r.uid = mPm.getPackageUidAsUser(r.pkg, changeUserId);
-                        mRestoredWithoutUids.remove(unrestoredPackageKey(pkg, changeUserId));
-                        synchronized (mPackagePreferences) {
+                try {
+                    // Package install
+                    int uid = mPm.getPackageUidAsUser(pkg, changeUserId);
+                    PackagePermission p = null;
+                    synchronized (mLock) {
+                        final PackagePreferences r =
+                                mRestoredWithoutUids.get(unrestoredPackageKey(pkg, changeUserId));
+                        if (r != null) {
+                            r.uid = uid;
+                            mRestoredWithoutUids.remove(unrestoredPackageKey(pkg, changeUserId));
                             mPackagePreferences.put(packagePreferencesKey(r.pkg, r.uid), r);
 
                             // Try to restore any unrestored sound resources
@@ -2715,32 +2723,29 @@
                                     channel.setSound(restoredUri, channel.getAudioAttributes());
                                 }
                             }
-                        }
-                        if (r.migrateToPm) {
-                            try {
-                                PackagePermission p = new PackagePermission(
+
+                            if (r.migrateToPm) {
+                                p = new PackagePermission(
                                         r.pkg, UserHandle.getUserId(r.uid),
                                         r.importance != IMPORTANCE_NONE,
                                         hasUserConfiguredSettings(r));
-                                mPermissionHelper.setNotificationPermission(p);
-                            } catch (Exception e) {
-                                Slog.e(TAG, "could not migrate setting for " + r.pkg, e);
                             }
+                            updated = true;
                         }
-                        updated = true;
-                    } catch (Exception e) {
-                        Slog.e(TAG, "could not restore " + r.pkg, e);
                     }
+                    if (p != null) {
+                        mPermissionHelper.setNotificationPermission(p);
+                    }
+                } catch (Exception e) {
+                    Slog.e(TAG, "could not restore " + pkg, e);
                 }
                 // Package upgrade
                 try {
-                    synchronized (mPackagePreferences) {
-                        PackagePreferences fullPackagePreferences = getPackagePreferencesLocked(pkg,
-                                mPm.getPackageUidAsUser(pkg, changeUserId));
-                        if (fullPackagePreferences != null) {
-                            updated |= createDefaultChannelIfNeededLocked(fullPackagePreferences);
-                            updated |= deleteDefaultChannelIfNeededLocked(fullPackagePreferences);
-                        }
+                    PackagePreferences fullPackagePreferences = getPackagePreferencesLocked(pkg,
+                            mPm.getPackageUidAsUser(pkg, changeUserId));
+                    if (fullPackagePreferences != null) {
+                        updated |= createDefaultChannelIfNeededLocked(fullPackagePreferences);
+                        updated |= deleteDefaultChannelIfNeededLocked(fullPackagePreferences);
                     }
                 } catch (PackageManager.NameNotFoundException e) {
                 }
@@ -2754,7 +2759,7 @@
     }
 
     public void clearData(String pkg, int uid) {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             PackagePreferences p = getPackagePreferencesLocked(pkg, uid);
             if (p != null) {
                 p.channels = new ArrayMap<>();
@@ -2941,7 +2946,7 @@
     }
 
     public void unlockAllNotificationChannels() {
-        synchronized (mPackagePreferences) {
+        synchronized (mLock) {
             final int numPackagePreferences = mPackagePreferences.size();
             for (int i = 0; i < numPackagePreferences; i++) {
                 final PackagePreferences r = mPackagePreferences.valueAt(i);
@@ -2958,7 +2963,7 @@
                     PackageManager.PackageInfoFlags.of(PackageManager.MATCH_ALL),
                     user.getUserHandle().getIdentifier());
             for (PackageInfo pi : packages) {
-                synchronized (mPackagePreferences) {
+                synchronized (mLock) {
                     PackagePreferences p = getOrCreatePackagePreferencesLocked(
                             pi.packageName, pi.applicationInfo.uid);
                     if (p.migrateToPm && p.uid != UNKNOWN_UID) {
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index efe07dc..b0f92e8 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -6535,9 +6535,7 @@
             // and the token could be null.
             return;
         }
-        if (r.mDisplayContent.mActivityRefresher != null) {
-            r.mDisplayContent.mActivityRefresher.onActivityRefreshed(r);
-        }
+        r.mDisplayContent.mAppCompatCameraPolicy.onActivityRefreshed(r);
     }
 
     static void splashScreenAttachedLocked(IBinder token) {
@@ -8194,7 +8192,7 @@
     }
 
     void setRequestedOrientation(@ActivityInfo.ScreenOrientation int requestedOrientation) {
-        if (mAppCompatController.getAppCompatOrientationOverrides()
+        if (mAppCompatController.getOrientationPolicy()
                 .shouldIgnoreRequestedOrientation(requestedOrientation)) {
             return;
         }
@@ -10030,16 +10028,6 @@
         return updateReportedConfigurationAndSend();
     }
 
-    /**
-     * @return {@code true} if the Camera is active for the current activity
-     */
-    boolean isCameraActive() {
-        return mDisplayContent != null
-                && mDisplayContent.getDisplayRotationCompatPolicy() != null
-                && mDisplayContent.getDisplayRotationCompatPolicy()
-                    .isCameraActive(this, /* mustBeFullscreen */ true);
-    }
-
     boolean updateReportedConfigurationAndSend() {
         if (isConfigurationDispatchPaused()) {
             Slog.wtf(TAG, "trying to update reported(client) config while dispatch is paused");
@@ -10187,11 +10175,10 @@
 
     private void notifyActivityRefresherAboutConfigurationChange(
             Configuration newConfig, Configuration lastReportedConfig) {
-        if (mDisplayContent.mActivityRefresher == null
-                || !shouldBeResumed(/* activeActivity */ null)) {
+        if (!shouldBeResumed(/* activeActivity */ null)) {
             return;
         }
-        mDisplayContent.mActivityRefresher.onActivityConfigurationChanging(
+        mDisplayContent.mAppCompatCameraPolicy.onActivityConfigurationChanging(
                 this, newConfig, lastReportedConfig);
     }
 
diff --git a/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java b/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
index c0e5005..0d108e1 100644
--- a/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
@@ -97,8 +97,7 @@
      * </ul>
      */
     boolean shouldOverrideMinAspectRatioForCamera() {
-        return mActivityRecord.isCameraActive()
-                && mAllowMinAspectRatioOverrideOptProp
+        return isCameraActive() && mAllowMinAspectRatioOverrideOptProp
                 .shouldEnableWithOptInOverrideAndOptOutProperty(
                         isCompatChangeEnabled(OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA));
     }
@@ -174,6 +173,15 @@
     }
 
     /**
+     * @return {@code true} if the Camera is active for the current activity
+     */
+    boolean isCameraActive() {
+        return mActivityRecord.mDisplayContent != null
+                && mActivityRecord.mDisplayContent.mAppCompatCameraPolicy
+                    .isCameraActive(mActivityRecord, /* mustBeFullscreen */ true);
+    }
+
+    /**
      * @return {@code true} if the configuration needs to be recomputed after a camera state update.
      */
     boolean shouldRecomputeConfigurationForCameraCompat() {
diff --git a/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java b/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java
index ee523a2..53729a2 100644
--- a/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java
+++ b/services/core/java/com/android/server/wm/AppCompatCameraPolicy.java
@@ -16,34 +16,158 @@
 
 package com.android.server.wm;
 
-import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_ATM;
-import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
+import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.content.pm.ActivityInfo.ScreenOrientation;
+import android.content.res.Configuration;
+import android.widget.Toast;
+
+import com.android.window.flags.Flags;
 
 /**
- * Encapsulate the app compat logic related to camera.
+ * Encapsulate policy logic related to app compat display rotation.
  */
 class AppCompatCameraPolicy {
 
-    private static final String TAG = TAG_WITH_CLASS_NAME
-            ? "AppCompatCameraPolicy" : TAG_ATM;
+    @Nullable
+    private final CameraStateMonitor mCameraStateMonitor;
+    @Nullable
+    private final ActivityRefresher mActivityRefresher;
+    @Nullable
+    final DisplayRotationCompatPolicy mDisplayRotationCompatPolicy;
+    @Nullable
+    final CameraCompatFreeformPolicy mCameraCompatFreeformPolicy;
 
-    @NonNull
-    private final ActivityRecord mActivityRecord;
-
-    @NonNull
-    private final AppCompatCameraOverrides mAppCompatCameraOverrides;
-
-    AppCompatCameraPolicy(@NonNull ActivityRecord activityRecord,
-            @NonNull AppCompatCameraOverrides appCompatCameraOverrides) {
-        mActivityRecord = activityRecord;
-        mAppCompatCameraOverrides = appCompatCameraOverrides;
-    }
-
-    void recomputeConfigurationForCameraCompatIfNeeded() {
-        if (mAppCompatCameraOverrides.shouldRecomputeConfigurationForCameraCompat()) {
-            mActivityRecord.recomputeConfiguration();
+    AppCompatCameraPolicy(@NonNull WindowManagerService wmService,
+            @NonNull DisplayContent displayContent) {
+        // Not checking DeviceConfig value here to allow enabling via DeviceConfig
+        // without the need to restart the device.
+        final boolean needsDisplayRotationCompatPolicy =
+                wmService.mLetterboxConfiguration.isCameraCompatTreatmentEnabledAtBuildTime();
+        final boolean needsCameraCompatFreeformPolicy = Flags.cameraCompatForFreeform()
+                && DesktopModeLaunchParamsModifier.canEnterDesktopMode(wmService.mContext);
+        if (needsDisplayRotationCompatPolicy || needsCameraCompatFreeformPolicy) {
+            mCameraStateMonitor = new CameraStateMonitor(displayContent, wmService.mH);
+            mActivityRefresher = new ActivityRefresher(wmService, wmService.mH);
+            mDisplayRotationCompatPolicy =
+                    needsDisplayRotationCompatPolicy ? new DisplayRotationCompatPolicy(
+                            displayContent, mCameraStateMonitor, mActivityRefresher) : null;
+            mCameraCompatFreeformPolicy =
+                    needsCameraCompatFreeformPolicy ? new CameraCompatFreeformPolicy(displayContent,
+                            mCameraStateMonitor, mActivityRefresher) : null;
+        } else {
+            mDisplayRotationCompatPolicy = null;
+            mCameraCompatFreeformPolicy = null;
+            mCameraStateMonitor = null;
+            mActivityRefresher = null;
         }
     }
+
+    void onActivityRefreshed(@NonNull ActivityRecord activity) {
+        if (mActivityRefresher != null) {
+            mActivityRefresher.onActivityRefreshed(activity);
+        }
+    }
+
+    /**
+     * "Refreshes" activity by going through "stopped -> resumed" or "paused -> resumed" cycle.
+     * This allows to clear cached values in apps (e.g. display or camera rotation) that influence
+     * camera preview and can lead to sideways or stretching issues persisting even after force
+     * rotation.
+     */
+    void onActivityConfigurationChanging(@NonNull ActivityRecord activity,
+            @NonNull Configuration newConfig, @NonNull Configuration lastReportedConfig) {
+        if (mActivityRefresher != null) {
+            mActivityRefresher.onActivityConfigurationChanging(activity, newConfig,
+                    lastReportedConfig);
+        }
+    }
+
+    /**
+     * Notifies that animation in {@link ScreenRotationAnimation} has finished.
+     *
+     * <p>This class uses this signal as a trigger for notifying the user about forced rotation
+     * reason with the {@link Toast}.
+     */
+    void onScreenRotationAnimationFinished() {
+        if (mDisplayRotationCompatPolicy != null) {
+            mDisplayRotationCompatPolicy.onScreenRotationAnimationFinished();
+        }
+    }
+
+    boolean isActivityEligibleForOrientationOverride(@NonNull ActivityRecord activity) {
+        if (mDisplayRotationCompatPolicy != null) {
+            return mDisplayRotationCompatPolicy.isActivityEligibleForOrientationOverride(activity);
+        }
+        return false;
+    }
+
+    /**
+     * Whether camera compat treatment is applicable for the given activity.
+     *
+     * <p>Conditions that need to be met:
+     * <ul>
+     *     <li>Camera is active for the package.
+     *     <li>The activity is in fullscreen
+     *     <li>The activity has fixed orientation but not "locked" or "nosensor" one.
+     * </ul>
+     */
+    boolean isTreatmentEnabledForActivity(@Nullable ActivityRecord activity) {
+        if (mDisplayRotationCompatPolicy != null) {
+            return mDisplayRotationCompatPolicy.isTreatmentEnabledForActivity(activity);
+        }
+        return false;
+    }
+
+    void start() {
+        if (mCameraCompatFreeformPolicy != null) {
+            mCameraCompatFreeformPolicy.start();
+        }
+        if (mCameraStateMonitor != null) {
+            mCameraStateMonitor.startListeningToCameraState();
+        }
+    }
+
+    void dispose() {
+        if (mDisplayRotationCompatPolicy != null) {
+            mDisplayRotationCompatPolicy.dispose();
+        }
+        if (mCameraCompatFreeformPolicy != null) {
+            mCameraCompatFreeformPolicy.dispose();
+        }
+        if (mCameraStateMonitor != null) {
+            mCameraStateMonitor.dispose();
+        }
+    }
+
+    boolean hasDisplayRotationCompatPolicy() {
+        return mDisplayRotationCompatPolicy != null;
+    }
+
+    boolean hasCameraCompatFreeformPolicy() {
+        return mCameraCompatFreeformPolicy != null;
+    }
+
+    @ScreenOrientation
+    int getOrientation() {
+        return mDisplayRotationCompatPolicy != null
+                ? mDisplayRotationCompatPolicy.getOrientation()
+                : SCREEN_ORIENTATION_UNSPECIFIED;
+    }
+
+    boolean isCameraActive(@NonNull ActivityRecord activity, boolean mustBeFullscreen) {
+        return mDisplayRotationCompatPolicy != null
+                && mDisplayRotationCompatPolicy.isCameraActive(activity, mustBeFullscreen);
+    }
+
+    @Nullable
+    String getSummaryForDisplayRotationHistoryRecord() {
+        if (mDisplayRotationCompatPolicy != null) {
+            return mDisplayRotationCompatPolicy.getSummaryForDisplayRotationHistoryRecord();
+        }
+        return null;
+    }
+
 }
diff --git a/services/core/java/com/android/server/wm/AppCompatController.java b/services/core/java/com/android/server/wm/AppCompatController.java
index d8c0c17..16d3787 100644
--- a/services/core/java/com/android/server/wm/AppCompatController.java
+++ b/services/core/java/com/android/server/wm/AppCompatController.java
@@ -16,6 +16,7 @@
 package com.android.server.wm;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.pm.PackageManager;
 
 import com.android.server.wm.utils.OptPropFactory;
@@ -26,16 +27,17 @@
 class AppCompatController {
 
     @NonNull
+    private final ActivityRecord mActivityRecord;
+    @NonNull
     private final TransparentPolicy mTransparentPolicy;
     @NonNull
     private final AppCompatOrientationPolicy mOrientationPolicy;
     @NonNull
     private final AppCompatOverrides mAppCompatOverrides;
-    @NonNull
-    private final AppCompatCameraPolicy mAppCompatCameraPolicy;
 
     AppCompatController(@NonNull WindowManagerService wmService,
                         @NonNull ActivityRecord activityRecord) {
+        mActivityRecord = activityRecord;
         final PackageManager packageManager = wmService.mContext.getPackageManager();
         final OptPropFactory optPropBuilder = new OptPropFactory(packageManager,
                 activityRecord.packageName);
@@ -49,8 +51,6 @@
                 mAppCompatOverrides, tmpController::shouldApplyUserFullscreenOverride,
                 tmpController::shouldApplyUserMinAspectRatioOverride,
                 tmpController::isSystemOverrideToFullscreenEnabled);
-        mAppCompatCameraPolicy = new AppCompatCameraPolicy(activityRecord,
-                mAppCompatOverrides.getAppCompatCameraOverrides());
     }
 
     @NonNull
@@ -64,11 +64,6 @@
     }
 
     @NonNull
-    AppCompatCameraPolicy getAppCompatCameraPolicy() {
-        return mAppCompatCameraPolicy;
-    }
-
-    @NonNull
     AppCompatOverrides getAppCompatOverrides() {
         return mAppCompatOverrides;
     }
@@ -82,4 +77,12 @@
     AppCompatCameraOverrides getAppCompatCameraOverrides() {
         return mAppCompatOverrides.getAppCompatCameraOverrides();
     }
+
+    @Nullable
+    AppCompatCameraPolicy getAppCompatCameraPolicy() {
+        if (mActivityRecord.mDisplayContent != null) {
+            return mActivityRecord.mDisplayContent.mAppCompatCameraPolicy;
+        }
+        return null;
+    }
 }
diff --git a/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java b/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java
index b0fdbb5..155e246 100644
--- a/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatOrientationOverrides.java
@@ -22,7 +22,6 @@
 import static android.content.pm.ActivityInfo.OVERRIDE_LANDSCAPE_ORIENTATION_TO_REVERSE_LANDSCAPE;
 import static android.content.pm.ActivityInfo.OVERRIDE_UNDEFINED_ORIENTATION_TO_NOSENSOR;
 import static android.content.pm.ActivityInfo.OVERRIDE_UNDEFINED_ORIENTATION_TO_PORTRAIT;
-import static android.content.pm.ActivityInfo.screenOrientationToString;
 import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_IGNORING_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED;
 import static android.view.WindowManager.PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION;
 
@@ -31,8 +30,6 @@
 import static com.android.server.wm.AppCompatUtils.asLazy;
 
 import android.annotation.NonNull;
-import android.content.pm.ActivityInfo;
-import android.util.Slog;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.wm.utils.OptPropFactory;
@@ -51,6 +48,8 @@
 
     @NonNull
     private final ActivityRecord mActivityRecord;
+    @NonNull
+    private final AppCompatCameraOverrides mAppCompatCameraOverrides;
 
     @NonNull
     private final OptPropFactory.OptProp mIgnoreRequestedOrientationOptProp;
@@ -62,8 +61,10 @@
 
     AppCompatOrientationOverrides(@NonNull ActivityRecord activityRecord,
             @NonNull LetterboxConfiguration letterboxConfiguration,
-            @NonNull OptPropFactory optPropBuilder) {
+            @NonNull OptPropFactory optPropBuilder,
+            @NonNull AppCompatCameraOverrides appCompatCameraOverrides) {
         mActivityRecord = activityRecord;
+        mAppCompatCameraOverrides = appCompatCameraOverrides;
         mOrientationOverridesState = new OrientationOverridesState(mActivityRecord,
                 System::currentTimeMillis);
         final BooleanSupplier isPolicyForIgnoringRequestedOrientationEnabled = asLazy(
@@ -76,59 +77,9 @@
                 isPolicyForIgnoringRequestedOrientationEnabled);
     }
 
-    /**
-     * Whether should ignore app requested orientation in response to an app
-     * calling {@link android.app.Activity#setRequestedOrientation}.
-     *
-     * <p>This is needed to avoid getting into {@link android.app.Activity#setRequestedOrientation}
-     * loop when {@link DisplayContent#getIgnoreOrientationRequest} is enabled or device has
-     * landscape natural orientation which app developers don't expect. For example, the loop can
-     * look like this:
-     * <ol>
-     *     <li>App sets default orientation to "unspecified" at runtime
-     *     <li>App requests to "portrait" after checking some condition (e.g. display rotation).
-     *     <li>(2) leads to fullscreen -> letterboxed bounds change and activity relaunch because
-     *     app can't handle the corresponding config changes.
-     *     <li>Loop goes back to (1)
-     * </ol>
-     *
-     * <p>This treatment is enabled when the following conditions are met:
-     * <ul>
-     *     <li>Flag gating the treatment is enabled
-     *     <li>Opt-out component property isn't enabled
-     *     <li>Opt-in component property or per-app override are enabled
-     *     <li>Activity is relaunched after {@link android.app.Activity#setRequestedOrientation}
-     *     call from an app or camera compat force rotation treatment is active for the activity.
-     *     <li>Orientation request loop detected and is not letterboxed for fixed orientation
-     * </ul>
-     */
-    boolean shouldIgnoreRequestedOrientation(
-            @ActivityInfo.ScreenOrientation int requestedOrientation) {
-        if (mIgnoreRequestedOrientationOptProp.shouldEnableWithOverrideAndProperty(
-                isCompatChangeEnabled(OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION))) {
-            if (mOrientationOverridesState.mIsRelaunchingAfterRequestedOrientationChanged) {
-                Slog.w(TAG, "Ignoring orientation update to "
-                        + screenOrientationToString(requestedOrientation)
-                        + " due to relaunching after setRequestedOrientation for "
-                        + mActivityRecord);
-                return true;
-            }
-            if (isCameraCompatTreatmentActive()) {
-                Slog.w(TAG, "Ignoring orientation update to "
-                        + screenOrientationToString(requestedOrientation)
-                        + " due to camera compat treatment for " + mActivityRecord);
-                return true;
-            }
-        }
-
-        if (shouldIgnoreOrientationRequestLoop()) {
-            Slog.w(TAG, "Ignoring orientation update to "
-                    + screenOrientationToString(requestedOrientation)
-                    + " as orientation request loop was detected for "
-                    + mActivityRecord);
-            return true;
-        }
-        return false;
+    boolean shouldEnableIgnoreOrientationRequest() {
+        return mIgnoreRequestedOrientationOptProp.shouldEnableWithOverrideAndProperty(
+                isCompatChangeEnabled(OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION));
     }
 
     /**
@@ -183,20 +134,6 @@
         return mActivityRecord.info.isChangeEnabled(overrideChangeId);
     }
 
-    /**
-     * @return {@code true} if the App Compat Camera Policy is active for the current activity.
-     */
-    // TODO(b/346253439): Remove after defining dependency with Camera capabilities.
-    private boolean isCameraCompatTreatmentActive() {
-        DisplayContent displayContent = mActivityRecord.mDisplayContent;
-        if (displayContent == null) {
-            return false;
-        }
-        return displayContent.mDisplayRotationCompatPolicy != null
-                && displayContent.mDisplayRotationCompatPolicy
-                .isTreatmentEnabledForActivity(mActivityRecord);
-    }
-
     static class OrientationOverridesState {
         // Corresponds to OVERRIDE_UNDEFINED_ORIENTATION_TO_NOSENSOR
         final boolean mIsOverrideToNosensorOrientationEnabled;
diff --git a/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java b/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java
index 960ef5a..69ba59b 100644
--- a/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java
+++ b/services/core/java/com/android/server/wm/AppCompatOrientationPolicy.java
@@ -46,7 +46,6 @@
 
     @NonNull
     private final AppCompatOverrides mAppCompatOverrides;
-
     @NonNull
     private final BooleanSupplier mShouldApplyUserFullscreenOverride;
     @NonNull
@@ -78,7 +77,7 @@
                 // often results in sideways or stretched previews. As the camera compat treatment
                 // targets fixed-orientation activities, overriding the orientation disables the
                 // treatment.
-                && !mActivityRecord.isCameraActive()) {
+                && !mAppCompatOverrides.getAppCompatCameraOverrides().isCameraActive()) {
             Slog.v(TAG, "Requested orientation " + screenOrientationToString(candidate)
                     + " for " + mActivityRecord + " is overridden to "
                     + screenOrientationToString(SCREEN_ORIENTATION_USER)
@@ -103,11 +102,11 @@
             return candidate;
         }
 
-        if (displayContent != null && mAppCompatOverrides.getAppCompatCameraOverrides()
-                .isOverrideOrientationOnlyForCameraEnabled()
-                    && (displayContent.mDisplayRotationCompatPolicy == null
-                    || !displayContent.mDisplayRotationCompatPolicy
-                        .isActivityEligibleForOrientationOverride(mActivityRecord))) {
+        if (displayContent != null
+                && mAppCompatOverrides.getAppCompatCameraOverrides()
+                    .isOverrideOrientationOnlyForCameraEnabled()
+                && !displayContent.mAppCompatCameraPolicy
+                    .isActivityEligibleForOrientationOverride(mActivityRecord)) {
             return candidate;
         }
 
@@ -120,7 +119,7 @@
                 // often results in sideways or stretched previews. As the camera compat treatment
                 // targets fixed-orientation activities, overriding the orientation disables the
                 // treatment.
-                && !mActivityRecord.isCameraActive()) {
+                && !mAppCompatOverrides.getAppCompatCameraOverrides().isCameraActive()) {
             Slog.v(TAG, "Requested orientation  " + screenOrientationToString(candidate)
                     + " for " + mActivityRecord + " is overridden to "
                     + screenOrientationToString(SCREEN_ORIENTATION_USER));
@@ -161,4 +160,62 @@
         return candidate;
     }
 
+    /**
+     * Whether should ignore app requested orientation in response to an app
+     * calling {@link android.app.Activity#setRequestedOrientation}.
+     *
+     * <p>This is needed to avoid getting into {@link android.app.Activity#setRequestedOrientation}
+     * loop when {@link DisplayContent#getIgnoreOrientationRequest} is enabled or device has
+     * landscape natural orientation which app developers don't expect. For example, the loop can
+     * look like this:
+     * <ol>
+     *     <li>App sets default orientation to "unspecified" at runtime
+     *     <li>App requests to "portrait" after checking some condition (e.g. display rotation).
+     *     <li>(2) leads to fullscreen -> letterboxed bounds change and activity relaunch because
+     *     app can't handle the corresponding config changes.
+     *     <li>Loop goes back to (1)
+     * </ol>
+     *
+     * <p>This treatment is enabled when the following conditions are met:
+     * <ul>
+     *     <li>Flag gating the treatment is enabled
+     *     <li>Opt-out component property isn't enabled
+     *     <li>Opt-in component property or per-app override are enabled
+     *     <li>Activity is relaunched after {@link android.app.Activity#setRequestedOrientation}
+     *     call from an app or camera compat force rotation treatment is active for the activity.
+     *     <li>Orientation request loop detected and is not letterboxed for fixed orientation
+     * </ul>
+     */
+    boolean shouldIgnoreRequestedOrientation(
+            @ActivityInfo.ScreenOrientation int requestedOrientation) {
+        final AppCompatOrientationOverrides orientationOverrides =
+                mAppCompatOverrides.getAppCompatOrientationOverrides();
+        if (orientationOverrides.shouldEnableIgnoreOrientationRequest()) {
+            if (orientationOverrides.getIsRelaunchingAfterRequestedOrientationChanged()) {
+                Slog.w(TAG, "Ignoring orientation update to "
+                        + screenOrientationToString(requestedOrientation)
+                        + " due to relaunching after setRequestedOrientation for "
+                        + mActivityRecord);
+                return true;
+            }
+            final AppCompatCameraPolicy cameraPolicy = mActivityRecord.mAppCompatController
+                    .getAppCompatCameraPolicy();
+            if (cameraPolicy != null
+                    && cameraPolicy.isTreatmentEnabledForActivity(mActivityRecord)) {
+                Slog.w(TAG, "Ignoring orientation update to "
+                        + screenOrientationToString(requestedOrientation)
+                        + " due to camera compat treatment for " + mActivityRecord);
+                return true;
+            }
+        }
+        if (orientationOverrides.shouldIgnoreOrientationRequestLoop()) {
+            Slog.w(TAG, "Ignoring orientation update to "
+                    + screenOrientationToString(requestedOrientation)
+                    + " as orientation request loop was detected for "
+                    + mActivityRecord);
+            return true;
+        }
+        return false;
+    }
+
 }
diff --git a/services/core/java/com/android/server/wm/AppCompatOverrides.java b/services/core/java/com/android/server/wm/AppCompatOverrides.java
index c20da7c..94c6ba9 100644
--- a/services/core/java/com/android/server/wm/AppCompatOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatOverrides.java
@@ -79,10 +79,10 @@
         mLetterboxConfiguration = letterboxConfiguration;
         mActivityRecord = activityRecord;
 
-        mAppCompatOrientationOverrides = new AppCompatOrientationOverrides(mActivityRecord,
-                mLetterboxConfiguration, optPropBuilder);
         mAppCompatCameraOverrides = new AppCompatCameraOverrides(mActivityRecord,
                 mLetterboxConfiguration, optPropBuilder);
+        mAppCompatOrientationOverrides = new AppCompatOrientationOverrides(mActivityRecord,
+                mLetterboxConfiguration, optPropBuilder, mAppCompatCameraOverrides);
 
         mFakeFocusOptProp = optPropBuilder.create(PROPERTY_COMPAT_ENABLE_FAKE_FOCUS,
                 mLetterboxConfiguration::isCompatFakeFocusEnabled);
@@ -113,19 +113,6 @@
                 mLetterboxConfiguration::isUserAppAspectRatioFullscreenEnabled);
     }
 
-    /**
-     * @return {@code true} if the App Compat Camera Policy is active for the current activity.
-     */
-    boolean isCameraCompatTreatmentActive() {
-        final DisplayContent displayContent = mActivityRecord.mDisplayContent;
-        if (displayContent == null) {
-            return false;
-        }
-        return displayContent.mDisplayRotationCompatPolicy != null
-                && displayContent.mDisplayRotationCompatPolicy
-                    .isTreatmentEnabledForActivity(mActivityRecord);
-    }
-
     @NonNull
     AppCompatOrientationOverrides getAppCompatOrientationOverrides() {
         return mAppCompatOrientationOverrides;
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index b5b9377..a8aa0ba 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -263,7 +263,6 @@
 import com.android.server.wm.utils.RegionUtils;
 import com.android.server.wm.utils.RotationCache;
 import com.android.server.wm.utils.WmDisplayCutout;
-import com.android.window.flags.Flags;
 
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
@@ -475,14 +474,8 @@
     private final DisplayPolicy mDisplayPolicy;
     private final DisplayRotation mDisplayRotation;
 
-    @Nullable
-    final DisplayRotationCompatPolicy mDisplayRotationCompatPolicy;
-    @Nullable
-    final CameraCompatFreeformPolicy mCameraCompatFreeformPolicy;
-    @Nullable
-    final CameraStateMonitor mCameraStateMonitor;
-    @Nullable
-    final ActivityRefresher mActivityRefresher;
+    @NonNull
+    AppCompatCameraPolicy mAppCompatCameraPolicy;
 
     DisplayFrames mDisplayFrames;
     final DisplayUpdater mDisplayUpdater;
@@ -1191,6 +1184,7 @@
 
         mDeviceStateController = deviceStateController;
 
+        mAppCompatCameraPolicy = new AppCompatCameraPolicy(mWmService, this);
         mDisplayPolicy = new DisplayPolicy(mWmService, this);
         mDisplayRotation = new DisplayRotation(mWmService, this, mDisplayInfo.address,
                 mDeviceStateController, root.getDisplayRotationCoordinator());
@@ -1231,40 +1225,6 @@
         onDisplayChanged(this);
         updateDisplayAreaOrganizers();
 
-        // Not checking DeviceConfig value here to allow enabling via DeviceConfig
-        // without the need to restart the device.
-        final boolean shouldCreateDisplayRotationCompatPolicy =
-                mWmService.mLetterboxConfiguration.isCameraCompatTreatmentEnabledAtBuildTime();
-        final boolean shouldCreateCameraCompatFreeformPolicy = Flags.cameraCompatForFreeform()
-                && DesktopModeLaunchParamsModifier.canEnterDesktopMode(mWmService.mContext);
-        if (shouldCreateDisplayRotationCompatPolicy || shouldCreateCameraCompatFreeformPolicy) {
-            mCameraStateMonitor = new CameraStateMonitor(this, mWmService.mH);
-            mActivityRefresher = new ActivityRefresher(mWmService, mWmService.mH);
-            if (shouldCreateDisplayRotationCompatPolicy) {
-                mDisplayRotationCompatPolicy = new DisplayRotationCompatPolicy(this,
-                        mCameraStateMonitor, mActivityRefresher);
-                mDisplayRotationCompatPolicy.start();
-            } else {
-                mDisplayRotationCompatPolicy = null;
-            }
-
-            if (shouldCreateCameraCompatFreeformPolicy) {
-                mCameraCompatFreeformPolicy = new CameraCompatFreeformPolicy(this,
-                        mCameraStateMonitor, mActivityRefresher);
-                mCameraCompatFreeformPolicy.start();
-            } else {
-                mCameraCompatFreeformPolicy = null;
-            }
-
-            mCameraStateMonitor.startListeningToCameraState();
-        } else {
-            // These are to satisfy the `final` check.
-            mCameraStateMonitor = null;
-            mActivityRefresher = null;
-            mDisplayRotationCompatPolicy = null;
-            mCameraCompatFreeformPolicy = null;
-        }
-
         mRotationReversionController = new DisplayRotationReversionController(this);
 
         mInputMonitor = new InputMonitor(mWmService, this);
@@ -1280,6 +1240,7 @@
                 R.bool.config_defaultInTouchMode);
         mWmService.mInputManager.setInTouchMode(mInTouchMode, mWmService.MY_PID, mWmService.MY_UID,
                 /* hasPermission= */ true, mDisplayId);
+        mAppCompatCameraPolicy.start();
     }
 
     private void beginHoldScreenUpdate() {
@@ -1314,15 +1275,6 @@
         }
     }
 
-    /**
-     * @return The {@link DisplayRotationCompatPolicy} for this DisplayContent
-     */
-    // TODO(b/335387481) Allow access to DisplayRotationCompatPolicy only with getters
-    @Nullable
-    DisplayRotationCompatPolicy getDisplayRotationCompatPolicy() {
-        return mDisplayRotationCompatPolicy;
-    }
-
     @Override
     void migrateToNewSurfaceControl(Transaction t) {
         t.remove(mSurfaceControl);
@@ -2889,12 +2841,10 @@
             }
         }
 
-        if (mDisplayRotationCompatPolicy != null) {
-            int compatOrientation = mDisplayRotationCompatPolicy.getOrientation();
-            if (compatOrientation != SCREEN_ORIENTATION_UNSPECIFIED) {
-                mLastOrientationSource = null;
-                return compatOrientation;
-            }
+        final int compatOrientation = mAppCompatCameraPolicy.getOrientation();
+        if (compatOrientation != SCREEN_ORIENTATION_UNSPECIFIED) {
+            mLastOrientationSource = null;
+            return compatOrientation;
         }
 
         final int orientation = super.getOrientation();
@@ -3364,17 +3314,7 @@
         getPendingTransaction().apply();
         mWmService.mWindowPlacerLocked.requestTraversal();
 
-        if (mDisplayRotationCompatPolicy != null) {
-            mDisplayRotationCompatPolicy.dispose();
-        }
-
-        if (mCameraCompatFreeformPolicy != null) {
-            mCameraCompatFreeformPolicy.dispose();
-        }
-
-        if (mCameraStateMonitor != null) {
-            mCameraStateMonitor.dispose();
-        }
+        mAppCompatCameraPolicy.dispose();
     }
 
     /** Returns true if a removal action is still being deferred. */
diff --git a/services/core/java/com/android/server/wm/DisplayRotation.java b/services/core/java/com/android/server/wm/DisplayRotation.java
index f3ccc3b..c67928a 100644
--- a/services/core/java/com/android/server/wm/DisplayRotation.java
+++ b/services/core/java/com/android/server/wm/DisplayRotation.java
@@ -2294,10 +2294,8 @@
                     mInHalfFoldTransition = false;
                     mDeviceState = DeviceStateController.DeviceState.UNKNOWN;
                 }
-                mDisplayRotationCompatPolicySummary = dc.mDisplayRotationCompatPolicy == null
-                        ? null
-                        : dc.mDisplayRotationCompatPolicy
-                                .getSummaryForDisplayRotationHistoryRecord();
+                mDisplayRotationCompatPolicySummary = dc.mAppCompatCameraPolicy
+                        .getSummaryForDisplayRotationHistoryRecord();
                 mRotationReversionSlots =
                         dr.mDisplayContent.getRotationReversionController().getSlotsCopy();
             }
diff --git a/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java b/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java
index 3d71e95..9998e1a 100644
--- a/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayRotationCompatPolicy.java
@@ -299,8 +299,7 @@
         // Checking whether an activity in fullscreen rather than the task as this camera
         // compat treatment doesn't cover activity embedding.
         if (cameraActivity.getWindowingMode() == WINDOWING_MODE_FULLSCREEN) {
-            cameraActivity.mAppCompatController
-                    .getAppCompatCameraPolicy().recomputeConfigurationForCameraCompatIfNeeded();
+            recomputeConfigurationForCameraCompatIfNeeded(cameraActivity);
             mDisplayContent.updateOrientation();
             return true;
         }
@@ -367,8 +366,7 @@
                 || topActivity.getWindowingMode() != WINDOWING_MODE_FULLSCREEN) {
             return true;
         }
-        topActivity.mAppCompatController
-                .getAppCompatCameraPolicy().recomputeConfigurationForCameraCompatIfNeeded();
+        recomputeConfigurationForCameraCompatIfNeeded(topActivity);
         mDisplayContent.updateOrientation();
         return true;
     }
@@ -383,4 +381,12 @@
         }
         return mActivityRefresher.isActivityRefreshing(topActivity);
     }
+
+    private void recomputeConfigurationForCameraCompatIfNeeded(
+            @NonNull ActivityRecord activityRecord) {
+        if (activityRecord.mAppCompatController.getAppCompatCameraOverrides()
+                .shouldRecomputeConfigurationForCameraCompat()) {
+            activityRecord.recomputeConfiguration();
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/wm/DisplayRotationReversionController.java b/services/core/java/com/android/server/wm/DisplayRotationReversionController.java
index f94b8c4..b955738 100644
--- a/services/core/java/com/android/server/wm/DisplayRotationReversionController.java
+++ b/services/core/java/com/android/server/wm/DisplayRotationReversionController.java
@@ -61,7 +61,7 @@
     }
 
     boolean isRotationReversionEnabled() {
-        return mDisplayContent.mDisplayRotationCompatPolicy != null
+        return mDisplayContent.mAppCompatCameraPolicy.hasDisplayRotationCompatPolicy()
                 || mDisplayContent.getDisplayRotation().mFoldController != null
                 || mDisplayContent.getIgnoreOrientationRequest();
     }
diff --git a/services/core/java/com/android/server/wm/LetterboxUiController.java b/services/core/java/com/android/server/wm/LetterboxUiController.java
index e924fb6..a3550bc 100644
--- a/services/core/java/com/android/server/wm/LetterboxUiController.java
+++ b/services/core/java/com/android/server/wm/LetterboxUiController.java
@@ -462,12 +462,16 @@
         final boolean isTabletopMode = isDisplayFullScreenAndInPosture(/* isTabletop */ true);
         final boolean isLandscape = isFixedOrientationLandscape(
                 mActivityRecord.getOverrideOrientation());
-
+        final AppCompatCameraOverrides appCompatCameraOverrides =
+                mActivityRecord.mAppCompatController.getAppCompatCameraOverrides();
+        final AppCompatCameraPolicy cameraPolicy =
+                mActivityRecord.mAppCompatController.getAppCompatCameraPolicy();
+        final boolean isCameraCompatTreatmentActive = cameraPolicy != null
+                && cameraPolicy.isTreatmentEnabledForActivity(mActivityRecord);
         // Don't resize to split screen size when in book mode if letterbox position is centered
         return (isBookMode && isNotCenteredHorizontally || isTabletopMode && isLandscape)
-                    || mActivityRecord.mAppCompatController.getAppCompatCameraOverrides()
-                            .isCameraCompatSplitScreenAspectRatioAllowed()
-                                && getAppCompatOverrides().isCameraCompatTreatmentActive();
+                    || (appCompatCameraOverrides.isCameraCompatSplitScreenAspectRatioAllowed()
+                    && isCameraCompatTreatmentActive);
     }
 
     private float getDefaultMinAspectRatioForUnresizableApps() {
diff --git a/services/core/java/com/android/server/wm/ScreenRotationAnimation.java b/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
index 3eb3218..31fda77 100644
--- a/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
+++ b/services/core/java/com/android/server/wm/ScreenRotationAnimation.java
@@ -815,10 +815,8 @@
                 if (mDisplayContent.getRotationAnimation() == ScreenRotationAnimation.this) {
                     // It also invokes kill().
                     mDisplayContent.setRotationAnimation(null);
-                    if (mDisplayContent.mDisplayRotationCompatPolicy != null) {
-                        mDisplayContent.mDisplayRotationCompatPolicy
-                                .onScreenRotationAnimationFinished();
-                    }
+                    mDisplayContent.mAppCompatCameraPolicy
+                            .onScreenRotationAnimationFinished();
                 } else {
                     kill();
                 }
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index f839ed6..47af6fc 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -105,8 +105,8 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.graphics.ColorUtils;
 import com.android.internal.policy.TransitionAnimation;
-import com.android.internal.protolog.ProtoLogGroup;
 import com.android.internal.protolog.ProtoLog;
+import com.android.internal.protolog.ProtoLogGroup;
 import com.android.internal.util.function.pooled.PooledLambda;
 import com.android.server.inputmethod.InputMethodManagerInternal;
 import com.android.server.statusbar.StatusBarManagerInternal;
@@ -1443,11 +1443,11 @@
                 asyncRotationController.onTransitionFinished();
             }
             dc.onTransitionFinished();
-            if (hasParticipatedDisplay && dc.mDisplayRotationCompatPolicy != null) {
+            if (hasParticipatedDisplay) {
                 final ChangeInfo changeInfo = mChanges.get(dc);
                 if (changeInfo != null
                         && changeInfo.mRotation != dc.getWindowConfiguration().getRotation()) {
-                    dc.mDisplayRotationCompatPolicy.onScreenRotationAnimationFinished();
+                    dc.mAppCompatCameraPolicy.onScreenRotationAnimationFinished();
                 }
             }
             if (mTransientLaunches != null) {
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index edd118d..1d02f1c 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -179,7 +179,7 @@
 
     // List of children for this window container. List is in z-order as the children appear on
     // screen with the top-most window container at the tail of the list.
-    protected final WindowList<E> mChildren = new WindowList<E>();
+    protected final ArrayList<E> mChildren = new ArrayList<E>();
 
     // The specified orientation for this window container.
     // Shouldn't be accessed directly since subclasses can override getOverrideOrientation.
@@ -855,7 +855,7 @@
             mSurfaceFreezer.unfreeze(getSyncTransaction());
         }
         while (!mChildren.isEmpty()) {
-            final E child = mChildren.peekLast();
+            final E child = mChildren.getLast();
             child.removeImmediately();
             // Need to do this after calling remove on the child because the child might try to
             // remove/detach itself from its parent which will cause an exception if we remove
@@ -979,7 +979,7 @@
 
         switch (position) {
             case POSITION_TOP:
-                if (mChildren.peekLast() != child) {
+                if (getTopChild() != child) {
                     mChildren.remove(child);
                     mChildren.add(child);
                     onChildPositionChanged(child);
@@ -990,7 +990,7 @@
                 }
                 break;
             case POSITION_BOTTOM:
-                if (mChildren.peekFirst() != child) {
+                if (getBottomChild() != child) {
                     mChildren.remove(child);
                     mChildren.addFirst(child);
                     onChildPositionChanged(child);
@@ -1445,7 +1445,13 @@
 
     /** Returns the top child container. */
     E getTopChild() {
-        return mChildren.peekLast();
+        final int n = mChildren.size();
+        return n == 0 ? null : mChildren.get(n - 1);
+    }
+
+    E getBottomChild() {
+        final int n = mChildren.size();
+        return n == 0 ? null : mChildren.get(0);
     }
 
     /**
@@ -2550,7 +2556,7 @@
         }
 
         if (mParent != null && mParent == other.mParent) {
-            final WindowList<WindowContainer> list = mParent.mChildren;
+            final ArrayList<WindowContainer> list = mParent.mChildren;
             return list.indexOf(this) > list.indexOf(other) ? 1 : -1;
         }
 
@@ -2587,7 +2593,7 @@
 
             // The position of the first non-common ancestor in the common ancestor list determines
             // which is greater the which.
-            final WindowList<WindowContainer> list = commonAncestor.mChildren;
+            final ArrayList<WindowContainer> list = commonAncestor.mChildren;
             return list.indexOf(thisParentChain.peekLast()) > list.indexOf(otherParentChain.peekLast())
                     ? 1 : -1;
         } finally {
diff --git a/services/core/java/com/android/server/wm/WindowList.java b/services/core/java/com/android/server/wm/WindowList.java
deleted file mode 100644
index 1e888f5..0000000
--- a/services/core/java/com/android/server/wm/WindowList.java
+++ /dev/null
@@ -1,38 +0,0 @@
-/*
- * Copyright (C) 2017 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License
- */
-
-package com.android.server.wm;
-
-import java.util.ArrayList;
-
-/**
- * An {@link ArrayList} with extended functionality to be used as the children data structure in
- * {@link WindowContainer}.
- */
-class WindowList<E> extends ArrayList<E> {
-
-    public void addFirst(E e) {
-        add(0, e);
-    }
-
-    E peekLast() {
-        return size() > 0 ? get(size() - 1) : null;
-    }
-
-    E peekFirst() {
-        return size() > 0 ? get(0) : null;
-    }
-}
diff --git a/services/core/java/com/android/server/wm/WindowOrganizerController.java b/services/core/java/com/android/server/wm/WindowOrganizerController.java
index deb7098..e900488 100644
--- a/services/core/java/com/android/server/wm/WindowOrganizerController.java
+++ b/services/core/java/com/android/server/wm/WindowOrganizerController.java
@@ -1596,7 +1596,7 @@
             case OP_TYPE_REORDER_TO_BOTTOM_OF_TASK: {
                 final Task task = taskFragment.getTask();
                 if (task != null) {
-                    if (task.mChildren.peekFirst() != taskFragment) {
+                    if (task.getBottomChild() != taskFragment) {
                         task.mChildren.remove(taskFragment);
                         task.mChildren.add(0, taskFragment);
                         if (!taskFragment.hasChild()) {
@@ -1612,7 +1612,7 @@
             case OP_TYPE_REORDER_TO_TOP_OF_TASK: {
                 final Task task = taskFragment.getTask();
                 if (task != null) {
-                    if (task.mChildren.peekLast() != taskFragment) {
+                    if (task.getTopChild() != taskFragment) {
                         task.mChildren.remove(taskFragment);
                         task.mChildren.add(taskFragment);
                         if (!taskFragment.hasChild()) {
diff --git a/services/core/jni/Android.bp b/services/core/jni/Android.bp
index 7649a4e..3cd5f76 100644
--- a/services/core/jni/Android.bp
+++ b/services/core/jni/Android.bp
@@ -210,6 +210,7 @@
         "android.system.suspend-V1-ndk",
         "server_configurable_flags",
         "service.incremental",
+        "android.companion.virtualdevice.flags-aconfig-cc",
     ],
 
     static_libs: [
diff --git a/services/core/jni/com_android_server_companion_virtual_InputController.cpp b/services/core/jni/com_android_server_companion_virtual_InputController.cpp
index 5c4db24..a32b0f1 100644
--- a/services/core/jni/com_android_server_companion_virtual_InputController.cpp
+++ b/services/core/jni/com_android_server_companion_virtual_InputController.cpp
@@ -19,6 +19,7 @@
 #include <android-base/unique_fd.h>
 #include <android/input.h>
 #include <android/keycodes.h>
+#include <android_companion_virtualdevice_flags.h>
 #include <errno.h>
 #include <fcntl.h>
 #include <input/Input.h>
@@ -37,6 +38,8 @@
 
 namespace android {
 
+namespace vd_flags = android::companion::virtualdevice::flags;
+
 static constexpr jlong INVALID_PTR = 0;
 
 enum class DeviceType {
@@ -88,6 +91,10 @@
             ioctl(fd, UI_SET_RELBIT, REL_Y);
             ioctl(fd, UI_SET_RELBIT, REL_WHEEL);
             ioctl(fd, UI_SET_RELBIT, REL_HWHEEL);
+            if (vd_flags::high_resolution_scroll()) {
+                ioctl(fd, UI_SET_RELBIT, REL_WHEEL_HI_RES);
+                ioctl(fd, UI_SET_RELBIT, REL_HWHEEL_HI_RES);
+            }
             break;
         case DeviceType::TOUCHSCREEN:
             ioctl(fd, UI_SET_EVBIT, EV_ABS);
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
index d151345..559c324 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/PreferencesHelperTest.java
@@ -188,6 +188,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ThreadLocalRandom;
 
 @SmallTest
@@ -489,6 +490,34 @@
         when(mPm.getPackageUidAsUser(eq(packageName), anyInt())).thenReturn(uid);
     }
 
+    private static void testThreadSafety(Runnable operationToTest, int nThreads,
+            int nRunsPerThread) throws Exception {
+        final CountDownLatch startLatch = new CountDownLatch(1);
+        final CountDownLatch doneLatch = new CountDownLatch(nThreads);
+
+        for (int i = 0; i < nThreads; i++) {
+            Runnable threadRunnable = () -> {
+                try {
+                    startLatch.await();
+                    for (int j = 0; j < nRunsPerThread; j++) {
+                        operationToTest.run();
+                    }
+                } catch (InterruptedException e) {
+                    e.printStackTrace();
+                } finally {
+                    doneLatch.countDown();
+                }
+            };
+            new Thread(threadRunnable, "Test Thread #" + i).start();
+        }
+
+        // Ready set go
+        startLatch.countDown();
+
+        // Wait for all test threads to be done.
+        doneLatch.await();
+    }
+
     @Test
     public void testWriteXml_onlyBackupsTargetUser() throws Exception {
         // Setup package notifications.
@@ -6193,6 +6222,36 @@
                 .isEqualTo(IMPORTANCE_LOW);
     }
 
+
+    @Test
+    public void testRestoredWithoutUid_threadSafety() throws Exception {
+        when(mPm.getPackageUidAsUser(anyString(), anyInt())).thenReturn(UNKNOWN_UID);
+        when(mPm.getApplicationInfoAsUser(anyString(), anyInt(), anyInt())).thenThrow(
+                new PackageManager.NameNotFoundException());
+        when(mClock.millis()).thenReturn(System.currentTimeMillis());
+        testThreadSafety(() -> {
+            String id = "id";
+            String xml = "<ranking version=\"1\">\n"
+                    + "<package name=\"" + Thread.currentThread()+ "\" show_badge=\"true\">\n"
+                    + "<channel id=\"" + id + "\" name=\"name\" importance=\"2\" "
+                    + "show_badge=\"true\" />\n"
+                    + "</package>\n"
+                    + "<package name=\"" + PKG_P + "\" show_badge=\"true\">\n"
+                    + "</package>\n"
+                    + "</ranking>\n";
+
+            try {
+                loadByteArrayXml(xml.getBytes(), true, USER_SYSTEM);
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+
+            // trigger a removal from the list
+            mXmlHelper.onPackagesChanged(true, USER_SYSTEM, new String[]{PKG_P},
+                    new int[]{UNKNOWN_UID});
+        }, 20, 50);
+    }
+
     private static NotificationChannel cloneChannel(NotificationChannel original) {
         Parcel parcel = Parcel.obtain();
         try {
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index ef3df6c..eb8825c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -116,7 +116,6 @@
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.when;
 
 import android.app.ActivityOptions;
 import android.app.AppOpsManager;
@@ -2482,10 +2481,10 @@
         assertTrue(activity.mChildren.contains(win4));
 
         // The starting window should be on-top of all other windows.
-        assertEquals(startingWin, activity.mChildren.peekLast());
+        assertEquals(startingWin, activity.getTopChild());
 
         // The base application window should be below all other windows.
-        assertEquals(baseWin, activity.mChildren.peekFirst());
+        assertEquals(baseWin, activity.getBottomChild());
         activity.removeImmediately();
     }
 
@@ -3508,23 +3507,6 @@
     }
 
     @Test
-    public void testIsCameraActive() {
-        final WindowState app = createWindow(null, TYPE_APPLICATION, "app");
-        final DisplayRotationCompatPolicy displayRotationCompatPolicy = mock(
-                DisplayRotationCompatPolicy.class);
-        when(mDisplayContent.getDisplayRotationCompatPolicy()).thenReturn(
-                displayRotationCompatPolicy);
-
-        when(displayRotationCompatPolicy.isCameraActive(any(ActivityRecord.class),
-                anyBoolean())).thenReturn(false);
-        assertFalse(app.mActivityRecord.isCameraActive());
-
-        when(displayRotationCompatPolicy.isCameraActive(any(ActivityRecord.class),
-                anyBoolean())).thenReturn(true);
-        assertTrue(app.mActivityRecord.isCameraActive());
-    }
-
-    @Test
     public void testUpdateCameraCompatStateFromUser_clickedOnDismiss() throws RemoteException {
         final ActivityRecord activity = createActivityWithTask();
         // Mock a flag being enabled.
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java
index 467050e..f79cdc1 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatActivityRobot.java
@@ -69,6 +69,7 @@
 
     private final int mDisplayWidth;
     private final int mDisplayHeight;
+    private DisplayContent mDisplayContent;
 
     AppCompatActivityRobot(@NonNull WindowManagerService wm,
             @NonNull ActivityTaskManagerService atm, @NonNull ActivityTaskSupervisor supervisor,
@@ -79,6 +80,7 @@
         mDisplayHeight = displayHeight;
         mActivityStack = new TestComponentStack<>();
         mTaskStack = new TestComponentStack<>();
+        createNewDisplay();
     }
 
     AppCompatActivityRobot(@NonNull WindowManagerService wm,
@@ -87,13 +89,19 @@
     }
 
     void createActivityWithComponent() {
-        createActivityWithComponentInNewTask(/* inNewTask */ mTaskStack.isEmpty());
+        createActivityWithComponentInNewTask(/* inNewTask */ mTaskStack.isEmpty(),
+                /* inNewDisplay */ false);
     }
 
     void createActivityWithComponentInNewTask() {
-        createActivityWithComponentInNewTask(/* inNewTask */ true);
+        createActivityWithComponentInNewTask(/* inNewTask */ true, /* inNewDisplay */ false);
     }
 
+    void createActivityWithComponentInNewTaskAndDisplay() {
+        createActivityWithComponentInNewTask(/* inNewTask */ true, /* inNewDisplay */ true);
+    }
+
+
     void configureTopActivity(float minAspect, float maxAspect, int screenOrientation,
             boolean isUnresizable) {
         prepareLimitedBounds(mActivityStack.top(), minAspect, maxAspect, screenOrientation,
@@ -110,12 +118,22 @@
                 /* isUnresizable */ true);
     }
 
+    void activateCameraInPolicy(boolean isCameraActive) {
+        doReturn(isCameraActive).when(mDisplayContent.mAppCompatCameraPolicy)
+                .isCameraActive(any(ActivityRecord.class), anyBoolean());
+    }
+
     @NonNull
     ActivityRecord top() {
         return mActivityStack.top();
     }
 
     @NonNull
+    DisplayContent displayContent() {
+        return mDisplayContent;
+    }
+
+    @NonNull
     ActivityRecord getFromTop(int fromTop) {
         return mActivityStack.getFromTop(fromTop);
     }
@@ -130,7 +148,7 @@
     }
 
     void enableTreatmentForTopActivity(boolean enabled) {
-        doReturn(enabled).when(getTopDisplayRotationCompatPolicy())
+        doReturn(enabled).when(mDisplayContent.mAppCompatCameraPolicy)
                 .isTreatmentEnabledForActivity(eq(mActivityStack.top()));
     }
 
@@ -164,7 +182,7 @@
     }
 
     void setIgnoreOrientationRequest(boolean enabled) {
-        mActivityStack.top().mDisplayContent.setIgnoreOrientationRequest(enabled);
+        mDisplayContent.setIgnoreOrientationRequest(enabled);
     }
 
     void setTopActivityAsEmbedded(boolean embedded) {
@@ -179,20 +197,22 @@
         mActivityStack.applyTo(/* fromTop */ fromTop, ActivityRecord::removeImmediately);
     }
 
+    void createNewDisplay() {
+        mDisplayContent = new TestDisplayContent.Builder(mAtm, mDisplayWidth, mDisplayHeight)
+                .build();
+        spyOnAppCompatCameraPolicy();
+    }
+
     void createNewTask() {
-        final DisplayContent displayContent = new TestDisplayContent
-                .Builder(mAtm, mDisplayWidth, mDisplayHeight).build();
         final Task newTask = new WindowTestsBase.TaskBuilder(mSupervisor)
-                .setDisplay(displayContent).build();
+                .setDisplay(mDisplayContent).build();
         mTaskStack.push(newTask);
     }
 
     void createNewTaskWithBaseActivity() {
-        final DisplayContent displayContent = new TestDisplayContent
-                .Builder(mAtm, mDisplayWidth, mDisplayHeight).build();
         final Task newTask = new WindowTestsBase.TaskBuilder(mSupervisor)
                 .setCreateActivity(true)
-                .setDisplay(displayContent).build();
+                .setDisplay(mDisplayContent).build();
         mTaskStack.push(newTask);
         pushActivity(newTask.getTopNonFinishingActivity());
     }
@@ -319,7 +339,10 @@
         pushActivity(newActivity);
     }
 
-    private void createActivityWithComponentInNewTask(boolean inNewTask) {
+    private void createActivityWithComponentInNewTask(boolean inNewTask, boolean inNewDisplay) {
+        if (inNewDisplay) {
+            createNewDisplay();
+        }
         if (inNewTask) {
             createNewTask();
         }
@@ -369,7 +392,8 @@
     }
 
     private DisplayRotationCompatPolicy getTopDisplayRotationCompatPolicy() {
-        return mActivityStack.top().mDisplayContent.mDisplayRotationCompatPolicy;
+        return mActivityStack.top().mDisplayContent
+                .mAppCompatCameraPolicy.mDisplayRotationCompatPolicy;
     }
 
     // We add the activity to the stack and spyOn() on its properties.
@@ -377,10 +401,16 @@
         mActivityStack.push(activity);
         spyOn(activity);
         spyOn(activity.mAppCompatController.getTransparentPolicy());
-        if (activity.mDisplayContent != null
-                && activity.mDisplayContent.mDisplayRotationCompatPolicy != null) {
-            spyOn(activity.mDisplayContent.mDisplayRotationCompatPolicy);
-        }
         spyOn(activity.mLetterboxUiController);
     }
+
+    private void spyOnAppCompatCameraPolicy() {
+        spyOn(mDisplayContent.mAppCompatCameraPolicy);
+        if (mDisplayContent.mAppCompatCameraPolicy.hasDisplayRotationCompatPolicy()) {
+            spyOn(mDisplayContent.mAppCompatCameraPolicy.mDisplayRotationCompatPolicy);
+        }
+        if (mDisplayContent.mAppCompatCameraPolicy.hasCameraCompatFreeformPolicy()) {
+            spyOn(mDisplayContent.mAppCompatCameraPolicy.mCameraCompatFreeformPolicy);
+        }
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java
index 9263b4f..2d94b34 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java
@@ -26,7 +26,6 @@
 import static android.view.WindowManager.PROPERTY_CAMERA_COMPAT_ALLOW_REFRESH;
 import static android.view.WindowManager.PROPERTY_CAMERA_COMPAT_ENABLE_REFRESH_VIA_PAUSE;
 
-import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.window.flags.Flags.FLAG_CAMERA_COMPAT_FOR_FREEFORM;
 
@@ -264,14 +263,29 @@
     public void testShouldRecomputeConfigurationForCameraCompat() {
         runTestScenario((robot) -> {
             robot.conf().enableCameraCompatSplitScreenAspectRatio(true);
-            robot.activity().createActivityWithComponentInNewTask();
-            robot.activateCamera(true);
-            robot.activity().setShouldCreateCompatDisplayInsets(false);
+            robot.applyOnActivity((a) -> {
+                a.createActivityWithComponentInNewTask();
+                a.activateCameraInPolicy(true);
+                a.setShouldCreateCompatDisplayInsets(false);
+            });
 
             robot.checkShouldApplyFreeformTreatmentForCameraCompat(true);
         });
     }
 
+    @Test
+    public void testIsCameraActive() {
+        runTestScenario((robot) -> {
+            robot.applyOnActivity((a) -> {
+                a.createActivityWithComponent();
+                a.activateCameraInPolicy(/* isCameraActive */ false);
+                robot.checkIsCameraActive(/* active */ false);
+                a.activateCameraInPolicy(/* isCameraActive */ true);
+                robot.checkIsCameraActive(/* active */ true);
+            });
+        });
+    }
+
     /**
      * Runs a test scenario providing a Robot.
      */
@@ -289,10 +303,6 @@
             super(wm, atm, supervisor);
         }
 
-        void activateCamera(boolean isCameraActive) {
-            doReturn(isCameraActive).when(activity().top()).isCameraActive();
-        }
-
         void checkShouldRefreshActivityForCameraCompat(boolean expected) {
             Assert.assertEquals(getAppCompatCameraOverrides()
                     .shouldRefreshActivityForCameraCompat(), expected);
@@ -313,6 +323,10 @@
                     .shouldApplyFreeformTreatmentForCameraCompat(), expected);
         }
 
+        void checkIsCameraActive(boolean active) {
+            Assert.assertEquals(getAppCompatCameraOverrides().isCameraActive(), active);
+        }
+
         private AppCompatCameraOverrides getAppCompatCameraOverrides() {
             return activity().top().mAppCompatController.getAppCompatCameraOverrides();
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraPolicyTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraPolicyTest.java
index 4116313..006b370 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraPolicyTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraPolicyTest.java
@@ -16,23 +16,20 @@
 
 package com.android.server.wm;
 
-import static android.content.pm.ActivityInfo.OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA;
-import static android.content.pm.ActivityInfo.OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA;
-
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.window.flags.Flags.FLAG_CAMERA_COMPAT_FOR_FREEFORM;
 
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
+import static org.mockito.ArgumentMatchers.any;
 
 import android.compat.testing.PlatformCompatChangeRule;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.annotation.NonNull;
 
-import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
-import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
-
+import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestRule;
@@ -54,96 +51,128 @@
     public TestRule compatChangeRule = new PlatformCompatChangeRule();
 
     @Test
-    @DisableCompatChanges({OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA})
-    public void testRecomputeConfigurationForCameraCompatIfNeeded_allDisabledNoRecompute() {
+    public void testDisplayRotationCompatPolicy_presentWhenEnabled() {
         runTestScenario((robot) -> {
-            robot.activity().createActivityWithComponent();
-            robot.conf().enableCameraCompatSplitScreenAspectRatio(false);
-            robot.activateCamera(/* isCameraActive */ false);
-
-            robot.recomputeConfigurationForCameraCompatIfNeeded();
-            robot.checkRecomputeConfigurationInvoked(/* invoked */ false);
-
+            robot.conf().enableCameraCompatTreatmentAtBuildTime(true);
+            robot.activity().createActivityWithComponentInNewTaskAndDisplay();
+            robot.checkTopActivityHasDisplayRotationCompatPolicy(true);
         });
     }
 
     @Test
-    @EnableCompatChanges({OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA})
-    public void testRecomputeConfigurationForCameraCompatIfNeeded_cameraEnabledRecompute() {
+    public void testDisplayRotationCompatPolicy_notPresentWhenDisabled() {
         runTestScenario((robot) -> {
-            robot.activity().createActivityWithComponent();
-            robot.conf().enableCameraCompatSplitScreenAspectRatio(false);
-            robot.activateCamera(/* isCameraActive */ false);
-
-            robot.recomputeConfigurationForCameraCompatIfNeeded();
-            robot.checkRecomputeConfigurationInvoked(/* invoked */ true);
+            robot.conf().enableCameraCompatTreatmentAtBuildTime(false);
+            robot.activity().createActivityWithComponentInNewTaskAndDisplay();
+            robot.checkTopActivityHasDisplayRotationCompatPolicy(false);
         });
     }
 
     @Test
-    @DisableCompatChanges({OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA})
-    public void testRecomputeConfigurationForCameraSplitScreenCompatIfNeeded_recompute() {
+    @EnableFlags(FLAG_CAMERA_COMPAT_FOR_FREEFORM)
+    public void testCameraCompatFreeformPolicy_presentWhenEnabledAndDW() {
         runTestScenario((robot) -> {
-            robot.activity().createActivityWithComponent();
-            robot.conf().enableCameraCompatSplitScreenAspectRatio(true);
-            robot.activateCamera(/* isCameraActive */ false);
-
-            robot.recomputeConfigurationForCameraCompatIfNeeded();
-            robot.checkRecomputeConfigurationInvoked(/* invoked */ true);
+            robot.allowEnterDesktopMode(true);
+            robot.activity().createActivityWithComponentInNewTaskAndDisplay();
+            robot.checkTopActivityHasCameraCompatFreeformPolicy(true);
         });
     }
 
     @Test
-    @DisableCompatChanges({OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA})
-    @EnableCompatChanges({OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA})
-    public void testRecomputeConfigurationForCameraSplitScreenCompatIfNeededWithCamera_recompute() {
+    @EnableFlags(FLAG_CAMERA_COMPAT_FOR_FREEFORM)
+    public void testCameraCompatFreeformPolicy_notPresentWhenNoDW() {
         runTestScenario((robot) -> {
-            robot.activity().createActivityWithComponent();
-            robot.conf().enableCameraCompatSplitScreenAspectRatio(false);
-            robot.activateCamera(/* isCameraActive */ true);
-
-            robot.recomputeConfigurationForCameraCompatIfNeeded();
-            robot.checkRecomputeConfigurationInvoked(/* invoked */ true);
+            robot.allowEnterDesktopMode(false);
+            robot.activity().createActivityWithComponentInNewTaskAndDisplay();
+            robot.checkTopActivityHasCameraCompatFreeformPolicy(false);
         });
     }
 
-    void runTestScenario(@NonNull Consumer<CameraPolicyRobotTest> consumer) {
+    @Test
+    @DisableFlags(FLAG_CAMERA_COMPAT_FOR_FREEFORM)
+    public void testCameraCompatFreeformPolicy_notPresentWhenNoFlag() {
+        runTestScenario((robot) -> {
+            robot.allowEnterDesktopMode(true);
+            robot.activity().createActivityWithComponentInNewTaskAndDisplay();
+            robot.checkTopActivityHasCameraCompatFreeformPolicy(false);
+        });
+    }
+
+    @Test
+    @EnableFlags(FLAG_CAMERA_COMPAT_FOR_FREEFORM)
+    public void testCameraCompatFreeformPolicy_notPresentWhenNoFlagAndNoDW() {
+        runTestScenario((robot) -> {
+            robot.allowEnterDesktopMode(false);
+            robot.activity().createActivityWithComponentInNewTaskAndDisplay();
+            robot.checkTopActivityHasCameraCompatFreeformPolicy(false);
+        });
+    }
+
+    /**
+     * Runs a test scenario providing a Robot.
+     */
+    void runTestScenario(@NonNull Consumer<DisplayRotationPolicyRobotTest> consumer) {
         spyOn(mWm.mLetterboxConfiguration);
-        final CameraPolicyRobotTest robot = new CameraPolicyRobotTest(mWm, mAtm, mSupervisor);
+        final DisplayRotationPolicyRobotTest robot =
+                new DisplayRotationPolicyRobotTest(mWm, mAtm, mSupervisor);
         consumer.accept(robot);
     }
 
-    private static class CameraPolicyRobotTest extends AppCompatRobotBase {
+    @Test
+    public void testIsCameraCompatTreatmentActive_whenTreatmentForTopActivityIsEnabled() {
+        runTestScenario((robot) -> {
+            robot.applyOnActivity((a)-> {
+                a.createActivityWithComponent();
+                a.enableTreatmentForTopActivity(/* enabled */ true);
+            });
 
-        private final WindowManagerService mWm;
+            robot.checkIsCameraCompatTreatmentActiveForTopActivity(/* active */ true);
+        });
+    }
 
-        CameraPolicyRobotTest(@NonNull WindowManagerService wm,
+    @Test
+    public void testIsCameraCompatTreatmentNotActive_whenTreatmentForTopActivityIsDisabled() {
+        runTestScenario((robot) -> {
+            robot.applyOnActivity((a)-> {
+                a.createActivityWithComponent();
+                a.enableTreatmentForTopActivity(/* enabled */ false);
+            });
+
+            robot.checkIsCameraCompatTreatmentActiveForTopActivity(/* active */ false);
+        });
+    }
+
+    private static class DisplayRotationPolicyRobotTest extends AppCompatRobotBase {
+
+        DisplayRotationPolicyRobotTest(@NonNull WindowManagerService wm,
                 @NonNull ActivityTaskManagerService atm,
                 @NonNull ActivityTaskSupervisor supervisor) {
             super(wm, atm, supervisor);
-            mWm = wm;
-            spyOn(mWm);
         }
 
-        void activateCamera(boolean isCameraActive) {
-            doReturn(isCameraActive).when(activity().top()).isCameraActive();
+        void checkTopActivityHasDisplayRotationCompatPolicy(boolean exists) {
+            Assert.assertEquals(exists, activity().top().mDisplayContent
+                    .mAppCompatCameraPolicy.hasDisplayRotationCompatPolicy());
         }
 
-        void recomputeConfigurationForCameraCompatIfNeeded() {
-            getAppCompatCameraPolicy().recomputeConfigurationForCameraCompatIfNeeded();
+        void checkTopActivityHasCameraCompatFreeformPolicy(boolean exists) {
+            Assert.assertEquals(exists, activity().top().mDisplayContent
+                    .mAppCompatCameraPolicy.hasCameraCompatFreeformPolicy());
         }
 
-        void checkRecomputeConfigurationInvoked(boolean invoked) {
-            if (invoked) {
-                verify(activity().top()).recomputeConfiguration();
-            } else {
-                verify(activity().top(), never()).recomputeConfiguration();
-            }
+        void checkIsCameraCompatTreatmentActiveForTopActivity(boolean active) {
+            Assert.assertEquals(getTopAppCompatCameraPolicy()
+                    .isTreatmentEnabledForActivity(activity().top()), active);
         }
 
-        private AppCompatCameraPolicy getAppCompatCameraPolicy() {
-            return activity().top().mAppCompatController.getAppCompatCameraPolicy();
+        // TODO(b/350460645): Create Desktop Windowing Robot to reuse common functionalities.
+        void allowEnterDesktopMode(boolean isAllowed) {
+            doReturn(isAllowed).when(() ->
+                    DesktopModeLaunchParamsModifier.canEnterDesktopMode(any()));
+        }
+
+        private AppCompatCameraPolicy getTopAppCompatCameraPolicy() {
+            return activity().top().mDisplayContent.mAppCompatCameraPolicy;
         }
     }
-
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java
index 1720b64..35c2ee0 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationOverridesTest.java
@@ -15,12 +15,8 @@
  */
 package com.android.server.wm;
 
-import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH;
 import static android.content.pm.ActivityInfo.OVERRIDE_ENABLE_COMPAT_IGNORE_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED;
-import static android.content.pm.ActivityInfo.OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION;
-import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED;
 import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_IGNORING_ORIENTATION_REQUEST_WHEN_LOOP_DETECTED;
-import static android.view.WindowManager.PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.server.wm.AppCompatOrientationOverrides.OrientationOverridesState.MIN_COUNT_TO_IGNORE_REQUEST_IN_LOOP;
@@ -31,7 +27,6 @@
 import static org.junit.Assert.assertTrue;
 
 import android.compat.testing.PlatformCompatChangeRule;
-import android.content.res.Configuration;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.annotation.NonNull;
@@ -62,82 +57,6 @@
     public TestRule compatChangeRule = new PlatformCompatChangeRule();
 
     @Test
-    @EnableCompatChanges({OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION})
-    public void testShouldIgnoreRequestedOrientation_activityRelaunching_returnsTrue() {
-        runTestScenario((robot) -> {
-            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
-            robot.activity().createActivityWithComponent();
-            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
-
-            robot.checkShouldIgnoreRequestedOrientation(/* expected */ true,
-                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
-        });
-    }
-
-    @Test
-    @EnableCompatChanges({OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION})
-    public void testShouldIgnoreRequestedOrientation_cameraCompatTreatment_returnsTrue() {
-        runTestScenario((robot) -> {
-            robot.applyOnConf((c) -> {
-                c.enableCameraCompatTreatment(true);
-                c.enableCameraCompatTreatmentAtBuildTime(true);
-                c.enablePolicyForIgnoringRequestedOrientation(true);
-            });
-            robot.applyOnActivity((a) -> {
-                a.createActivityWithComponentInNewTask();
-                a.enableTreatmentForTopActivity(true);
-            });
-            robot.prepareRelaunchingAfterRequestedOrientationChanged(false);
-
-            robot.checkShouldIgnoreRequestedOrientation(/* expected */ true,
-                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
-        });
-    }
-
-    @Test
-    public void testShouldIgnoreRequestedOrientation_overrideDisabled_returnsFalse() {
-        runTestScenario((robot) -> {
-            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
-
-            robot.activity().createActivityWithComponent();
-            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
-
-            robot.checkShouldIgnoreRequestedOrientation(/* expected */ false,
-                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
-        });
-    }
-
-    @Test
-    public void testShouldIgnoreRequestedOrientation_propertyIsTrue_returnsTrue() {
-        runTestScenario((robot) -> {
-            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
-            robot.prop().enable(PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION);
-
-            robot.activity().createActivityWithComponent();
-            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
-
-            robot.checkShouldIgnoreRequestedOrientation(/* expected */ true,
-                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
-        });
-    }
-
-    @Test
-    @EnableCompatChanges({OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION})
-    public void testShouldIgnoreRequestedOrientation_propertyIsFalseAndOverride_returnsFalse()
-            throws Exception {
-        runTestScenario((robot) -> {
-            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
-            robot.prop().disable(PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION);
-
-            robot.activity().createActivityWithComponent();
-            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
-
-            robot.checkShouldIgnoreRequestedOrientation(/* expected */ false,
-                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
-        });
-    }
-
-    @Test
     public void testShouldIgnoreOrientationRequestLoop_overrideDisabled_returnsFalse() {
         runTestScenario((robot) -> {
             robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
@@ -239,21 +158,6 @@
         });
     }
 
-    @Test
-    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH})
-    public void testShouldIgnoreRequestedOrientation_flagIsDisabled_returnsFalse() {
-        runTestScenario((robot) -> {
-            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
-            robot.applyOnActivity((a) -> {
-                a.createActivityWithComponent();
-                a.setLetterboxedForFixedOrientationAndAspectRatio(false);
-            });
-
-            robot.checkShouldIgnoreRequestedOrientation(/* expected */ false,
-                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
-        });
-    }
-
     /**
      * Runs a test scenario providing a Robot.
      */
@@ -276,10 +180,6 @@
             mTestCurrentTimeMillisSupplier = new CurrentTimeMillisSupplierFake();
         }
 
-        void prepareRelaunchingAfterRequestedOrientationChanged(boolean enabled) {
-            getTopOrientationOverrides().setRelaunchingAfterRequestedOrientationChanged(enabled);
-        }
-
         // Useful to reduce timeout during tests
         void prepareMockedTime() {
             getTopOrientationOverrides().mOrientationOverridesState.mCurrentTimeMillisSupplier =
@@ -290,12 +190,6 @@
             mTestCurrentTimeMillisSupplier.delay(SET_ORIENTATION_REQUEST_COUNTER_TIMEOUT_MS);
         }
 
-        void checkShouldIgnoreRequestedOrientation(boolean expected,
-                @Configuration.Orientation int requestedOrientation) {
-            assertEquals(expected, getTopOrientationOverrides()
-                    .shouldIgnoreRequestedOrientation(requestedOrientation));
-        }
-
         void checkExpectedLoopCount(int expectedCount) {
             assertEquals(expectedCount, getTopOrientationOverrides()
                     .getSetOrientationRequestCounter());
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationPolicyTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationPolicyTest.java
index 9885a2d..aa520e9 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationPolicyTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatOrientationPolicyTest.java
@@ -18,6 +18,8 @@
 
 import static android.content.pm.ActivityInfo.OVERRIDE_ANY_ORIENTATION;
 import static android.content.pm.ActivityInfo.OVERRIDE_ANY_ORIENTATION_TO_USER;
+import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH;
+import static android.content.pm.ActivityInfo.OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION;
 import static android.content.pm.ActivityInfo.OVERRIDE_LANDSCAPE_ORIENTATION_TO_REVERSE_LANDSCAPE;
 import static android.content.pm.ActivityInfo.OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA;
 import static android.content.pm.ActivityInfo.OVERRIDE_UNDEFINED_ORIENTATION_TO_NOSENSOR;
@@ -33,13 +35,16 @@
 import static android.content.pm.PackageManager.USER_MIN_ASPECT_RATIO_FULLSCREEN;
 import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_ORIENTATION_OVERRIDE;
 import static android.view.WindowManager.PROPERTY_COMPAT_ALLOW_USER_ASPECT_RATIO_FULLSCREEN_OVERRIDE;
+import static android.view.WindowManager.PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.verify;
 
 import android.compat.testing.PlatformCompatChangeRule;
 import android.content.pm.ActivityInfo;
+import android.content.res.Configuration;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.annotation.NonNull;
@@ -266,7 +271,7 @@
                 c.enableCameraCompatTreatmentAtBuildTime(true);
             });
             robot.applyOnActivity((a) -> {
-                a.createActivityWithComponentInNewTask();
+                a.createActivityWithComponentInNewTaskAndDisplay();
                 a.setTopActivityEligibleForOrientationOverride(false);
             });
 
@@ -285,7 +290,7 @@
                 c.enableCameraCompatTreatmentAtBuildTime(true);
             });
             robot.applyOnActivity((a) -> {
-                a.createActivityWithComponentInNewTask();
+                a.createActivityWithComponentInNewTaskAndDisplay();
                 a.setTopActivityEligibleForOrientationOverride(true);
             });
 
@@ -315,7 +320,7 @@
                 c.enableCameraCompatTreatmentAtBuildTime(true);
             });
             robot.applyOnActivity((a) -> {
-                a.createActivityWithComponentInNewTask();
+                a.createActivityWithComponentInNewTaskAndDisplay();
                 a.setTopActivityCameraActive(false);
             });
 
@@ -398,6 +403,97 @@
         });
     }
 
+    @Test
+    @EnableCompatChanges({OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION})
+    public void testShouldIgnoreRequestedOrientation_activityRelaunching_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
+            robot.activity().createActivityWithComponent();
+            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
+
+            robot.checkShouldIgnoreRequestedOrientation(/* expected */ true,
+                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
+        });
+    }
+
+    @Test
+    @EnableCompatChanges({OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION})
+    public void testShouldIgnoreRequestedOrientation_cameraCompatTreatment_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.applyOnConf((c) -> {
+                c.enableCameraCompatTreatment(true);
+                c.enableCameraCompatTreatmentAtBuildTime(true);
+                c.enablePolicyForIgnoringRequestedOrientation(true);
+            });
+            robot.applyOnActivity((a) -> {
+                a.createActivityWithComponentInNewTask();
+                a.enableTreatmentForTopActivity(true);
+            });
+            robot.prepareRelaunchingAfterRequestedOrientationChanged(false);
+
+            robot.checkShouldIgnoreRequestedOrientation(/* expected */ true,
+                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
+        });
+    }
+
+    @Test
+    public void testShouldIgnoreRequestedOrientation_overrideDisabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
+
+            robot.activity().createActivityWithComponent();
+            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
+
+            robot.checkShouldIgnoreRequestedOrientation(/* expected */ false,
+                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
+        });
+    }
+
+    @Test
+    public void testShouldIgnoreRequestedOrientation_propertyIsTrue_returnsTrue() {
+        runTestScenario((robot) -> {
+            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
+            robot.prop().enable(PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION);
+
+            robot.activity().createActivityWithComponent();
+            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
+
+            robot.checkShouldIgnoreRequestedOrientation(/* expected */ true,
+                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
+        });
+    }
+
+    @Test
+    @EnableCompatChanges({OVERRIDE_ENABLE_COMPAT_IGNORE_REQUESTED_ORIENTATION})
+    public void testShouldIgnoreRequestedOrientation_propertyIsFalseAndOverride_returnsFalse()
+            throws Exception {
+        runTestScenario((robot) -> {
+            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
+            robot.prop().disable(PROPERTY_COMPAT_IGNORE_REQUESTED_ORIENTATION);
+
+            robot.activity().createActivityWithComponent();
+            robot.prepareRelaunchingAfterRequestedOrientationChanged(true);
+
+            robot.checkShouldIgnoreRequestedOrientation(/* expected */ false,
+                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
+        });
+    }
+
+    @Test
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH})
+    public void testShouldIgnoreRequestedOrientation_flagIsDisabled_returnsFalse() {
+        runTestScenario((robot) -> {
+            robot.conf().enablePolicyForIgnoringRequestedOrientation(true);
+            robot.applyOnActivity((a) -> {
+                a.createActivityWithComponent();
+                a.setLetterboxedForFixedOrientationAndAspectRatio(false);
+            });
+
+            robot.checkShouldIgnoreRequestedOrientation(/* expected */ false,
+                    /* requestedOrientation */ SCREEN_ORIENTATION_UNSPECIFIED);
+        });
+    }
+
 
     /**
      * Runs a test scenario with an existing activity providing a Robot.
@@ -440,6 +536,10 @@
             }
         }
 
+        void prepareRelaunchingAfterRequestedOrientationChanged(boolean enabled) {
+            getTopOrientationOverrides().setRelaunchingAfterRequestedOrientationChanged(enabled);
+        }
+
         int overrideOrientationIfNeeded(@ActivityInfo.ScreenOrientation int candidate) {
             return activity().top().mAppCompatController.getOrientationPolicy()
                     .overrideOrientationIfNeeded(candidate);
@@ -451,12 +551,27 @@
 
         void checkOverrideOrientation(@ActivityInfo.ScreenOrientation int candidate,
                                       @ActivityInfo.ScreenOrientation int expected) {
-            Assert.assertEquals(expected, overrideOrientationIfNeeded(candidate));
+            assertEquals(expected, overrideOrientationIfNeeded(candidate));
         }
 
         void checkOverrideOrientationIsNot(@ActivityInfo.ScreenOrientation int candidate,
                                            @ActivityInfo.ScreenOrientation int notExpected) {
             Assert.assertNotEquals(notExpected, overrideOrientationIfNeeded(candidate));
         }
+
+        void checkShouldIgnoreRequestedOrientation(boolean expected,
+                @Configuration.Orientation int requestedOrientation) {
+            assertEquals(expected, getTopAppCompatOrientationPolicy()
+                    .shouldIgnoreRequestedOrientation(requestedOrientation));
+        }
+
+        private AppCompatOrientationOverrides getTopOrientationOverrides() {
+            return activity().top().mAppCompatController.getAppCompatOverrides()
+                    .getAppCompatOrientationOverrides();
+        }
+
+        private AppCompatOrientationPolicy getTopAppCompatOrientationPolicy() {
+            return activity().top().mAppCompatController.getOrientationPolicy();
+        }
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index 6957502..5739a04 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -117,8 +117,8 @@
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
-import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.util.ArraySet;
 import android.view.Display;
@@ -2832,7 +2832,7 @@
         doReturn(true).when(() ->
                 DesktopModeLaunchParamsModifier.canEnterDesktopMode(any()));
 
-        assertNotNull(createNewDisplay().mCameraCompatFreeformPolicy);
+        assertTrue(createNewDisplay().mAppCompatCameraPolicy.hasCameraCompatFreeformPolicy());
     }
 
     @DisableFlags(FLAG_CAMERA_COMPAT_FOR_FREEFORM)
@@ -2841,14 +2841,14 @@
         doReturn(true).when(() ->
                 DesktopModeLaunchParamsModifier.canEnterDesktopMode(any()));
 
-        assertNull(createNewDisplay().mCameraCompatFreeformPolicy);
+        assertFalse(createNewDisplay().mAppCompatCameraPolicy.hasCameraCompatFreeformPolicy());
     }
 
     @EnableFlags(FLAG_CAMERA_COMPAT_FOR_FREEFORM)
     @DisableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_MODE)
     @Test
     public void desktopWindowingFlagNotEnabled_cameraCompatFreeformPolicyIsNull() {
-        assertNull(createNewDisplay().mCameraCompatFreeformPolicy);
+        assertFalse(createNewDisplay().mAppCompatCameraPolicy.hasCameraCompatFreeformPolicy());
     }
 
     private void removeRootTaskTests(Runnable runnable) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayRotationTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayRotationTests.java
index e3a8542..2e488d8 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayRotationTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayRotationTests.java
@@ -1607,6 +1607,7 @@
                     .thenReturn(mMockDeviceStateManager);
 
             mDeviceStateController = mock(DeviceStateController.class);
+            mMockDisplayContent.mAppCompatCameraPolicy = mock(AppCompatCameraPolicy.class);
             mTarget = new TestDisplayRotation(mMockDisplayContent, mMockDisplayAddress,
                     mMockDisplayPolicy, mMockDisplayWindowSettings, mMockContext,
                     mDeviceStateController);
diff --git a/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java
index 74e2d44..51b3c48 100644
--- a/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/LetterboxUiControllerTest.java
@@ -636,8 +636,9 @@
     @Test
     @EnableCompatChanges({OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA})
     public void shouldOverrideMinAspectRatioForCamera_overrideEnabled_returnsTrue() {
-        doReturn(true).when(mActivity).isCameraActive();
-        mController = new LetterboxUiController(mWm, mActivity);
+        mActivity = setUpActivityWithComponent();
+        doReturn(true).when(mActivity.mAppCompatController
+                .getAppCompatCameraOverrides()).isCameraActive();
 
         assertTrue(mActivity.mAppCompatController.getAppCompatCameraOverrides()
                 .shouldOverrideMinAspectRatioForCamera());
@@ -647,9 +648,10 @@
     @EnableCompatChanges({OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA})
     public void shouldOverrideMinAspectRatioForCamera_propertyTrue_overrideEnabled_returnsTrue()
             throws Exception {
-        doReturn(true).when(mActivity).isCameraActive();
         mockThatProperty(PROPERTY_COMPAT_ALLOW_MIN_ASPECT_RATIO_OVERRIDE, /* value */ true);
-        mController = new LetterboxUiController(mWm, mActivity);
+        mActivity = setUpActivityWithComponent();
+        doReturn(true).when(mActivity.mAppCompatController
+                .getAppCompatCameraOverrides()).isCameraActive();
 
         assertTrue(mActivity.mAppCompatController.getAppCompatCameraOverrides()
                 .shouldOverrideMinAspectRatioForCamera());
@@ -659,9 +661,10 @@
     @EnableCompatChanges({OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA})
     public void shouldOverrideMinAspectRatioForCamera_propertyTrue_overrideEnabled_returnsFalse()
             throws Exception {
-        doReturn(false).when(mActivity).isCameraActive();
         mockThatProperty(PROPERTY_COMPAT_ALLOW_MIN_ASPECT_RATIO_OVERRIDE, /* value */ true);
-        mController = new LetterboxUiController(mWm, mActivity);
+        mActivity = setUpActivityWithComponent();
+        doReturn(false).when(mActivity.mAppCompatController
+                .getAppCompatCameraOverrides()).isCameraActive();
 
         assertFalse(mActivity.mAppCompatController.getAppCompatCameraOverrides()
                 .shouldOverrideMinAspectRatioForCamera());
@@ -671,9 +674,10 @@
     @DisableCompatChanges({OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA})
     public void shouldOverrideMinAspectRatioForCamera_propertyTrue_overrideDisabled_returnsFalse()
             throws Exception {
-        doReturn(true).when(mActivity).isCameraActive();
         mockThatProperty(PROPERTY_COMPAT_ALLOW_MIN_ASPECT_RATIO_OVERRIDE, /* value */ true);
-        mController = new LetterboxUiController(mWm, mActivity);
+        mActivity = setUpActivityWithComponent();
+        doReturn(true).when(mActivity.mAppCompatController
+                .getAppCompatCameraOverrides()).isCameraActive();
 
         assertFalse(mActivity.mAppCompatController.getAppCompatCameraOverrides()
                 .shouldOverrideMinAspectRatioForCamera());
@@ -682,8 +686,9 @@
     @Test
     @DisableCompatChanges({OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA})
     public void shouldOverrideMinAspectRatioForCamera_overrideDisabled_returnsFalse() {
-        doReturn(true).when(mActivity).isCameraActive();
-        mController = new LetterboxUiController(mWm, mActivity);
+        mActivity = setUpActivityWithComponent();
+        doReturn(true).when(mActivity.mAppCompatController
+                .getAppCompatCameraOverrides()).isCameraActive();
 
         assertFalse(mActivity.mAppCompatController.getAppCompatCameraOverrides()
                 .shouldOverrideMinAspectRatioForCamera());
@@ -694,7 +699,7 @@
     public void shouldOverrideMinAspectRatioForCamera_propertyFalse_overrideEnabled_returnsFalse()
             throws Exception {
         mockThatProperty(PROPERTY_COMPAT_ALLOW_MIN_ASPECT_RATIO_OVERRIDE, /* value */ false);
-        mController = new LetterboxUiController(mWm, mActivity);
+        mActivity = setUpActivityWithComponent();
 
         assertFalse(mActivity.mAppCompatController.getAppCompatCameraOverrides()
                 .shouldOverrideMinAspectRatioForCamera());
@@ -705,8 +710,11 @@
     public void shouldOverrideMinAspectRatioForCamera_propertyFalse_noOverride_returnsFalse()
             throws Exception {
         mockThatProperty(PROPERTY_COMPAT_ALLOW_MIN_ASPECT_RATIO_OVERRIDE, /* value */ false);
-        doReturn(true).when(mActivity).isCameraActive();
-        mController = new LetterboxUiController(mWm, mActivity);
+
+        mActivity = setUpActivityWithComponent();
+
+        doReturn(true).when(mActivity.mAppCompatController
+                .getAppCompatCameraOverrides()).isCameraActive();
 
         assertFalse(mActivity.mAppCompatController.getAppCompatCameraOverrides()
                 .shouldOverrideMinAspectRatioForCamera());
@@ -848,8 +856,8 @@
         assertEquals(1.5f, mController.getFixedOrientationLetterboxAspectRatio(
                 mActivity.getParent().getConfiguration()), /* delta */ 0.01);
 
-        spyOn(mDisplayContent.mDisplayRotationCompatPolicy);
-        doReturn(true).when(mDisplayContent.mDisplayRotationCompatPolicy)
+        spyOn(mDisplayContent.mAppCompatCameraPolicy);
+        doReturn(true).when(mDisplayContent.mAppCompatCameraPolicy)
                 .isTreatmentEnabledForActivity(eq(mActivity));
 
         assertEquals(mController.getSplitScreenAspectRatio(),
@@ -980,6 +988,7 @@
                 .setComponent(ComponentName.createRelative(mContext,
                         com.android.server.wm.LetterboxUiControllerTest.class.getName()))
                 .build();
+        spyOn(activity.mAppCompatController.getAppCompatCameraOverrides());
         return activity;
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index c962a3f..4220f31 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -442,15 +442,7 @@
                 dc.getDisplayPolicy().release();
                 // Unregister SensorEventListener (foldable device may register for hinge angle).
                 dc.getDisplayRotation().onDisplayRemoved();
-                if (dc.mDisplayRotationCompatPolicy != null) {
-                    dc.mDisplayRotationCompatPolicy.dispose();
-                }
-                if (dc.mCameraCompatFreeformPolicy != null) {
-                    dc.mCameraCompatFreeformPolicy.dispose();
-                }
-                if (dc.mCameraStateMonitor != null) {
-                    dc.mCameraStateMonitor.dispose();
-                }
+                dc.mAppCompatCameraPolicy.dispose();
             }
         }