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 609e5af..9256c60 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/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;